Sankofa/api/src/services/training-orchestrator.ts

import { Context } from '../types/context'
import { logger } from '../lib/logger'
import * as k8s from '@kubernetes/client-node'

export interface TrainingJob {
  id: string
  name: string
  status: string
  createdAt: Date
  namespace?: string
  jobName?: string
}

/**
 * Creates a Kubernetes job for model training
 */
export async function createTrainingJob(
  context: Context,
  input: {
    name: string
    image: string
    namespace?: string
    resources?: {
      cpu?: string
      memory?: string
      gpu?: number
    }
    env?: Record<string, string>
    command?: string[]
    args?: string[]
    timeout?: number
    restartPolicy?: 'Never' | 'OnFailure'
  }
): Promise<TrainingJob> {
  try {
    const kc = new k8s.KubeConfig()
    kc.loadFromDefault()
    const k8sBatchApi = kc.makeApiClient(k8s.BatchV1Api)
    const k8sCoreApi = kc.makeApiClient(k8s.CoreV1Api)

    const namespace = input.namespace || 'training'
    const jobName = `training-${input.name.toLowerCase().replace(/[^a-z0-9-]/g, '-')}-${Date.now()}`

    // Ensure namespace exists
    try {
      await k8sCoreApi.readNamespace(namespace)
    } catch (error) {
      // Namespace doesn't exist, create it
      const ns: k8s.V1Namespace = {
        apiVersion: 'v1',
        kind: 'Namespace',
        metadata: {
          name: namespace,
          labels: {
            'app.kubernetes.io/name': 'training',
          },
        },
      }
      await k8sCoreApi.createNamespace(ns)
    }

    // Create job
    const job: k8s.V1Job = {
      apiVersion: 'batch/v1',
      kind: 'Job',
      metadata: {
        name: jobName,
        namespace,
        labels: {
          'app': jobName,
          'component': 'training',
          'job-name': input.name,
        },
      },
      spec: {
        backoffLimit: 3,
        completions: 1,
        parallelism: 1,
        ttlSecondsAfterFinished: input.timeout || 3600, // Clean up after 1 hour by default
        template: {
          metadata: {
            labels: {
              app: jobName,
            },
          },
          spec: {
            restartPolicy: input.restartPolicy || 'Never',
            containers: [
              {
                name: 'training',
                image: input.image,
                command: input.command,
                args: input.args,
                env: Object.entries(input.env || {}).map(([key, value]) => ({
                  name: key,
                  value,
                })),
                resources: {
                  requests: {
                    cpu: input.resources?.cpu || '1000m',
                    memory: input.resources?.memory || '2Gi',
                  },
                  limits: {
                    cpu: input.resources?.cpu ? `${parseFloat(input.resources.cpu) * 2}${input.resources.cpu.slice(-1)}` : '4000m',
                    memory: input.resources?.memory ? `${parseFloat(input.resources.memory) * 2}${input.resources.memory.slice(-2)}` : '4Gi',
                  },
                },
              },
            ],
          },
        },
      },
    }

    // Add GPU support if specified
    if (input.resources?.gpu && input.resources.gpu > 0) {
      job.spec!.template!.spec!.containers![0].resources!.limits!['nvidia.com/gpu'] = input.resources.gpu.toString()
    }

    const jobResult = await k8sBatchApi.createNamespacedJob(namespace, job)

    const jobId = `${namespace}/${jobName}`

    return {
      id: jobId,
      name: input.name,
      status: 'PENDING',
      createdAt: new Date(),
      namespace,
      jobName,
    }
  } catch (error) {
    logger.error('Error creating training job', { error })
    throw new Error(`Failed to create training job: ${error instanceof Error ? error.message : 'Unknown error'}`)
  }
}