modelgrid/ts/containers/container-manager.ts

/**
 * Container Manager
 *
 * Orchestrates multiple AI model containers.
 */

import type {
  IContainerConfig,
  IContainerStatus,
  IContainerEndpoint,
  TContainerType,
} from '../interfaces/container.ts';
import { logger } from '../logger.ts';
import { DockerManager } from '../docker/docker-manager.ts';
import { BaseContainer } from './base-container.ts';
import { OllamaContainer } from './ollama.ts';
import { VllmContainer } from './vllm.ts';
import { TgiContainer } from './tgi.ts';

/**
 * Container Manager - orchestrates all containers
 */
export class ContainerManager {
  private containers: Map<string, BaseContainer>;
  private dockerManager: DockerManager;

  constructor() {
    this.containers = new Map();
    this.dockerManager = new DockerManager();
  }

  /**
   * Initialize container manager
   */
  public async initialize(): Promise<void> {
    // Ensure Docker is running
    if (!await this.dockerManager.isRunning()) {
      throw new Error('Docker is not running');
    }

    // Create network if it doesn't exist
    await this.dockerManager.createNetwork();
  }

  /**
   * Create a container instance from config
   */
  private createContainerInstance(config: IContainerConfig): BaseContainer {
    switch (config.type) {
      case 'ollama':
        return new OllamaContainer(config);
      case 'vllm':
        return new VllmContainer(config);
      case 'tgi':
        return new TgiContainer(config);
      default:
        throw new Error(`Unknown container type: ${config.type}`);
    }
  }

  /**
   * Add a container
   */
  public addContainer(config: IContainerConfig): BaseContainer {
    if (this.containers.has(config.id)) {
      throw new Error(`Container with ID ${config.id} already exists`);
    }

    const container = this.createContainerInstance(config);
    this.containers.set(config.id, container);
    return container;
  }

  /**
   * Remove a container
   */
  public async removeContainer(containerId: string): Promise<boolean> {
    const container = this.containers.get(containerId);
    if (!container) {
      return false;
    }

    await container.remove();
    this.containers.delete(containerId);
    return true;
  }

  /**
   * Get a container by ID
   */
  public getContainer(containerId: string): BaseContainer | undefined {
    return this.containers.get(containerId);
  }

  /**
   * Get all containers
   */
  public getAllContainers(): BaseContainer[] {
    return Array.from(this.containers.values());
  }

  /**
   * Load containers from configuration
   */
  public loadFromConfig(configs: IContainerConfig[]): void {
    this.containers.clear();
    for (const config of configs) {
      try {
        this.addContainer(config);
      } catch (error) {
        logger.warn(`Failed to load container ${config.id}: ${error instanceof Error ? error.message : String(error)}`);
      }
    }
  }

  /**
   * Start all containers
   */
  public async startAll(): Promise<Map<string, boolean>> {
    const results = new Map<string, boolean>();

    for (const [id, container] of this.containers) {
      if (!container.getConfig().autoStart) {
        continue;
      }

      try {
        const success = await container.start();
        results.set(id, success);
      } catch (error) {
        logger.error(`Failed to start container ${id}: ${error instanceof Error ? error.message : String(error)}`);
        results.set(id, false);
      }
    }

    return results;
  }

  /**
   * Stop all containers
   */
  public async stopAll(): Promise<Map<string, boolean>> {
    const results = new Map<string, boolean>();

    for (const [id, container] of this.containers) {
      try {
        const success = await container.stop();
        results.set(id, success);
      } catch (error) {
        logger.error(`Failed to stop container ${id}: ${error instanceof Error ? error.message : String(error)}`);
        results.set(id, false);
      }
    }

    return results;
  }

  /**
   * Get status of all containers
   */
  public async getAllStatus(): Promise<Map<string, IContainerStatus>> {
    const statuses = new Map<string, IContainerStatus>();

    for (const [id, container] of this.containers) {
      try {
        const status = await container.getStatus();
        statuses.set(id, status);
      } catch (error) {
        logger.warn(`Failed to get status for container ${id}: ${error instanceof Error ? error.message : String(error)}`);
      }
    }

    return statuses;
  }

  /**
   * Get available endpoints for a model
   */
  public async getEndpointsForModel(modelName: string): Promise<IContainerEndpoint[]> {
    const endpoints: IContainerEndpoint[] = [];

    for (const [_id, container] of this.containers) {
      try {
        const status = await container.getStatus();

        if (!status.running) {
          continue;
        }

        // Check if container has this model
        const models = await container.listModels();
        if (!models.includes(modelName)) {
          continue;
        }

        endpoints.push({
          containerId: container.getConfig().id,
          type: container.type,
          url: container.getEndpoint(),
          models,
          healthy: status.health === 'healthy',
          priority: 0, // Could be based on load
        });
      } catch {
        // Skip containers that fail to respond
      }
    }

    return endpoints;
  }

  /**
   * Find best container for a model
   */
  public async findContainerForModel(modelName: string): Promise<BaseContainer | null> {
    const endpoints = await this.getEndpointsForModel(modelName);

    // Filter to healthy endpoints
    const healthy = endpoints.filter((e) => e.healthy);
    if (healthy.length === 0) {
      return null;
    }

    // Return first healthy endpoint (could add load balancing)
    const endpoint = healthy[0];
    return this.containers.get(endpoint.containerId) || null;
  }

  /**
   * Get all available models across all containers
   */
  public async getAllAvailableModels(): Promise<Map<string, IContainerEndpoint[]>> {
    const modelMap = new Map<string, IContainerEndpoint[]>();

    for (const container of this.containers.values()) {
      try {
        const status = await container.getStatus();
        if (!status.running) continue;

        const models = await container.listModels();

        for (const model of models) {
          if (!modelMap.has(model)) {
            modelMap.set(model, []);
          }

          modelMap.get(model)!.push({
            containerId: container.getConfig().id,
            type: container.type,
            url: container.getEndpoint(),
            models,
            healthy: status.health === 'healthy',
            priority: 0,
          });
        }
      } catch {
        // Skip failed containers
      }
    }

    return modelMap;
  }

  /**
   * Pull a model to a specific container type
   */
  public async pullModel(
    modelName: string,
    containerType: TContainerType = 'ollama',
    containerId?: string,
  ): Promise<boolean> {
    // Find or create appropriate container
    let container: BaseContainer | undefined;

    if (containerId) {
      container = this.containers.get(containerId);
    } else {
      // Find first container of the specified type
      for (const c of this.containers.values()) {
        if (c.type === containerType) {
          container = c;
          break;
        }
      }
    }

    if (!container) {
      logger.error(`No ${containerType} container available to pull model`);
      return false;
    }

    return container.pullModel(modelName, (progress) => {
      const percent = progress.percent !== undefined ? ` (${progress.percent}%)` : '';
      logger.dim(`  ${progress.status}${percent}`);
    });
  }

  /**
   * Health check all containers
   */
  public async healthCheck(): Promise<Map<string, boolean>> {
    const results = new Map<string, boolean>();

    for (const [id, container] of this.containers) {
      try {
        const healthy = await container.isHealthy();
        results.set(id, healthy);
      } catch {
        results.set(id, false);
      }
    }

    return results;
  }

  /**
   * Print container status summary
   */
  public async printStatus(): Promise<void> {
    const statuses = await this.getAllStatus();

    if (statuses.size === 0) {
      logger.logBox('Containers', ['No containers configured'], 50, 'warning');
      return;
    }

    logger.logBoxTitle('Container Status', 70, 'info');

    for (const [id, status] of statuses) {
      const runningStr = status.running ? 'Running' : 'Stopped';
      const healthStr = status.health;
      const modelsStr = status.loadedModels.length > 0
        ? status.loadedModels.join(', ')
        : 'None';

      logger.logBoxLine(`${status.name} (${id})`);
      logger.logBoxLine(`  Type: ${status.type} | Status: ${runningStr} | Health: ${healthStr}`);
      logger.logBoxLine(`  Models: ${modelsStr}`);
      logger.logBoxLine(`  Endpoint: ${status.endpoint}`);

      if (status.gpuUtilization !== undefined) {
        logger.logBoxLine(`  GPU: ${status.gpuUtilization}% | Memory: ${status.memoryUsage || 0}MB`);
      }
      logger.logBoxLine('');
    }

    logger.logBoxEnd();
  }
}