弹性 GPU 调度

// pkg/elastic/types.go
package elastic

import (
    "time"

    v1 "k8s.io/api/core/v1"
)

// PriorityClass GPU 工作负载优先级
type PriorityClass string

const (
    // PriorityProduction 生产级：不可抢占
    PriorityProduction PriorityClass = "production"
    // PriorityHighPriority 高优先级：可抢占低优先级
    PriorityHighPriority PriorityClass = "high-priority"
    // PriorityNormal 普通优先级
    PriorityNormal PriorityClass = "normal"
    // PriorityBestEffort 尽力而为：可被抢占
    PriorityBestEffort PriorityClass = "best-effort"
    // PrioritySpot 竞价实例：随时可被抢占
    PrioritySpot PriorityClass = "spot"
)

// PriorityValue 优先级数值（越大越高）
var PriorityValue = map[PriorityClass]int32{
    PriorityProduction:   1000000,
    PriorityHighPriority: 100000,
    PriorityNormal:       10000,
    PriorityBestEffort:   1000,
    PrioritySpot:         100,
}

// ElasticJob 弹性作业定义
type ElasticJob struct {
    // 基本信息
    Name      string
    Namespace string

    // GPU 需求
    MinGPUs int32 // 最小 GPU 数
    MaxGPUs int32 // 最大 GPU 数
    GPUType string // GPU 型号要求

    // 优先级
    Priority PriorityClass

    // 队列
    Queue string

    // 弹性配置
    ElasticConfig *ElasticConfig

    // 调度约束
    Constraints *SchedulingConstraints

    // 状态
    Status *ElasticJobStatus
}

// ElasticConfig 弹性配置
type ElasticConfig struct {
    // 是否启用弹性伸缩
    Enabled bool
    // 扩容策略
    ScaleUpPolicy *ScalePolicy
    // 缩容策略
    ScaleDownPolicy *ScalePolicy
    // 扩缩容冷却时间
    CooldownPeriod time.Duration
    // 是否允许被抢占
    Preemptible bool
    // 抢占恢复策略
    PreemptionRecovery PreemptionRecovery
}

// ScalePolicy 扩缩容策略
type ScalePolicy struct {
    // 策略类型
    Type ScalePolicyType
    // 步长
    Step int32
    // 百分比（用于 Percent 类型）
    Percent int32
    // 稳定窗口
    StabilizationWindow time.Duration
}

// ScalePolicyType 策略类型
type ScalePolicyType string

const (
    ScalePolicyStep    ScalePolicyType = "Step"
    ScalePolicyPercent ScalePolicyType = "Percent"
)

// PreemptionRecovery 抢占恢复策略
type PreemptionRecovery string

const (
    // RecoveryRestart 重启任务
    RecoveryRestart PreemptionRecovery = "Restart"
    // RecoveryResume 从检查点恢复
    RecoveryResume PreemptionRecovery = "Resume"
    // RecoveryGraceful 优雅降级（减少 GPU 继续运行）
    RecoveryGraceful PreemptionRecovery = "Graceful"
    // RecoveryFail 标记失败
    RecoveryFail PreemptionRecovery = "Fail"
)

// SchedulingConstraints 调度约束
type SchedulingConstraints struct {
    // 拓扑约束
    TopologyConstraint string
    // 节点选择
    NodeSelector map[string]string
    // 亲和性
    Affinity *v1.Affinity
    // 容忍
    Tolerations []v1.Toleration
    // 最大等待时间
    MaxWaitTime time.Duration
}

// ElasticJobStatus 作业状态
type ElasticJobStatus struct {
    // 当前阶段
    Phase JobPhase
    // 分配的 GPU 数量
    AllocatedGPUs int32
    // 运行的 Pod 数量
    RunningPods int32
    // 等待的 Pod 数量
    PendingPods int32
    // 开始时间
    StartTime *time.Time
    // 完成时间
    CompletionTime *time.Time
    // 上次扩缩容时间
    LastScaleTime *time.Time
    // 被抢占次数
    PreemptionCount int32
    // 条件
    Conditions []JobCondition
}

// JobPhase 作业阶段
type JobPhase string

const (
    JobPhasePending   JobPhase = "Pending"
    JobPhaseRunning   JobPhase = "Running"
    JobPhaseScaling   JobPhase = "Scaling"
    JobPhasePreempted JobPhase = "Preempted"
    JobPhaseSucceeded JobPhase = "Succeeded"
    JobPhaseFailed    JobPhase = "Failed"
)

// JobCondition 作业条件
type JobCondition struct {
    Type               string
    Status             v1.ConditionStatus
    LastTransitionTime time.Time
    Reason             string
    Message            string
}

// pkg/elastic/priority_queue.go
package elastic

import (
    "container/heap"
    "sync"
    "time"
)

// PriorityQueue 优先级队列
type PriorityQueue struct {
    mu      sync.RWMutex
    items   priorityHeap
    itemSet map[string]*QueueItem
}

// QueueItem 队列项
type QueueItem struct {
    Job            *ElasticJob
    EnqueueTime    time.Time
    Priority       int32
    Index          int // heap 索引
}

// priorityHeap 实现 heap.Interface
type priorityHeap []*QueueItem

func (h priorityHeap) Len() int { return len(h) }

func (h priorityHeap) Less(i, j int) bool {
    // 优先级高的排前面
    if h[i].Priority != h[j].Priority {
        return h[i].Priority > h[j].Priority
    }
    // 优先级相同，先入队的排前面（FIFO）
    return h[i].EnqueueTime.Before(h[j].EnqueueTime)
}

func (h priorityHeap) Swap(i, j int) {
    h[i], h[j] = h[j], h[i]
    h[i].Index = i
    h[j].Index = j
}

func (h *priorityHeap) Push(x interface{}) {
    item := x.(*QueueItem)
    item.Index = len(*h)
    *h = append(*h, item)
}

func (h *priorityHeap) Pop() interface{} {
    old := *h
    n := len(old)
    item := old[n-1]
    old[n-1] = nil
    item.Index = -1
    *h = old[0 : n-1]
    return item
}

// NewPriorityQueue 创建优先级队列
func NewPriorityQueue() *PriorityQueue {
    return &PriorityQueue{
        items:   make(priorityHeap, 0),
        itemSet: make(map[string]*QueueItem),
    }
}

// Enqueue 入队
func (q *PriorityQueue) Enqueue(job *ElasticJob) {
    q.mu.Lock()
    defer q.mu.Unlock()

    key := job.Namespace + "/" + job.Name

    // 已存在则更新
    if existing, ok := q.itemSet[key]; ok {
        existing.Job = job
        existing.Priority = PriorityValue[job.Priority]
        heap.Fix(&q.items, existing.Index)
        return
    }

    // 新增
    item := &QueueItem{
        Job:         job,
        EnqueueTime: time.Now(),
        Priority:    PriorityValue[job.Priority],
    }
    heap.Push(&q.items, item)
    q.itemSet[key] = item
}

// Dequeue 出队
func (q *PriorityQueue) Dequeue() *ElasticJob {
    q.mu.Lock()
    defer q.mu.Unlock()

    if q.items.Len() == 0 {
        return nil
    }

    item := heap.Pop(&q.items).(*QueueItem)
    delete(q.itemSet, item.Job.Namespace+"/"+item.Job.Name)
    return item.Job
}

// Peek 查看队首
func (q *PriorityQueue) Peek() *ElasticJob {
    q.mu.RLock()
    defer q.mu.RUnlock()

    if q.items.Len() == 0 {
        return nil
    }
    return q.items[0].Job
}

// Remove 移除
func (q *PriorityQueue) Remove(namespace, name string) *ElasticJob {
    q.mu.Lock()
    defer q.mu.Unlock()

    key := namespace + "/" + name
    item, ok := q.itemSet[key]
    if !ok {
        return nil
    }

    heap.Remove(&q.items, item.Index)
    delete(q.itemSet, key)
    return item.Job
}

// Len 长度
func (q *PriorityQueue) Len() int {
    q.mu.RLock()
    defer q.mu.RUnlock()
    return q.items.Len()
}

// GetJobsAbovePriority 获取高于指定优先级的作业
func (q *PriorityQueue) GetJobsAbovePriority(priority PriorityClass) []*ElasticJob {
    q.mu.RLock()
    defer q.mu.RUnlock()

    threshold := PriorityValue[priority]
    var jobs []*ElasticJob

    for _, item := range q.items {
        if item.Priority > threshold {
            jobs = append(jobs, item.Job)
        }
    }

    return jobs
}

// pkg/elastic/queue_manager.go
package elastic

import (
    "fmt"
    "sync"
)

// QueueConfig 队列配置
type QueueConfig struct {
    Name string
    // 权重（用于公平调度）
    Weight int32
    // 最大 GPU 配额
    MaxGPUs int32
    // 最小保证 GPU
    MinGPUs int32
    // 允许的优先级
    AllowedPriorities []PriorityClass
    // 是否允许抢占
    AllowPreemption bool
    // 是否允许借用资源
    AllowBorrowing bool
}

// QueueManager 队列管理器
type QueueManager struct {
    mu     sync.RWMutex
    queues map[string]*Queue
    // 全局配置
    defaultQueue string
}

// Queue 队列
type Queue struct {
    Config      QueueConfig
    PriorityQ   *PriorityQueue
    AllocatedGPUs int32
    UsedGPUs    int32
}

// NewQueueManager 创建队列管理器
func NewQueueManager(configs []QueueConfig) *QueueManager {
    qm := &QueueManager{
        queues:       make(map[string]*Queue),
        defaultQueue: "default",
    }

    for _, config := range configs {
        qm.queues[config.Name] = &Queue{
            Config:    config,
            PriorityQ: NewPriorityQueue(),
        }
    }

    return qm
}

// Submit 提交作业到队列
func (qm *QueueManager) Submit(job *ElasticJob) error {
    qm.mu.Lock()
    defer qm.mu.Unlock()

    queueName := job.Queue
    if queueName == "" {
        queueName = qm.defaultQueue
    }

    queue, ok := qm.queues[queueName]
    if !ok {
        return fmt.Errorf("queue %s not found", queueName)
    }

    // 检查优先级是否允许
    if !qm.isPriorityAllowed(queue, job.Priority) {
        return fmt.Errorf("priority %s not allowed in queue %s",
            job.Priority, queueName)
    }

    queue.PriorityQ.Enqueue(job)
    return nil
}

// SelectNext 选择下一个要调度的作业
func (qm *QueueManager) SelectNext(availableGPUs int32) *ElasticJob {
    qm.mu.Lock()
    defer qm.mu.Unlock()

    // 使用加权公平队列选择
    var selectedJob *ElasticJob
    var selectedQueue *Queue
    bestScore := float64(-1)

    for _, queue := range qm.queues {
        job := queue.PriorityQ.Peek()
        if job == nil {
            continue
        }

        // 检查配额
        if queue.UsedGPUs >= queue.Config.MaxGPUs {
            // 检查是否可以借用
            if !queue.Config.AllowBorrowing {
                continue
            }
        }

        // 检查资源是否足够
        if int32(job.MinGPUs) > availableGPUs {
            continue
        }

        // 计算公平调度分数
        score := qm.calculateFairScore(queue, job)
        if score > bestScore {
            bestScore = score
            selectedJob = job
            selectedQueue = queue
        }
    }

    if selectedJob != nil {
        selectedQueue.PriorityQ.Dequeue()
    }

    return selectedJob
}

// calculateFairScore 计算公平调度分数
func (qm *QueueManager) calculateFairScore(queue *Queue, job *ElasticJob) float64 {
    // 基础分数 = 权重 / (已使用 + 1)
    baseScore := float64(queue.Config.Weight) / float64(queue.UsedGPUs+1)

    // 优先级加成
    priorityBonus := float64(PriorityValue[job.Priority]) / 1000000.0

    // 等待时间加成
    // waitingBonus := time.Since(job.Status.StartTime).Minutes() / 60.0

    return baseScore + priorityBonus
}

// isPriorityAllowed 检查优先级是否允许
func (qm *QueueManager) isPriorityAllowed(queue *Queue, priority PriorityClass) bool {
    if len(queue.Config.AllowedPriorities) == 0 {
        return true
    }
    for _, p := range queue.Config.AllowedPriorities {
        if p == priority {
            return true
        }
    }
    return false
}

// GetQueueStatus 获取队列状态
func (qm *QueueManager) GetQueueStatus() map[string]*QueueStatus {
    qm.mu.RLock()
    defer qm.mu.RUnlock()

    status := make(map[string]*QueueStatus)
    for name, queue := range qm.queues {
        status[name] = &QueueStatus{
            Name:           name,
            PendingJobs:    queue.PriorityQ.Len(),
            AllocatedGPUs:  queue.AllocatedGPUs,
            UsedGPUs:       queue.UsedGPUs,
            MaxGPUs:        queue.Config.MaxGPUs,
            MinGPUs:        queue.Config.MinGPUs,
            Weight:         queue.Config.Weight,
        }
    }
    return status
}

// QueueStatus 队列状态
type QueueStatus struct {
    Name          string
    PendingJobs   int
    AllocatedGPUs int32
    UsedGPUs      int32
    MaxGPUs       int32
    MinGPUs       int32
    Weight        int32
}

// pkg/elastic/preemption.go
package elastic

import (
    "context"
    "fmt"
    "sort"
    "time"

    v1 "k8s.io/api/core/v1"
    metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
    "k8s.io/client-go/kubernetes"
)

// PreemptionManager 抢占管理器
type PreemptionManager struct {
    kubeClient   kubernetes.Interface
    stateManager *GPUStateManager
    // 抢占策略
    strategy PreemptionStrategy
    // 优雅终止时间
    gracePeriod time.Duration
    // 最小运行时间保护
    minRuntime time.Duration
}

// PreemptionStrategy 抢占策略
type PreemptionStrategy string

const (
    // StrategyLowestPriority 抢占最低优先级
    StrategyLowestPriority PreemptionStrategy = "LowestPriority"
    // StrategyNewestFirst 抢占最新启动的
    StrategyNewestFirst PreemptionStrategy = "NewestFirst"
    // StrategyMinimumVictims 最小受害者
    StrategyMinimumVictims PreemptionStrategy = "MinimumVictims"
    // StrategyMinimumGPUWaste 最小 GPU 浪费
    StrategyMinimumGPUWaste PreemptionStrategy = "MinimumGPUWaste"
)

// PreemptionCandidate 抢占候选者
type PreemptionCandidate struct {
    Pod       *v1.Pod
    Job       *ElasticJob
    GPUs      int32
    Priority  int32
    StartTime time.Time
    NodeName  string
}

// PreemptionPlan 抢占计划
type PreemptionPlan struct {
    // 请求者
    Requester *ElasticJob
    // 受害者列表
    Victims []*PreemptionCandidate
    // 释放的 GPU 数量
    FreedGPUs int32
    // 目标节点
    TargetNodes []string
}

// NewPreemptionManager 创建抢占管理器
func NewPreemptionManager(kubeClient kubernetes.Interface,
    stateManager *GPUStateManager) *PreemptionManager {
    return &PreemptionManager{
        kubeClient:   kubeClient,
        stateManager: stateManager,
        strategy:     StrategyLowestPriority,
        gracePeriod:  30 * time.Second,
        minRuntime:   5 * time.Minute,
    }
}

// FindPreemptionCandidates 找到抢占候选者
func (pm *PreemptionManager) FindPreemptionCandidates(ctx context.Context,
    requester *ElasticJob, neededGPUs int32) (*PreemptionPlan, error) {

    requesterPriority := PriorityValue[requester.Priority]

    // 获取所有运行中的 GPU Pod
    runningPods, err := pm.getGPUPods(ctx)
    if err != nil {
        return nil, err
    }

    // 筛选可抢占的候选者
    var candidates []*PreemptionCandidate
    for _, pod := range runningPods {
        candidate := pm.evaluateCandidate(pod, requesterPriority)
        if candidate != nil {
            candidates = append(candidates, candidate)
        }
    }

    if len(candidates) == 0 {
        return nil, fmt.Errorf("no preemptible candidates found")
    }

    // 根据策略选择受害者
    victims := pm.selectVictims(candidates, neededGPUs)
    if victims == nil {
        return nil, fmt.Errorf("cannot find enough GPUs to preempt")
    }

    // 计算释放的 GPU 总数
    var freedGPUs int32
    targetNodes := make(map[string]bool)
    for _, v := range victims {
        freedGPUs += v.GPUs
        targetNodes[v.NodeName] = true
    }

    nodes := make([]string, 0, len(targetNodes))
    for node := range targetNodes {
        nodes = append(nodes, node)
    }

    return &PreemptionPlan{
        Requester:   requester,
        Victims:     victims,
        FreedGPUs:   freedGPUs,
        TargetNodes: nodes,
    }, nil
}

// evaluateCandidate 评估候选者
func (pm *PreemptionManager) evaluateCandidate(pod *v1.Pod,
    requesterPriority int32) *PreemptionCandidate {

    // 检查优先级
    podPriority := int32(0)
    if pod.Spec.Priority != nil {
        podPriority = *pod.Spec.Priority
    }

    // 只能抢占优先级更低的
    if podPriority >= requesterPriority {
        return nil
    }

    // 检查是否可抢占
    if pm.isProtected(pod) {
        return nil
    }

    // 获取 GPU 数量
    gpus := pm.getPodGPUs(pod)
    if gpus == 0 {
        return nil
    }

    // 检查最小运行时间
    startTime := pod.Status.StartTime
    if startTime != nil {
        runtime := time.Since(startTime.Time)
        if runtime < pm.minRuntime {
            return nil
        }
    }

    return &PreemptionCandidate{
        Pod:       pod,
        GPUs:      gpus,
        Priority:  podPriority,
        StartTime: startTime.Time,
        NodeName:  pod.Spec.NodeName,
    }
}

// selectVictims 选择受害者
func (pm *PreemptionManager) selectVictims(candidates []*PreemptionCandidate,
    neededGPUs int32) []*PreemptionCandidate {

    switch pm.strategy {
    case StrategyLowestPriority:
        return pm.selectLowestPriority(candidates, neededGPUs)
    case StrategyNewestFirst:
        return pm.selectNewestFirst(candidates, neededGPUs)
    case StrategyMinimumVictims:
        return pm.selectMinimumVictims(candidates, neededGPUs)
    case StrategyMinimumGPUWaste:
        return pm.selectMinimumWaste(candidates, neededGPUs)
    default:
        return pm.selectLowestPriority(candidates, neededGPUs)
    }
}

// selectLowestPriority 按优先级从低到高选择
func (pm *PreemptionManager) selectLowestPriority(candidates []*PreemptionCandidate,
    neededGPUs int32) []*PreemptionCandidate {

    // 按优先级升序排序
    sort.Slice(candidates, func(i, j int) bool {
        if candidates[i].Priority != candidates[j].Priority {
            return candidates[i].Priority < candidates[j].Priority
        }
        // 同优先级，选择启动时间较晚的
        return candidates[i].StartTime.After(candidates[j].StartTime)
    })

    var victims []*PreemptionCandidate
    var freedGPUs int32

    for _, c := range candidates {
        victims = append(victims, c)
        freedGPUs += c.GPUs
        if freedGPUs >= neededGPUs {
            return victims
        }
    }

    return nil // 无法满足需求
}

// selectNewestFirst 优先选择最新启动的
func (pm *PreemptionManager) selectNewestFirst(candidates []*PreemptionCandidate,
    neededGPUs int32) []*PreemptionCandidate {

    // 按启动时间降序排序（最新的在前）
    sort.Slice(candidates, func(i, j int) bool {
        return candidates[i].StartTime.After(candidates[j].StartTime)
    })

    var victims []*PreemptionCandidate
    var freedGPUs int32

    for _, c := range candidates {
        victims = append(victims, c)
        freedGPUs += c.GPUs
        if freedGPUs >= neededGPUs {
            return victims
        }
    }

    return nil
}

// selectMinimumVictims 选择最少数量的受害者
func (pm *PreemptionManager) selectMinimumVictims(candidates []*PreemptionCandidate,
    neededGPUs int32) []*PreemptionCandidate {

    // 按 GPU 数量降序排序
    sort.Slice(candidates, func(i, j int) bool {
        return candidates[i].GPUs > candidates[j].GPUs
    })

    var victims []*PreemptionCandidate
    var freedGPUs int32

    for _, c := range candidates {
        victims = append(victims, c)
        freedGPUs += c.GPUs
        if freedGPUs >= neededGPUs {
            return victims
        }
    }

    return nil
}

// selectMinimumWaste 选择最小浪费的组合
func (pm *PreemptionManager) selectMinimumWaste(candidates []*PreemptionCandidate,
    neededGPUs int32) []*PreemptionCandidate {

    // 使用动态规划找最接近 neededGPUs 的组合
    n := len(candidates)
    target := int(neededGPUs)

    // dp[i] 表示是否可以达到 GPU 数量 i
    dp := make([]bool, target+1)
    dp[0] = true
    parent := make([]int, target+1)
    for i := range parent {
        parent[i] = -1
    }

    for idx, c := range candidates {
        gpus := int(c.GPUs)
        for i := target; i >= gpus; i-- {
            if dp[i-gpus] && !dp[i] {
                dp[i] = true
                parent[i] = idx
            }
        }
    }

    // 找到最小的 >= neededGPUs 的值
    for i := target; i <= target+10 && i <= n*8; i++ {
        if dp[i] {
            // 回溯找出选择的候选者
            var victims []*PreemptionCandidate
            remaining := i
            for remaining > 0 && parent[remaining] >= 0 {
                idx := parent[remaining]
                victims = append(victims, candidates[idx])
                remaining -= int(candidates[idx].GPUs)
            }
            return victims
        }
    }

    // 回退到最小受害者策略
    return pm.selectMinimumVictims(candidates, neededGPUs)
}

// ExecutePreemption 执行抢占
func (pm *PreemptionManager) ExecutePreemption(ctx context.Context,
    plan *PreemptionPlan) error {

    for _, victim := range plan.Victims {
        // 发送抢占信号
        if err := pm.preemptPod(ctx, victim); err != nil {
            return fmt.Errorf("failed to preempt pod %s: %v",
                victim.Pod.Name, err)
        }
    }

    return nil
}

// preemptPod 抢占单个 Pod
func (pm *PreemptionManager) preemptPod(ctx context.Context,
    victim *PreemptionCandidate) error {

    pod := victim.Pod

    // 添加抢占注解
    if pod.Annotations == nil {
        pod.Annotations = make(map[string]string)
    }
    pod.Annotations["elastic.gpu/preempted"] = "true"
    pod.Annotations["elastic.gpu/preemption-time"] = time.Now().Format(time.RFC3339)

    // 更新 Pod
    _, err := pm.kubeClient.CoreV1().Pods(pod.Namespace).Update(ctx, pod, metav1.UpdateOptions{})
    if err != nil {
        return err
    }

    // 删除 Pod（带优雅终止）
    gracePeriod := int64(pm.gracePeriod.Seconds())
    return pm.kubeClient.CoreV1().Pods(pod.Namespace).Delete(ctx, pod.Name,
        metav1.DeleteOptions{
            GracePeriodSeconds: &gracePeriod,
        })
}

// isProtected 检查 Pod 是否受保护
func (pm *PreemptionManager) isProtected(pod *v1.Pod) bool {
    // 检查注解
    if pod.Annotations != nil {
        if v, ok := pod.Annotations["elastic.gpu/preemptible"]; ok && v == "false" {
            return true
        }
    }

    // 检查 PodDisruptionBudget
    // ...

    return false
}

// getPodGPUs 获取 Pod 的 GPU 数量
func (pm *PreemptionManager) getPodGPUs(pod *v1.Pod) int32 {
    var total int32
    for _, c := range pod.Spec.Containers {
        if c.Resources.Limits != nil {
            if gpu, ok := c.Resources.Limits["nvidia.com/gpu"]; ok {
                total += int32(gpu.Value())
            }
        }
    }
    return total
}

// getGPUPods 获取所有 GPU Pod
func (pm *PreemptionManager) getGPUPods(ctx context.Context) ([]*v1.Pod, error) {
    podList, err := pm.kubeClient.CoreV1().Pods("").List(ctx, metav1.ListOptions{
        FieldSelector: "status.phase=Running",
    })
    if err != nil {
        return nil, err
    }

    var gpuPods []*v1.Pod
    for i := range podList.Items {
        pod := &podList.Items[i]
        if pm.getPodGPUs(pod) > 0 {
            gpuPods = append(gpuPods, pod)
        }
    }

    return gpuPods, nil
}

// pkg/elastic/recovery.go
package elastic

import (
    "context"
    "fmt"
    "time"

    batchv1 "k8s.io/api/batch/v1"
    v1 "k8s.io/api/core/v1"
    metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// RecoveryManager 恢复管理器
type RecoveryManager struct {
    kubeClient   kubernetes.Interface
    checkpointer Checkpointer
}

// Checkpointer 检查点接口
type Checkpointer interface {
    // SaveCheckpoint 保存检查点
    SaveCheckpoint(ctx context.Context, job *ElasticJob) (*Checkpoint, error)
    // LoadCheckpoint 加载检查点
    LoadCheckpoint(ctx context.Context, job *ElasticJob) (*Checkpoint, error)
    // ListCheckpoints 列出检查点
    ListCheckpoints(ctx context.Context, job *ElasticJob) ([]*Checkpoint, error)
}

// Checkpoint 检查点
type Checkpoint struct {
    ID        string
    JobName   string
    Namespace string
    // 保存时间
    Timestamp time.Time
    // 检查点路径
    Path string
    // 训练步数
    Step int64
    // 元数据
    Metadata map[string]string
}

// RecoverJob 恢复作业
func (rm *RecoveryManager) RecoverJob(ctx context.Context, job *ElasticJob) error {
    strategy := job.ElasticConfig.PreemptionRecovery

    switch strategy {
    case RecoveryRestart:
        return rm.recoverByRestart(ctx, job)
    case RecoveryResume:
        return rm.recoverByResume(ctx, job)
    case RecoveryGraceful:
        return rm.recoverGracefully(ctx, job)
    case RecoveryFail:
        return rm.markFailed(ctx, job)
    default:
        return rm.recoverByRestart(ctx, job)
    }
}

// recoverByRestart 通过重启恢复
func (rm *RecoveryManager) recoverByRestart(ctx context.Context, job *ElasticJob) error {
    // 简单重新创建 Pod
    return rm.recreateJob(ctx, job)
}

// recoverByResume 从检查点恢复
func (rm *RecoveryManager) recoverByResume(ctx context.Context, job *ElasticJob) error {
    // 获取最新检查点
    checkpoint, err := rm.checkpointer.LoadCheckpoint(ctx, job)
    if err != nil {
        // 没有检查点，回退到重启
        return rm.recoverByRestart(ctx, job)
    }

    // 使用检查点恢复
    return rm.recreateJobWithCheckpoint(ctx, job, checkpoint)
}

// recoverGracefully 优雅降级恢复
func (rm *RecoveryManager) recoverGracefully(ctx context.Context, job *ElasticJob) error {
    // 检查是否支持弹性训练
    if !job.ElasticConfig.Enabled {
        return rm.recoverByRestart(ctx, job)
    }

    // 获取当前可用资源
    availableGPUs := rm.getAvailableGPUs(ctx)

    // 计算新的 GPU 数量
    newGPUs := availableGPUs
    if newGPUs < job.MinGPUs {
        // 资源不足，等待
        return fmt.Errorf("insufficient GPUs for graceful recovery")
    }

    // 以更少的 GPU 继续运行
    return rm.recreateJobWithGPUs(ctx, job, newGPUs)
}

// markFailed 标记失败
func (rm *RecoveryManager) markFailed(ctx context.Context, job *ElasticJob) error {
    job.Status.Phase = JobPhaseFailed
    return rm.updateJobStatus(ctx, job)
}

// recreateJob 重建作业
func (rm *RecoveryManager) recreateJob(ctx context.Context, job *ElasticJob) error {
    // 创建新的 Pod/Job
    k8sJob := rm.buildK8sJob(job, nil)
    _, err := rm.kubeClient.BatchV1().Jobs(job.Namespace).Create(ctx, k8sJob, metav1.CreateOptions{})
    return err
}

// recreateJobWithCheckpoint 使用检查点重建
func (rm *RecoveryManager) recreateJobWithCheckpoint(ctx context.Context,
    job *ElasticJob, checkpoint *Checkpoint) error {

    // 构建带检查点的 Job
    k8sJob := rm.buildK8sJob(job, checkpoint)

    // 添加检查点环境变量
    for i := range k8sJob.Spec.Template.Spec.Containers {
        k8sJob.Spec.Template.Spec.Containers[i].Env = append(
            k8sJob.Spec.Template.Spec.Containers[i].Env,
            v1.EnvVar{Name: "CHECKPOINT_PATH", Value: checkpoint.Path},
            v1.EnvVar{Name: "RESUME_STEP", Value: fmt.Sprintf("%d", checkpoint.Step)},
        )
    }

    _, err := rm.kubeClient.BatchV1().Jobs(job.Namespace).Create(ctx, k8sJob, metav1.CreateOptions{})
    return err
}

// recreateJobWithGPUs 使用指定 GPU 数量重建
func (rm *RecoveryManager) recreateJobWithGPUs(ctx context.Context,
    job *ElasticJob, gpus int32) error {

    // 更新 GPU 数量
    originalGPUs := job.Status.AllocatedGPUs
    job.Status.AllocatedGPUs = gpus

    // 构建 Job
    k8sJob := rm.buildK8sJob(job, nil)

    // 添加弹性训练环境变量
    for i := range k8sJob.Spec.Template.Spec.Containers {
        k8sJob.Spec.Template.Spec.Containers[i].Env = append(
            k8sJob.Spec.Template.Spec.Containers[i].Env,
            v1.EnvVar{Name: "ELASTIC_GPUS", Value: fmt.Sprintf("%d", gpus)},
            v1.EnvVar{Name: "ORIGINAL_GPUS", Value: fmt.Sprintf("%d", originalGPUs)},
            v1.EnvVar{Name: "ELASTIC_RESIZE", Value: "true"},
        )
    }

    _, err := rm.kubeClient.BatchV1().Jobs(job.Namespace).Create(ctx, k8sJob, metav1.CreateOptions{})
    return err
}

// buildK8sJob 构建 K8s Job
func (rm *RecoveryManager) buildK8sJob(job *ElasticJob, checkpoint *Checkpoint) *batchv1.Job {
    gpus := job.Status.AllocatedGPUs
    if gpus == 0 {
        gpus = job.MinGPUs
    }

    parallelism := gpus // 每个 GPU 一个 Pod
    completions := parallelism

    return &batchv1.Job{
        ObjectMeta: metav1.ObjectMeta{
            Name:      job.Name,
            Namespace: job.Namespace,
            Labels: map[string]string{
                "elastic.gpu/job-name": job.Name,
                "elastic.gpu/queue":    job.Queue,
                "elastic.gpu/priority": string(job.Priority),
            },
            Annotations: map[string]string{
                "elastic.gpu/preemptible": fmt.Sprintf("%v", job.ElasticConfig.Preemptible),
            },
        },
        Spec: batchv1.JobSpec{
            Parallelism: &parallelism,
            Completions: &completions,
            Template: v1.PodTemplateSpec{
                ObjectMeta: metav1.ObjectMeta{
                    Labels: map[string]string{
                        "elastic.gpu/job-name": job.Name,
                    },
                },
                Spec: v1.PodSpec{
                    SchedulerName: "gpu-scheduler",
                    Containers: []v1.Container{
                        {
                            Name:  "main",
                            Image: "training-image:latest",
                            Resources: v1.ResourceRequirements{
                                Limits: v1.ResourceList{
                                    "nvidia.com/gpu": *resource.NewQuantity(1, resource.DecimalSI),
                                },
                                Requests: v1.ResourceList{
                                    "nvidia.com/gpu": *resource.NewQuantity(1, resource.DecimalSI),
                                },
                            },
                        },
                    },
                    RestartPolicy: v1.RestartPolicyNever,
                },
            },
        },
    }
}

// getAvailableGPUs 获取可用 GPU 数量
func (rm *RecoveryManager) getAvailableGPUs(ctx context.Context) int32 {
    // 从集群获取可用 GPU
    // 简化实现
    return 4
}

// updateJobStatus 更新作业状态
func (rm *RecoveryManager) updateJobStatus(ctx context.Context, job *ElasticJob) error {
    // 更新 CRD 状态
    return nil
}

// pkg/elastic/autoscaler.go
package elastic

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

// GPUAutoscaler GPU 自动伸缩器
type GPUAutoscaler struct {
    mu sync.RWMutex
    // 伸缩目标
    targets map[string]*ScaleTarget
    // 度量收集器
    metricsCollector MetricsCollector
    // 伸缩执行器
    scaleExecutor ScaleExecutor
    // 配置
    config *AutoscalerConfig
}

// ScaleTarget 伸缩目标
type ScaleTarget struct {
    Job       *ElasticJob
    CurrentGPUs int32
    DesiredGPUs int32
    // 上次伸缩时间
    LastScaleTime time.Time
    // 伸缩方向
    ScaleDirection ScaleDirection
    // 度量历史
    MetricsHistory []MetricsSample
}

// ScaleDirection 伸缩方向
type ScaleDirection string

const (
    ScaleUp   ScaleDirection = "up"
    ScaleDown ScaleDirection = "down"
    ScaleNone ScaleDirection = "none"
)

// MetricsSample 度量样本
type MetricsSample struct {
    Timestamp      time.Time
    GPUUtilization float64
    MemoryUsage    float64
    Throughput     float64
    QueueLength    int32
}

// MetricsCollector 度量收集器接口
type MetricsCollector interface {
    CollectMetrics(ctx context.Context, job *ElasticJob) (*MetricsSample, error)
}

// ScaleExecutor 伸缩执行器接口
type ScaleExecutor interface {
    ScaleUp(ctx context.Context, job *ElasticJob, delta int32) error
    ScaleDown(ctx context.Context, job *ElasticJob, delta int32) error
}

// AutoscalerConfig 自动伸缩配置
type AutoscalerConfig struct {
    // 扩容阈值（GPU 利用率）
    ScaleUpThreshold float64
    // 缩容阈值
    ScaleDownThreshold float64
    // 冷却时间
    CooldownPeriod time.Duration
    // 度量采集间隔
    MetricsInterval time.Duration
    // 稳定窗口
    StabilizationWindow time.Duration
    // 最大伸缩步长
    MaxScaleStep int32
}

// NewGPUAutoscaler 创建自动伸缩器
func NewGPUAutoscaler(metricsCollector MetricsCollector,
    scaleExecutor ScaleExecutor, config *AutoscalerConfig) *GPUAutoscaler {

    if config == nil {
        config = &AutoscalerConfig{
            ScaleUpThreshold:    0.9,
            ScaleDownThreshold:  0.3,
            CooldownPeriod:      5 * time.Minute,
            MetricsInterval:     30 * time.Second,
            StabilizationWindow: 3 * time.Minute,
            MaxScaleStep:        4,
        }
    }

    return &GPUAutoscaler{
        targets:          make(map[string]*ScaleTarget),
        metricsCollector: metricsCollector,
        scaleExecutor:    scaleExecutor,
        config:           config,
    }
}

// Register 注册伸缩目标
func (a *GPUAutoscaler) Register(job *ElasticJob) {
    a.mu.Lock()
    defer a.mu.Unlock()

    key := job.Namespace + "/" + job.Name
    a.targets[key] = &ScaleTarget{
        Job:         job,
        CurrentGPUs: job.Status.AllocatedGPUs,
        DesiredGPUs: job.Status.AllocatedGPUs,
    }
}

// Unregister 注销伸缩目标
func (a *GPUAutoscaler) Unregister(namespace, name string) {
    a.mu.Lock()
    defer a.mu.Unlock()
    delete(a.targets, namespace+"/"+name)
}

// Run 运行自动伸缩循环
func (a *GPUAutoscaler) Run(ctx context.Context) {
    ticker := time.NewTicker(a.config.MetricsInterval)
    defer ticker.Stop()

    for {
        select {
        case <-ctx.Done():
            return
        case <-ticker.C:
            a.reconcile(ctx)
        }
    }
}

// reconcile 协调循环
func (a *GPUAutoscaler) reconcile(ctx context.Context) {
    a.mu.Lock()
    targets := make([]*ScaleTarget, 0, len(a.targets))
    for _, t := range a.targets {
        targets = append(targets, t)
    }
    a.mu.Unlock()

    for _, target := range targets {
        a.processTarget(ctx, target)
    }
}

// processTarget 处理单个目标
func (a *GPUAutoscaler) processTarget(ctx context.Context, target *ScaleTarget) {
    job := target.Job

    // 检查是否启用弹性
    if !job.ElasticConfig.Enabled {
        return
    }

    // 收集度量
    metrics, err := a.metricsCollector.CollectMetrics(ctx, job)
    if err != nil {
        return
    }

    // 添加到历史
    target.MetricsHistory = append(target.MetricsHistory, *metrics)

    // 保留最近的样本
    windowSamples := int(a.config.StabilizationWindow / a.config.MetricsInterval)
    if len(target.MetricsHistory) > windowSamples {
        target.MetricsHistory = target.MetricsHistory[len(target.MetricsHistory)-windowSamples:]
    }

    // 计算伸缩决策
    decision := a.calculateDecision(target)

    // 执行伸缩
    if decision.Direction != ScaleNone {
        a.executeScale(ctx, target, decision)
    }
}

// ScaleDecision 伸缩决策
type ScaleDecision struct {
    Direction ScaleDirection
    Delta     int32
    Reason    string
}

// calculateDecision 计算伸缩决策
func (a *GPUAutoscaler) calculateDecision(target *ScaleTarget) *ScaleDecision {
    job := target.Job

    // 检查冷却期
    if time.Since(target.LastScaleTime) < a.config.CooldownPeriod {
        return &ScaleDecision{Direction: ScaleNone, Reason: "in cooldown period"}
    }

    // 计算平均利用率
    avgUtilization := a.calculateAverageUtilization(target.MetricsHistory)

    // 扩容判断
    if avgUtilization > a.config.ScaleUpThreshold {
        if target.CurrentGPUs >= job.MaxGPUs {
            return &ScaleDecision{Direction: ScaleNone, Reason: "at max GPUs"}
        }

        delta := a.calculateScaleUpDelta(target, avgUtilization)
        return &ScaleDecision{
            Direction: ScaleUp,
            Delta:     delta,
            Reason:    fmt.Sprintf("high utilization: %.2f%%", avgUtilization*100),
        }
    }

    // 缩容判断
    if avgUtilization < a.config.ScaleDownThreshold {
        if target.CurrentGPUs <= job.MinGPUs {
            return &ScaleDecision{Direction: ScaleNone, Reason: "at min GPUs"}
        }

        delta := a.calculateScaleDownDelta(target, avgUtilization)
        return &ScaleDecision{
            Direction: ScaleDown,
            Delta:     delta,
            Reason:    fmt.Sprintf("low utilization: %.2f%%", avgUtilization*100),
        }
    }

    return &ScaleDecision{Direction: ScaleNone, Reason: "utilization in normal range"}
}

// calculateAverageUtilization 计算平均利用率
func (a *GPUAutoscaler) calculateAverageUtilization(history []MetricsSample) float64 {
    if len(history) == 0 {
        return 0
    }

    var sum float64
    for _, sample := range history {
        sum += sample.GPUUtilization
    }
    return sum / float64(len(history))
}

// calculateScaleUpDelta 计算扩容数量
func (a *GPUAutoscaler) calculateScaleUpDelta(target *ScaleTarget, utilization float64) int32 {
    job := target.Job
    policy := job.ElasticConfig.ScaleUpPolicy

    var delta int32

    switch policy.Type {
    case ScalePolicyStep:
        delta = policy.Step
    case ScalePolicyPercent:
        delta = target.CurrentGPUs * policy.Percent / 100
        if delta < 1 {
            delta = 1
        }
    default:
        delta = 1
    }

    // 限制最大步长
    if delta > a.config.MaxScaleStep {
        delta = a.config.MaxScaleStep
    }

    // 不超过最大值
    if target.CurrentGPUs+delta > job.MaxGPUs {
        delta = job.MaxGPUs - target.CurrentGPUs
    }

    return delta
}

// calculateScaleDownDelta 计算缩容数量
func (a *GPUAutoscaler) calculateScaleDownDelta(target *ScaleTarget, utilization float64) int32 {
    job := target.Job
    policy := job.ElasticConfig.ScaleDownPolicy

    var delta int32

    switch policy.Type {
    case ScalePolicyStep:
        delta = policy.Step
    case ScalePolicyPercent:
        delta = target.CurrentGPUs * policy.Percent / 100
        if delta < 1 {
            delta = 1
        }
    default:
        delta = 1
    }

    // 限制最大步长
    if delta > a.config.MaxScaleStep {
        delta = a.config.MaxScaleStep
    }

    // 不低于最小值
    if target.CurrentGPUs-delta < job.MinGPUs {
        delta = target.CurrentGPUs - job.MinGPUs
    }

    return delta
}

// executeScale 执行伸缩
func (a *GPUAutoscaler) executeScale(ctx context.Context, target *ScaleTarget,
    decision *ScaleDecision) {

    var err error

    switch decision.Direction {
    case ScaleUp:
        err = a.scaleExecutor.ScaleUp(ctx, target.Job, decision.Delta)
    case ScaleDown:
        err = a.scaleExecutor.ScaleDown(ctx, target.Job, decision.Delta)
    }

    if err != nil {
        return
    }

    // 更新状态
    a.mu.Lock()
    defer a.mu.Unlock()

    switch decision.Direction {
    case ScaleUp:
        target.CurrentGPUs += decision.Delta
    case ScaleDown:
        target.CurrentGPUs -= decision.Delta
    }
    target.LastScaleTime = time.Now()
    target.ScaleDirection = decision.Direction
}