模型存储与管理

// model_access_pattern.go
package storage

import (
    "time"
)

// AccessPattern 模型文件访问模式
type AccessPattern struct {
    // 检查点写入模式
    CheckpointWrite CheckpointWritePattern
    // 检查点读取模式（恢复训练）
    CheckpointRead CheckpointReadPattern
    // 模型导出模式
    ModelExport ModelExportPattern
    // 模型加载模式（推理）
    ModelLoad ModelLoadPattern
}

// CheckpointWritePattern 检查点写入特征
type CheckpointWritePattern struct {
    // 写入频率：每 N 步保存一次
    SaveInterval int
    // 典型间隔：1000-10000 步
    TypicalSteps int
    // 写入大小：与模型规模成正比
    WriteSize int64
    // 并发写入：分布式训练时所有 rank 同时写
    ConcurrentWriters int
    // 写入带宽需求：高峰期需要高吞吐
    RequiredBandwidth int64 // bytes/s
    // 写入模式：顺序写为主
    Sequential bool
}

// CheckpointReadPattern 检查点读取特征
type CheckpointReadPattern struct {
    // 读取频率：训练恢复时
    Frequency string // "rare" - 仅故障恢复或继续训练
    // 读取大小：完整检查点
    ReadSize int64
    // 并发读取：所有 rank 同时读
    ConcurrentReaders int
    // 读取带宽需求
    RequiredBandwidth int64
    // 读取模式：顺序读
    Sequential bool
}

// ModelExportPattern 模型导出特征
type ModelExportPattern struct {
    // 导出格式
    Formats []string // ["pytorch", "safetensors", "onnx", "tensorrt"]
    // 导出频率：训练结束或里程碑
    Frequency string // "milestone"
    // 处理类型：可能需要格式转换
    RequiresConversion bool
}

// ModelLoadPattern 模型加载特征（推理场景）
type ModelLoadPattern struct {
    // 加载频率：服务启动、扩容
    Frequency string // "on_demand"
    // 加载速度要求：影响服务启动时间
    LatencyRequirement time.Duration
    // 并发加载：多实例同时加载
    ConcurrentLoaders int
    // 部分加载：分层加载支持
    SupportsPartialLoad bool
}

// CalculateStorageRequirements 计算存储需求
func CalculateStorageRequirements(
    modelParams int64, // 模型参数量
    trainingDays int,
    checkpointInterval int, // 小时
    keepCheckpoints int, // 保留数量
) StorageRequirements {

    // 每个参数占用字节（FP16 = 2, FP32 = 4, 混合精度约 6）
    bytesPerParam := int64(6)

    // 检查点大小 = 参数 × 每参数字节 × 2（参数+优化器状态）
    checkpointSize := modelParams * bytesPerParam * 2

    // 每天检查点数
    checkpointsPerDay := 24 / checkpointInterval

    // 总存储需求
    totalCheckpointStorage := checkpointSize * int64(keepCheckpoints)

    // 峰值写入带宽（假设检查点写入时间 < 10分钟）
    peakWriteBandwidth := checkpointSize / (10 * 60) // bytes/s

    return StorageRequirements{
        CheckpointSize:       checkpointSize,
        TotalStorage:         totalCheckpointStorage,
        PeakWriteBandwidth:   peakWriteBandwidth,
        CheckpointsPerDay:    checkpointsPerDay,
        RecommendedReplicas:  3, // 三副本保证可靠性
    }
}

type StorageRequirements struct {
    CheckpointSize      int64
    TotalStorage        int64
    PeakWriteBandwidth  int64
    CheckpointsPerDay   int
    RecommendedReplicas int
}

// parallel_fs.go
package storage

import (
    "context"
    "fmt"
    "io"
    "os"
    "path/filepath"
    "sync"
    "syscall"
)

// ParallelFileSystem 并行文件系统接口
type ParallelFileSystem interface {
    // 并行写入
    ParallelWrite(ctx context.Context, path string, data []byte, opts WriteOptions) error
    // 并行读取
    ParallelRead(ctx context.Context, path string, opts ReadOptions) ([]byte, error)
    // 条带化配置
    SetStriping(path string, stripeCount int, stripeSize int64) error
    // 获取文件布局
    GetLayout(path string) (*FileLayout, error)
}

// LustreFS Lustre 文件系统实现
type LustreFS struct {
    mountPoint   string
    defaultStripeCount int
    defaultStripeSize  int64
    maxStripeCount     int
}

// WriteOptions 写入选项
type WriteOptions struct {
    StripeCount   int   // 条带数量
    StripeSize    int64 // 条带大小
    Parallel      bool  // 是否并行写
    DirectIO      bool  // 是否使用 Direct I/O
    Compression   string // 压缩算法
}

// ReadOptions 读取选项
type ReadOptions struct {
    Parallel    bool
    DirectIO    bool
    ReadAhead   int64 // 预读大小
}

// FileLayout 文件布局信息
type FileLayout struct {
    StripeCount  int
    StripeSize   int64
    StripeOffset int64
    OSTs         []int // Object Storage Targets
}

func NewLustreFS(mountPoint string) *LustreFS {
    return &LustreFS{
        mountPoint:         mountPoint,
        defaultStripeCount: 4,
        defaultStripeSize:  1 << 20, // 1MB
        maxStripeCount:     128,
    }
}

// SetStriping 设置文件条带化
func (l *LustreFS) SetStriping(path string, stripeCount int, stripeSize int64) error {
    fullPath := filepath.Join(l.mountPoint, path)

    // 确保目录存在
    dir := filepath.Dir(fullPath)
    if err := os.MkdirAll(dir, 0755); err != nil {
        return fmt.Errorf("create directory: %w", err)
    }

    // 使用 lfs setstripe 命令设置条带
    // 实际实现中使用 Lustre API
    return l.setStripeViaAPI(fullPath, stripeCount, stripeSize)
}

func (l *LustreFS) setStripeViaAPI(path string, stripeCount int, stripeSize int64) error {
    // Lustre ioctl 调用设置条带
    // LOV_USER_MAGIC_V3 结构
    type lovUserMd struct {
        Magic       uint32
        Pattern     uint32
        ObjectID    uint64
        ObjectSeq   uint64
        StripeSize  uint32
        StripeCount uint16
        StripeOffset uint16
    }

    // 这里简化处理，实际需要调用 Lustre API
    return nil
}

// ParallelWrite 并行写入大文件
func (l *LustreFS) ParallelWrite(ctx context.Context, path string, data []byte, opts WriteOptions) error {
    fullPath := filepath.Join(l.mountPoint, path)

    // 设置条带化
    if opts.StripeCount > 0 {
        if err := l.SetStriping(path, opts.StripeCount, opts.StripeSize); err != nil {
            return err
        }
    }

    // 打开文件
    flags := os.O_CREATE | os.O_WRONLY | os.O_TRUNC
    if opts.DirectIO {
        flags |= syscall.O_DIRECT
    }

    file, err := os.OpenFile(fullPath, flags, 0644)
    if err != nil {
        return fmt.Errorf("open file: %w", err)
    }
    defer file.Close()

    if !opts.Parallel {
        // 顺序写入
        _, err = file.Write(data)
        return err
    }

    // 并行写入
    return l.parallelWriteChunks(ctx, file, data, opts)
}

func (l *LustreFS) parallelWriteChunks(ctx context.Context, file *os.File, data []byte, opts WriteOptions) error {
    chunkSize := opts.StripeSize
    if chunkSize == 0 {
        chunkSize = l.defaultStripeSize
    }

    numChunks := (int64(len(data)) + chunkSize - 1) / chunkSize

    var wg sync.WaitGroup
    errCh := make(chan error, numChunks)

    for i := int64(0); i < numChunks; i++ {
        start := i * chunkSize
        end := start + chunkSize
        if end > int64(len(data)) {
            end = int64(len(data))
        }

        wg.Add(1)
        go func(offset int64, chunk []byte) {
            defer wg.Done()

            select {
            case <-ctx.Done():
                errCh <- ctx.Err()
                return
            default:
            }

            _, err := file.WriteAt(chunk, offset)
            if err != nil {
                errCh <- fmt.Errorf("write at %d: %w", offset, err)
            }
        }(start, data[start:end])
    }

    wg.Wait()
    close(errCh)

    for err := range errCh {
        if err != nil {
            return err
        }
    }

    return file.Sync()
}

// ParallelRead 并行读取大文件
func (l *LustreFS) ParallelRead(ctx context.Context, path string, opts ReadOptions) ([]byte, error) {
    fullPath := filepath.Join(l.mountPoint, path)

    // 获取文件信息
    info, err := os.Stat(fullPath)
    if err != nil {
        return nil, fmt.Errorf("stat file: %w", err)
    }

    // 打开文件
    flags := os.O_RDONLY
    if opts.DirectIO {
        flags |= syscall.O_DIRECT
    }

    file, err := os.OpenFile(fullPath, flags, 0)
    if err != nil {
        return nil, fmt.Errorf("open file: %w", err)
    }
    defer file.Close()

    // 获取文件布局以优化读取
    layout, _ := l.GetLayout(path)

    data := make([]byte, info.Size())

    if !opts.Parallel {
        _, err = io.ReadFull(file, data)
        return data, err
    }

    // 根据条带布局并行读取
    return l.parallelReadChunks(ctx, file, data, layout, opts)
}

func (l *LustreFS) parallelReadChunks(
    ctx context.Context,
    file *os.File,
    data []byte,
    layout *FileLayout,
    opts ReadOptions,
) ([]byte, error) {
    chunkSize := int64(1 << 20) // 1MB default
    if layout != nil && layout.StripeSize > 0 {
        chunkSize = layout.StripeSize
    }

    numChunks := (int64(len(data)) + chunkSize - 1) / chunkSize

    var wg sync.WaitGroup
    errCh := make(chan error, numChunks)

    // 控制并发度
    sem := make(chan struct{}, 16)

    for i := int64(0); i < numChunks; i++ {
        start := i * chunkSize
        end := start + chunkSize
        if end > int64(len(data)) {
            end = int64(len(data))
        }

        wg.Add(1)
        go func(offset, length int64) {
            defer wg.Done()

            sem <- struct{}{}
            defer func() { <-sem }()

            select {
            case <-ctx.Done():
                errCh <- ctx.Err()
                return
            default:
            }

            _, err := file.ReadAt(data[offset:offset+length], offset)
            if err != nil && err != io.EOF {
                errCh <- fmt.Errorf("read at %d: %w", offset, err)
            }
        }(start, end-start)
    }

    wg.Wait()
    close(errCh)

    for err := range errCh {
        if err != nil {
            return nil, err
        }
    }

    return data, nil
}

// GetLayout 获取文件条带布局
func (l *LustreFS) GetLayout(path string) (*FileLayout, error) {
    // 通过 Lustre API 获取文件布局
    // 简化实现
    return &FileLayout{
        StripeCount:  l.defaultStripeCount,
        StripeSize:   l.defaultStripeSize,
        StripeOffset: 0,
        OSTs:         []int{0, 1, 2, 3},
    }, nil
}

// OptimalStripeConfig 计算最优条带配置
func (l *LustreFS) OptimalStripeConfig(fileSize int64, accessPattern string) (stripeCount int, stripeSize int64) {
    switch accessPattern {
    case "checkpoint_write":
        // 检查点写入：高并发，大文件
        // 使用更多条带以提高吞吐
        if fileSize > 10<<30 { // > 10GB
            return min(l.maxStripeCount, 64), 4 << 20 // 64 条带，4MB
        } else if fileSize > 1<<30 { // > 1GB
            return 32, 2 << 20 // 32 条带，2MB
        }
        return 16, 1 << 20 // 16 条带，1MB

    case "checkpoint_read":
        // 检查点读取：顺序读为主
        if fileSize > 10<<30 {
            return 32, 4 << 20
        }
        return 16, 2 << 20

    case "model_serve":
        // 模型加载：追求最快加载
        return min(l.maxStripeCount, 128), 1 << 20

    default:
        return l.defaultStripeCount, l.defaultStripeSize
    }
}

func min(a, b int) int {
    if a < b {
        return a
    }
    return b
}

// object_storage.go
package storage

import (
    "bytes"
    "context"
    "crypto/md5"
    "encoding/hex"
    "fmt"
    "io"
    "sync"
    "time"

    "github.com/aws/aws-sdk-go-v2/aws"
    "github.com/aws/aws-sdk-go-v2/service/s3"
    "github.com/aws/aws-sdk-go-v2/service/s3/types"
)

// ObjectStorage 对象存储接口
type ObjectStorage interface {
    // 上传对象
    Upload(ctx context.Context, bucket, key string, data []byte) error
    // 分片上传大对象
    MultipartUpload(ctx context.Context, bucket, key string, reader io.Reader, size int64) error
    // 下载对象
    Download(ctx context.Context, bucket, key string) ([]byte, error)
    // 并行下载大对象
    ParallelDownload(ctx context.Context, bucket, key string) ([]byte, error)
    // 列出对象
    List(ctx context.Context, bucket, prefix string) ([]ObjectInfo, error)
    // 删除对象
    Delete(ctx context.Context, bucket, key string) error
    // 复制对象
    Copy(ctx context.Context, srcBucket, srcKey, dstBucket, dstKey string) error
}

// ObjectInfo 对象信息
type ObjectInfo struct {
    Key          string
    Size         int64
    LastModified time.Time
    ETag         string
    StorageClass string
}

// S3Storage S3 对象存储实现
type S3Storage struct {
    client      *s3.Client
    partSize    int64
    concurrency int
}

func NewS3Storage(client *s3.Client) *S3Storage {
    return &S3Storage{
        client:      client,
        partSize:    100 << 20, // 100MB per part
        concurrency: 16,
    }
}

// MultipartUpload 分片上传大文件
func (s *S3Storage) MultipartUpload(
    ctx context.Context,
    bucket, key string,
    reader io.Reader,
    size int64,
) error {
    // 创建分片上传
    createResp, err := s.client.CreateMultipartUpload(ctx, &s3.CreateMultipartUploadInput{
        Bucket:       aws.String(bucket),
        Key:          aws.String(key),
        StorageClass: types.StorageClassStandardIa, // 检查点使用低频存储
    })
    if err != nil {
        return fmt.Errorf("create multipart upload: %w", err)
    }

    uploadID := *createResp.UploadId

    // 计算分片
    numParts := (size + s.partSize - 1) / s.partSize
    parts := make([]types.CompletedPart, numParts)

    var wg sync.WaitGroup
    errCh := make(chan error, numParts)
    sem := make(chan struct{}, s.concurrency)

    var mu sync.Mutex
    partNum := int32(1)

    for {
        // 读取一个分片
        buf := make([]byte, s.partSize)
        n, readErr := io.ReadFull(reader, buf)
        if n == 0 {
            break
        }

        currentPart := partNum
        partNum++
        data := buf[:n]

        wg.Add(1)
        go func(pn int32, d []byte) {
            defer wg.Done()
            sem <- struct{}{}
            defer func() { <-sem }()

            // 计算 MD5
            hash := md5.Sum(d)
            md5Str := hex.EncodeToString(hash[:])

            // 上传分片
            uploadResp, err := s.client.UploadPart(ctx, &s3.UploadPartInput{
                Bucket:     aws.String(bucket),
                Key:        aws.String(key),
                UploadId:   aws.String(uploadID),
                PartNumber: aws.Int32(pn),
                Body:       bytes.NewReader(d),
                ContentMD5: aws.String(md5Str),
            })
            if err != nil {
                errCh <- fmt.Errorf("upload part %d: %w", pn, err)
                return
            }

            mu.Lock()
            parts[pn-1] = types.CompletedPart{
                ETag:       uploadResp.ETag,
                PartNumber: aws.Int32(pn),
            }
            mu.Unlock()
        }(currentPart, data)

        if readErr == io.EOF || readErr == io.ErrUnexpectedEOF {
            break
        }
        if readErr != nil {
            // 取消上传
            s.abortMultipartUpload(ctx, bucket, key, uploadID)
            return fmt.Errorf("read data: %w", readErr)
        }
    }

    wg.Wait()
    close(errCh)

    // 检查错误
    for err := range errCh {
        if err != nil {
            s.abortMultipartUpload(ctx, bucket, key, uploadID)
            return err
        }
    }

    // 完成上传
    _, err = s.client.CompleteMultipartUpload(ctx, &s3.CompleteMultipartUploadInput{
        Bucket:   aws.String(bucket),
        Key:      aws.String(key),
        UploadId: aws.String(uploadID),
        MultipartUpload: &types.CompletedMultipartUpload{
            Parts: parts[:partNum-1],
        },
    })
    if err != nil {
        return fmt.Errorf("complete multipart upload: %w", err)
    }

    return nil
}

func (s *S3Storage) abortMultipartUpload(ctx context.Context, bucket, key, uploadID string) {
    s.client.AbortMultipartUpload(ctx, &s3.AbortMultipartUploadInput{
        Bucket:   aws.String(bucket),
        Key:      aws.String(key),
        UploadId: aws.String(uploadID),
    })
}

// ParallelDownload 并行下载大文件
func (s *S3Storage) ParallelDownload(ctx context.Context, bucket, key string) ([]byte, error) {
    // 获取对象大小
    headResp, err := s.client.HeadObject(ctx, &s3.HeadObjectInput{
        Bucket: aws.String(bucket),
        Key:    aws.String(key),
    })
    if err != nil {
        return nil, fmt.Errorf("head object: %w", err)
    }

    size := *headResp.ContentLength
    if size <= s.partSize {
        // 小文件直接下载
        return s.Download(ctx, bucket, key)
    }

    // 并行下载
    data := make([]byte, size)
    numParts := (size + s.partSize - 1) / s.partSize

    var wg sync.WaitGroup
    errCh := make(chan error, numParts)
    sem := make(chan struct{}, s.concurrency)

    for i := int64(0); i < numParts; i++ {
        start := i * s.partSize
        end := start + s.partSize - 1
        if end >= size {
            end = size - 1
        }

        wg.Add(1)
        go func(s3Start, s3End int64) {
            defer wg.Done()
            sem <- struct{}{}
            defer func() { <-sem }()

            rangeStr := fmt.Sprintf("bytes=%d-%d", s3Start, s3End)

            resp, err := s.client.GetObject(ctx, &s3.GetObjectInput{
                Bucket: aws.String(bucket),
                Key:    aws.String(key),
                Range:  aws.String(rangeStr),
            })
            if err != nil {
                errCh <- fmt.Errorf("get range %s: %w", rangeStr, err)
                return
            }
            defer resp.Body.Close()

            _, err = io.ReadFull(resp.Body, data[s3Start:s3End+1])
            if err != nil {
                errCh <- fmt.Errorf("read range %s: %w", rangeStr, err)
            }
        }(start, end)
    }

    wg.Wait()
    close(errCh)

    for err := range errCh {
        if err != nil {
            return nil, err
        }
    }

    return data, nil
}

// Download 下载对象
func (s *S3Storage) Download(ctx context.Context, bucket, key string) ([]byte, error) {
    resp, err := s.client.GetObject(ctx, &s3.GetObjectInput{
        Bucket: aws.String(bucket),
        Key:    aws.String(key),
    })
    if err != nil {
        return nil, fmt.Errorf("get object: %w", err)
    }
    defer resp.Body.Close()

    return io.ReadAll(resp.Body)
}

// Upload 上传对象
func (s *S3Storage) Upload(ctx context.Context, bucket, key string, data []byte) error {
    _, err := s.client.PutObject(ctx, &s3.PutObjectInput{
        Bucket: aws.String(bucket),
        Key:    aws.String(key),
        Body:   bytes.NewReader(data),
    })
    return err
}

// List 列出对象
func (s *S3Storage) List(ctx context.Context, bucket, prefix string) ([]ObjectInfo, error) {
    var objects []ObjectInfo

    paginator := s3.NewListObjectsV2Paginator(s.client, &s3.ListObjectsV2Input{
        Bucket: aws.String(bucket),
        Prefix: aws.String(prefix),
    })

    for paginator.HasMorePages() {
        page, err := paginator.NextPage(ctx)
        if err != nil {
            return nil, fmt.Errorf("list objects: %w", err)
        }

        for _, obj := range page.Contents {
            objects = append(objects, ObjectInfo{
                Key:          *obj.Key,
                Size:         *obj.Size,
                LastModified: *obj.LastModified,
                ETag:         *obj.ETag,
                StorageClass: string(obj.StorageClass),
            })
        }
    }

    return objects, nil
}

// Delete 删除对象
func (s *S3Storage) Delete(ctx context.Context, bucket, key string) error {
    _, err := s.client.DeleteObject(ctx, &s3.DeleteObjectInput{
        Bucket: aws.String(bucket),
        Key:    aws.String(key),
    })
    return err
}

// Copy 复制对象
func (s *S3Storage) Copy(ctx context.Context, srcBucket, srcKey, dstBucket, dstKey string) error {
    _, err := s.client.CopyObject(ctx, &s3.CopyObjectInput{
        Bucket:     aws.String(dstBucket),
        Key:        aws.String(dstKey),
        CopySource: aws.String(fmt.Sprintf("%s/%s", srcBucket, srcKey)),
    })
    return err
}

// checkpoint_manager.go
package checkpoint

import (
    "context"
    "encoding/json"
    "fmt"
    "path/filepath"
    "sort"
    "sync"
    "time"
)

// CheckpointManager 检查点管理器
type CheckpointManager struct {
    // 存储后端
    hotStorage  Storage // 本地/NVMe
    warmStorage Storage // 并行文件系统
    coldStorage Storage // 对象存储

    // 元数据存储
    metaStore MetadataStore

    // 配置
    config CheckpointConfig

    // 运行时状态
    mu             sync.RWMutex
    activeUploads  map[string]*UploadTask

    // 指标
    metrics *CheckpointMetrics
}

// CheckpointConfig 检查点配置
type CheckpointConfig struct {
    // 保存策略
    SaveInterval        int           // 每 N 步保存
    SaveIntervalMinutes int           // 或每 N 分钟保存

    // 保留策略
    HotRetention  int // Hot 存储保留数量
    WarmRetention int // Warm 存储保留数量
    ColdRetention int // Cold 存储保留数量（里程碑）

    // 存储路径
    HotPath  string
    WarmPath string
    ColdBucket string

    // 性能配置
    AsyncUpload     bool
    CompressionAlgo string // "none", "lz4", "zstd"
    ParallelIO      bool

    // 数据完整性
    VerifyChecksum  bool

    // 生命周期
    MilestoneSteps []int // 标记为里程碑的步数
}

// Checkpoint 检查点元数据
type Checkpoint struct {
    ID        string    `json:"id"`
    JobID     string    `json:"job_id"`
    Step      int64     `json:"step"`
    Epoch     int       `json:"epoch"`
    CreatedAt time.Time `json:"created_at"`

    // 存储位置
    StorageTier StorageTier `json:"storage_tier"`
    Path        string      `json:"path"`

    // 文件信息
    Files       []CheckpointFile `json:"files"`
    TotalSize   int64            `json:"total_size"`
    Checksum    string           `json:"checksum"`

    // 训练状态
    TrainingState TrainingState `json:"training_state"`

    // 标签
    IsMilestone bool              `json:"is_milestone"`
    Tags        map[string]string `json:"tags"`
}

// CheckpointFile 检查点文件
type CheckpointFile struct {
    Name     string `json:"name"`
    Size     int64  `json:"size"`
    Checksum string `json:"checksum"`
    Type     string `json:"type"` // "model", "optimizer", "scheduler", "rng"
}

// TrainingState 训练状态
type TrainingState struct {
    GlobalStep    int64             `json:"global_step"`
    Epoch         int               `json:"epoch"`
    LearningRate  float64           `json:"learning_rate"`
    Loss          float64           `json:"loss"`
    Metrics       map[string]float64 `json:"metrics"`
    RandomState   []byte            `json:"random_state"`
}

type StorageTier string

const (
    StorageTierHot  StorageTier = "hot"
    StorageTierWarm StorageTier = "warm"
    StorageTierCold StorageTier = "cold"
)

// NewCheckpointManager 创建检查点管理器
func NewCheckpointManager(
    hotStorage, warmStorage, coldStorage Storage,
    metaStore MetadataStore,
    config CheckpointConfig,
) *CheckpointManager {
    return &CheckpointManager{
        hotStorage:    hotStorage,
        warmStorage:   warmStorage,
        coldStorage:   coldStorage,
        metaStore:     metaStore,
        config:        config,
        activeUploads: make(map[string]*UploadTask),
        metrics:       NewCheckpointMetrics(),
    }
}

// Save 保存检查点
func (m *CheckpointManager) Save(
    ctx context.Context,
    jobID string,
    step int64,
    epoch int,
    state *TrainingState,
    data map[string][]byte, // 文件名 -> 数据
) (*Checkpoint, error) {

    startTime := time.Now()

    // 创建检查点元数据
    checkpoint := &Checkpoint{
        ID:            fmt.Sprintf("%s-step-%d", jobID, step),
        JobID:         jobID,
        Step:          step,
        Epoch:         epoch,
        CreatedAt:     time.Now(),
        StorageTier:   StorageTierHot,
        TrainingState: *state,
        IsMilestone:   m.isMilestone(step),
        Tags:          make(map[string]string),
    }

    // 计算检查点路径
    checkpoint.Path = m.checkpointPath(jobID, step)

    // 处理每个文件
    var files []CheckpointFile
    var totalSize int64

    for name, content := range data {
        // 可选压缩
        processedData := content
        if m.config.CompressionAlgo != "none" {
            var err error
            processedData, err = m.compress(content)
            if err != nil {
                return nil, fmt.Errorf("compress %s: %w", name, err)
            }
        }

        // 计算校验和
        checksum := m.calculateChecksum(processedData)

        // 保存到 Hot 存储
        filePath := filepath.Join(checkpoint.Path, name)
        if err := m.hotStorage.Write(ctx, filePath, processedData); err != nil {
            return nil, fmt.Errorf("write %s: %w", name, err)
        }

        files = append(files, CheckpointFile{
            Name:     name,
            Size:     int64(len(processedData)),
            Checksum: checksum,
            Type:     m.inferFileType(name),
        })
        totalSize += int64(len(processedData))
    }

    checkpoint.Files = files
    checkpoint.TotalSize = totalSize
    checkpoint.Checksum = m.calculateCheckpointChecksum(files)

    // 保存元数据
    if err := m.metaStore.SaveCheckpoint(ctx, checkpoint); err != nil {
        return nil, fmt.Errorf("save metadata: %w", err)
    }

    // 记录指标
    m.metrics.RecordSave(totalSize, time.Since(startTime))

    // 异步迁移到 Warm 存储
    if m.config.AsyncUpload {
        go m.migrateToWarm(context.Background(), checkpoint)
    }

    // 触发清理
    go m.cleanup(context.Background(), jobID)

    return checkpoint, nil
}

// Load 加载检查点
func (m *CheckpointManager) Load(
    ctx context.Context,
    checkpointID string,
) (map[string][]byte, *TrainingState, error) {

    startTime := time.Now()

    // 获取元数据
    checkpoint, err := m.metaStore.GetCheckpoint(ctx, checkpointID)
    if err != nil {
        return nil, nil, fmt.Errorf("get checkpoint metadata: %w", err)
    }

    // 选择存储后端
    storage := m.getStorage(checkpoint.StorageTier)

    // 加载所有文件
    data := make(map[string][]byte)

    for _, file := range checkpoint.Files {
        filePath := filepath.Join(checkpoint.Path, file.Name)

        content, err := storage.Read(ctx, filePath)
        if err != nil {
            return nil, nil, fmt.Errorf("read %s: %w", file.Name, err)
        }

        // 验证校验和
        if m.config.VerifyChecksum {
            if checksum := m.calculateChecksum(content); checksum != file.Checksum {
                return nil, nil, fmt.Errorf("checksum mismatch for %s", file.Name)
            }
        }

        // 解压
        if m.config.CompressionAlgo != "none" {
            content, err = m.decompress(content)
            if err != nil {
                return nil, nil, fmt.Errorf("decompress %s: %w", file.Name, err)
            }
        }

        data[file.Name] = content
    }

    // 记录指标
    m.metrics.RecordLoad(checkpoint.TotalSize, time.Since(startTime))

    return data, &checkpoint.TrainingState, nil
}

// LoadLatest 加载最新检查点
func (m *CheckpointManager) LoadLatest(
    ctx context.Context,
    jobID string,
) (map[string][]byte, *TrainingState, error) {

    // 获取最新检查点
    checkpoints, err := m.metaStore.ListCheckpoints(ctx, jobID)
    if err != nil {
        return nil, nil, fmt.Errorf("list checkpoints: %w", err)
    }

    if len(checkpoints) == 0 {
        return nil, nil, fmt.Errorf("no checkpoints found for job %s", jobID)
    }

    // 按步数排序，取最新
    sort.Slice(checkpoints, func(i, j int) bool {
        return checkpoints[i].Step > checkpoints[j].Step
    })

    return m.Load(ctx, checkpoints[0].ID)
}

// migrateToWarm 迁移到 Warm 存储
func (m *CheckpointManager) migrateToWarm(ctx context.Context, checkpoint *Checkpoint) {
    // 检查是否已在迁移
    m.mu.Lock()
    if _, exists := m.activeUploads[checkpoint.ID]; exists {
        m.mu.Unlock()
        return
    }
    task := &UploadTask{
        CheckpointID: checkpoint.ID,
        StartTime:    time.Now(),
        Status:       "in_progress",
    }
    m.activeUploads[checkpoint.ID] = task
    m.mu.Unlock()

    defer func() {
        m.mu.Lock()
        delete(m.activeUploads, checkpoint.ID)
        m.mu.Unlock()
    }()

    // 复制文件到 Warm 存储
    for _, file := range checkpoint.Files {
        srcPath := filepath.Join(checkpoint.Path, file.Name)
        dstPath := filepath.Join(m.config.WarmPath, checkpoint.JobID, checkpoint.ID, file.Name)

        data, err := m.hotStorage.Read(ctx, srcPath)
        if err != nil {
            m.metrics.RecordMigrationError("hot_to_warm")
            return
        }

        if err := m.warmStorage.Write(ctx, dstPath, data); err != nil {
            m.metrics.RecordMigrationError("hot_to_warm")
            return
        }
    }

    // 更新元数据
    checkpoint.StorageTier = StorageTierWarm
    checkpoint.Path = filepath.Join(m.config.WarmPath, checkpoint.JobID, checkpoint.ID)
    m.metaStore.SaveCheckpoint(ctx, checkpoint)

    // 如果是里程碑，还要迁移到 Cold 存储
    if checkpoint.IsMilestone {
        m.migrateToCold(ctx, checkpoint)
    }
}

// migrateToCold 迁移到 Cold 存储
func (m *CheckpointManager) migrateToCold(ctx context.Context, checkpoint *Checkpoint) {
    // 上传到对象存储
    for _, file := range checkpoint.Files {
        srcPath := filepath.Join(checkpoint.Path, file.Name)
        dstKey := fmt.Sprintf("%s/%s/%s", checkpoint.JobID, checkpoint.ID, file.Name)

        data, err := m.warmStorage.Read(ctx, srcPath)
        if err != nil {
            m.metrics.RecordMigrationError("warm_to_cold")
            return
        }

        if err := m.coldStorage.Write(ctx, dstKey, data); err != nil {
            m.metrics.RecordMigrationError("warm_to_cold")
            return
        }
    }

    // 更新元数据，记录 Cold 存储位置
    checkpoint.Tags["cold_path"] = fmt.Sprintf("s3://%s/%s/%s",
        m.config.ColdBucket, checkpoint.JobID, checkpoint.ID)
    m.metaStore.SaveCheckpoint(ctx, checkpoint)
}

// cleanup 清理过期检查点
func (m *CheckpointManager) cleanup(ctx context.Context, jobID string) {
    checkpoints, err := m.metaStore.ListCheckpoints(ctx, jobID)
    if err != nil {
        return
    }

    // 按存储层分组
    hotCheckpoints := filterByTier(checkpoints, StorageTierHot)
    warmCheckpoints := filterByTier(checkpoints, StorageTierWarm)

    // 按步数排序（降序）
    sort.Slice(hotCheckpoints, func(i, j int) bool {
        return hotCheckpoints[i].Step > hotCheckpoints[j].Step
    })
    sort.Slice(warmCheckpoints, func(i, j int) bool {
        return warmCheckpoints[i].Step > warmCheckpoints[j].Step
    })

    // 清理 Hot 存储中超出保留数量的检查点
    if len(hotCheckpoints) > m.config.HotRetention {
        for _, cp := range hotCheckpoints[m.config.HotRetention:] {
            m.deleteCheckpoint(ctx, cp, StorageTierHot)
        }
    }

    // 清理 Warm 存储中超出保留数量的非里程碑检查点
    nonMilestoneWarm := filterNonMilestone(warmCheckpoints)
    if len(nonMilestoneWarm) > m.config.WarmRetention {
        for _, cp := range nonMilestoneWarm[m.config.WarmRetention:] {
            m.deleteCheckpoint(ctx, cp, StorageTierWarm)
        }
    }
}

func (m *CheckpointManager) deleteCheckpoint(ctx context.Context, cp *Checkpoint, tier StorageTier) {
    storage := m.getStorage(tier)

    for _, file := range cp.Files {
        filePath := filepath.Join(cp.Path, file.Name)
        storage.Delete(ctx, filePath)
    }

    // 如果所有存储层都删除了，才删除元数据
    // 这里简化处理，实际需要跟踪每个层的删除状态
}

func (m *CheckpointManager) getStorage(tier StorageTier) Storage {
    switch tier {
    case StorageTierHot:
        return m.hotStorage
    case StorageTierWarm:
        return m.warmStorage
    case StorageTierCold:
        return m.coldStorage
    default:
        return m.hotStorage
    }
}

func (m *CheckpointManager) isMilestone(step int64) bool {
    for _, ms := range m.config.MilestoneSteps {
        if int64(ms) == step {
            return true
        }
    }
    // 或者每 N 步自动标记
    return step%10000 == 0
}

func (m *CheckpointManager) checkpointPath(jobID string, step int64) string {
    return filepath.Join(m.config.HotPath, jobID, fmt.Sprintf("step-%d", step))
}

func (m *CheckpointManager) compress(data []byte) ([]byte, error) {
    // 实现压缩逻辑
    return data, nil
}

func (m *CheckpointManager) decompress(data []byte) ([]byte, error) {
    // 实现解压逻辑
    return data, nil
}

func (m *CheckpointManager) calculateChecksum(data []byte) string {
    // 实现校验和计算
    return ""
}

func (m *CheckpointManager) calculateCheckpointChecksum(files []CheckpointFile) string {
    // 计算整个检查点的校验和
    return ""
}

func (m *CheckpointManager) inferFileType(name string) string {
    switch {
    case name == "model.pt" || name == "pytorch_model.bin":
        return "model"
    case name == "optimizer.pt":
        return "optimizer"
    case name == "scheduler.pt":
        return "scheduler"
    case name == "rng_state.pt":
        return "rng"
    default:
        return "other"
    }
}

type UploadTask struct {
    CheckpointID string
    StartTime    time.Time
    Status       string
}

func filterByTier(checkpoints []*Checkpoint, tier StorageTier) []*Checkpoint {
    var result []*Checkpoint
    for _, cp := range checkpoints {
        if cp.StorageTier == tier {
            result = append(result, cp)
        }
    }
    return result
}

func filterNonMilestone(checkpoints []*Checkpoint) []*Checkpoint {
    var result []*Checkpoint
    for _, cp := range checkpoints {
        if !cp.IsMilestone {
            result = append(result, cp)
        }
    }
    return result
}

// distributed_checkpoint.go
package checkpoint

import (
    "context"
    "fmt"
    "sync"
    "time"
)

// DistributedCheckpointManager 分布式检查点管理器
type DistributedCheckpointManager struct {
    // 本地检查点管理器
    localManager *CheckpointManager

    // 分布式协调
    coordinator Coordinator

    // 当前进程信息
    rank      int
    worldSize int

    // 配置
    config DistributedCheckpointConfig
}

// DistributedCheckpointConfig 分布式配置
type DistributedCheckpointConfig struct {
    // 分片策略
    ShardingStrategy string // "full", "sharded", "fsdp"

    // 聚合配置
    AggregateOnSave   bool // 保存时聚合到 rank 0
    AggregateOnLoad   bool // 加载时从聚合文件分发

    // 并行 IO
    ParallelRanks     int  // 并行保存的 rank 数

    // 超时
    BarrierTimeout    time.Duration
}

// DistributedCheckpoint 分布式检查点
type DistributedCheckpoint struct {
    *Checkpoint

    // 分片信息
    Shards []ShardInfo `json:"shards"`

    // 分布式状态
    WorldSize int `json:"world_size"`
}

// ShardInfo 分片信息
type ShardInfo struct {
    Rank       int      `json:"rank"`
    Files      []string `json:"files"`
    TotalSize  int64    `json:"total_size"`
    NodeID     string   `json:"node_id"`
}

// Coordinator 分布式协调接口
type Coordinator interface {
    Barrier(ctx context.Context, name string) error
    Broadcast(ctx context.Context, data []byte, root int) ([]byte, error)
    Gather(ctx context.Context, data []byte, root int) ([][]byte, error)
    AllGather(ctx context.Context, data []byte) ([][]byte, error)
}

// NewDistributedCheckpointManager 创建分布式检查点管理器
func NewDistributedCheckpointManager(
    localManager *CheckpointManager,
    coordinator Coordinator,
    rank, worldSize int,
    config DistributedCheckpointConfig,
) *DistributedCheckpointManager {
    return &DistributedCheckpointManager{
        localManager: localManager,
        coordinator:  coordinator,
        rank:         rank,
        worldSize:    worldSize,
        config:       config,
    }
}

// Save 分布式保存检查点
func (m *DistributedCheckpointManager) Save(
    ctx context.Context,
    jobID string,
    step int64,
    epoch int,
    state *TrainingState,
    localData map[string][]byte, // 本 rank 的数据
) (*DistributedCheckpoint, error) {

    // 阶段 1：准备（所有 rank 同步）
    if err := m.coordinator.Barrier(ctx, fmt.Sprintf("ckpt-prepare-%d", step)); err != nil {
        return nil, fmt.Errorf("barrier at prepare: %w", err)
    }

    var checkpoint *DistributedCheckpoint
    var saveErr error

    switch m.config.ShardingStrategy {
    case "full":
        // 每个 rank 保存完整检查点
        checkpoint, saveErr = m.saveFullReplicated(ctx, jobID, step, epoch, state, localData)

    case "sharded":
        // 每个 rank 保存自己的分片
        checkpoint, saveErr = m.saveSharded(ctx, jobID, step, epoch, state, localData)

    case "fsdp":
        // FSDP 风格：参数分片 + 优化器分片
        checkpoint, saveErr = m.saveFSDP(ctx, jobID, step, epoch, state, localData)
    }

    if saveErr != nil {
        return nil, saveErr
    }

    // 阶段 2：完成（所有 rank 同步）
    if err := m.coordinator.Barrier(ctx, fmt.Sprintf("ckpt-complete-%d", step)); err != nil {
        return nil, fmt.Errorf("barrier at complete: %w", err)
    }

    return checkpoint, nil
}

// saveSharded 分片保存
func (m *DistributedCheckpointManager) saveSharded(
    ctx context.Context,
    jobID string,
    step int64,
    epoch int,
    state *TrainingState,
    localData map[string][]byte,
) (*DistributedCheckpoint, error) {

    // 每个 rank 保存自己的分片
    // 使用 rank 作为文件名前缀
    shardedData := make(map[string][]byte)
    for name, data := range localData {
        shardedName := fmt.Sprintf("rank_%d_%s", m.rank, name)
        shardedData[shardedName] = data
    }

    // 调用本地管理器保存
    localCheckpoint, err := m.localManager.Save(ctx, jobID, step, epoch, state, shardedData)
    if err != nil {
        return nil, fmt.Errorf("save local shard: %w", err)
    }

    // 收集所有 rank 的分片信息
    shardInfo := ShardInfo{
        Rank:      m.rank,
        Files:     getFileNames(shardedData),
        TotalSize: calculateTotalSize(shardedData),
        NodeID:    getNodeID(),
    }

    shardInfoBytes, _ := json.Marshal(shardInfo)
    allShardInfoBytes, err := m.coordinator.AllGather(ctx, shardInfoBytes)
    if err != nil {
        return nil, fmt.Errorf("gather shard info: %w", err)
    }

    // 解析所有分片信息
    var allShards []ShardInfo
    for _, b := range allShardInfoBytes {
        var si ShardInfo
        json.Unmarshal(b, &si)
        allShards = append(allShards, si)
    }

    return &DistributedCheckpoint{
        Checkpoint: localCheckpoint,
        Shards:     allShards,
        WorldSize:  m.worldSize,
    }, nil
}

// saveFSDP FSDP 风格保存
func (m *DistributedCheckpointManager) saveFSDP(
    ctx context.Context,
    jobID string,
    step int64,
    epoch int,
    state *TrainingState,
    localData map[string][]byte,
) (*DistributedCheckpoint, error) {

    // FSDP 保存策略：
    // 1. 模型参数：每个 rank 保存自己的分片，或聚合后保存
    // 2. 优化器状态：每个 rank 保存自己的分片
    // 3. 其他状态：只在 rank 0 保存

    shardedData := make(map[string][]byte)

    for name, data := range localData {
        switch {
        case isModelParam(name):
            if m.config.AggregateOnSave {
                // 聚合模型参数到 rank 0
                aggregated, err := m.aggregateParams(ctx, data)
                if err != nil {
                    return nil, err
                }
                if m.rank == 0 {
                    shardedData[name] = aggregated
                }
            } else {
                // 分片保存
                shardedData[fmt.Sprintf("model_shard_%d.pt", m.rank)] = data
            }

        case isOptimizerState(name):
            // 优化器状态总是分片保存
            shardedData[fmt.Sprintf("optim_shard_%d.pt", m.rank)] = data

        default:
            // 其他状态只在 rank 0 保存
            if m.rank == 0 {
                shardedData[name] = data
            }
        }
    }

    return m.saveSharded(ctx, jobID, step, epoch, state, shardedData)
}

// saveFullReplicated 完全复制保存（每个 rank 保存完整）
func (m *DistributedCheckpointManager) saveFullReplicated(
    ctx context.Context,
    jobID string,
    step int64,
    epoch int,
    state *TrainingState,
    localData map[string][]byte,
) (*DistributedCheckpoint, error) {

    // 控制并发度：只有部分 rank 同时保存
    if m.rank < m.config.ParallelRanks {
        localCheckpoint, err := m.localManager.Save(ctx, jobID, step, epoch, state, localData)
        if err != nil {
            return nil, err
        }

        return &DistributedCheckpoint{
            Checkpoint: localCheckpoint,
            Shards: []ShardInfo{{
                Rank:      m.rank,
                TotalSize: localCheckpoint.TotalSize,
            }},
            WorldSize: m.worldSize,
        }, nil
    }

    // 其他 rank 等待
    return &DistributedCheckpoint{}, nil
}

// Load 分布式加载检查点
func (m *DistributedCheckpointManager) Load(
    ctx context.Context,
    checkpointID string,
) (map[string][]byte, *TrainingState, error) {

    // 阶段 1：准备
    if err := m.coordinator.Barrier(ctx, fmt.Sprintf("load-prepare-%s", checkpointID)); err != nil {
        return nil, nil, err
    }

    var data map[string][]byte
    var state *TrainingState
    var loadErr error

    switch m.config.ShardingStrategy {
    case "sharded", "fsdp":
        data, state, loadErr = m.loadSharded(ctx, checkpointID)
    case "full":
        data, state, loadErr = m.localManager.Load(ctx, checkpointID)
    }

    if loadErr != nil {
        return nil, nil, loadErr
    }

    // 阶段 2：完成
    if err := m.coordinator.Barrier(ctx, fmt.Sprintf("load-complete-%s", checkpointID)); err != nil {
        return nil, nil, err
    }

    return data, state, nil
}

// loadSharded 加载分片检查点
func (m *DistributedCheckpointManager) loadSharded(
    ctx context.Context,
    checkpointID string,
) (map[string][]byte, *TrainingState, error) {

    // 获取元数据
    metadata, err := m.localManager.metaStore.GetCheckpoint(ctx, checkpointID)
    if err != nil {
        return nil, nil, err
    }

    // 加载本 rank 的分片
    data := make(map[string][]byte)

    for _, file := range metadata.Files {
        // 检查是否是本 rank 的文件
        if isRankFile(file.Name, m.rank) || isSharedFile(file.Name) {
            content, err := m.localManager.hotStorage.Read(ctx,
                fmt.Sprintf("%s/%s", metadata.Path, file.Name))
            if err != nil {
                // 尝试从 warm 存储读取
                content, err = m.localManager.warmStorage.Read(ctx,
                    fmt.Sprintf("%s/%s", metadata.Path, file.Name))
                if err != nil {
                    return nil, nil, err
                }
            }

            // 去除 rank 前缀
            cleanName := removeRankPrefix(file.Name)
            data[cleanName] = content
        }
    }

    return data, &metadata.TrainingState, nil
}

// aggregateParams 聚合参数到 rank 0
func (m *DistributedCheckpointManager) aggregateParams(ctx context.Context, localParams []byte) ([]byte, error) {
    // 使用 Gather 收集到 rank 0
    allParams, err := m.coordinator.Gather(ctx, localParams, 0)
    if err != nil {
        return nil, err
    }

    if m.rank != 0 {
        return nil, nil
    }

    // 在 rank 0 上拼接所有参数
    var totalSize int
    for _, p := range allParams {
        totalSize += len(p)
    }

    result := make([]byte, 0, totalSize)
    for _, p := range allParams {
        result = append(result, p...)
    }

    return result, nil
}

// 辅助函数
func getFileNames(data map[string][]byte) []string {
    names := make([]string, 0, len(data))
    for name := range data {
        names = append(names, name)
    }
    return names
}

func calculateTotalSize(data map[string][]byte) int64 {
    var total int64
    for _, d := range data {
        total += int64(len(d))
    }
    return total
}

func getNodeID() string {
    // 返回节点标识
    return "node-1"
}

func isModelParam(name string) bool {
    return name == "model.pt" || name == "pytorch_model.bin"
}

func isOptimizerState(name string) bool {
    return name == "optimizer.pt"
}

func isRankFile(name string, rank int) bool {
    return fmt.Sprintf("rank_%d_", rank) == name[:len(fmt.Sprintf("rank_%d_", rank))]
}

func isSharedFile(name string) bool {
    // 共享文件不带 rank 前缀
    return len(name) > 0 && name[0] != 'r'
}

func removeRankPrefix(name string) string {
    // 去除 "rank_N_" 前缀
    for i := 0; i < len(name); i++ {
        if name[i] == '_' && i > 5 {
            return name[i+1:]
        }
    }
    return name
}