modelgrid/ts/modelgrid.ts

/**
 * ModelGrid
 *
 * Main coordinator class for the ModelGrid system.
 */

import type { IModelGridConfig } from './interfaces/config.ts';
import { logger } from './logger.ts';
import { PATHS, VERSION } from './constants.ts';
import { Systemd } from './systemd.ts';
import { Daemon } from './daemon.ts';
import { GpuDetector } from './hardware/gpu-detector.ts';
import { SystemInfo } from './hardware/system-info.ts';
import { DriverManager } from './drivers/driver-manager.ts';
import { DockerManager } from './docker/docker-manager.ts';
import { ContainerManager } from './containers/container-manager.ts';
import { ClusterCoordinator } from './cluster/coordinator.ts';
import { ClusterManager } from './cluster/cluster-manager.ts';
import { ModelRegistry } from './models/registry.ts';
import { ModelLoader } from './models/loader.ts';
import { GpuHandler } from './cli/gpu-handler.ts';
import { ContainerHandler } from './cli/container-handler.ts';
import { ClusterHandler } from './cli/cluster-handler.ts';
import { ModelHandler } from './cli/model-handler.ts';
import { ConfigHandler } from './cli/config-handler.ts';
import { ServiceHandler } from './cli/service-handler.ts';
import * as fs from 'node:fs/promises';

/**
 * ModelGrid - Main application coordinator
 */
export class ModelGrid {
  private config?: IModelGridConfig;
  private systemd: Systemd;
  private daemon: Daemon;
  private gpuDetector: GpuDetector;
  private systemInfo: SystemInfo;
  private driverManager: DriverManager;
  private dockerManager: DockerManager;
  private containerManager: ContainerManager;
  private clusterManager: ClusterManager;
  private clusterCoordinator?: ClusterCoordinator;
  private modelRegistry: ModelRegistry;
  private modelLoader?: ModelLoader;

  // CLI Handlers
  private gpuHandler: GpuHandler;
  private containerHandler: ContainerHandler;
  private clusterHandler: ClusterHandler;
  private modelHandler: ModelHandler;
  private configHandler: ConfigHandler;
  private serviceHandler: ServiceHandler;

  constructor() {
    // Initialize core components
    this.gpuDetector = new GpuDetector();
    this.systemInfo = new SystemInfo();
    this.driverManager = new DriverManager();
    this.dockerManager = new DockerManager();
    this.containerManager = new ContainerManager();
    this.clusterManager = new ClusterManager();
    this.modelRegistry = new ModelRegistry();
    this.systemd = new Systemd();
    this.daemon = new Daemon(this);

    // Initialize CLI handlers
    this.gpuHandler = new GpuHandler();
    this.containerHandler = new ContainerHandler(this.containerManager);
    this.clusterHandler = new ClusterHandler();
    this.modelHandler = new ModelHandler(
      this.containerManager,
      this.getClusterCoordinator(),
      this.modelRegistry,
    );
    this.configHandler = new ConfigHandler();
    this.serviceHandler = new ServiceHandler(this);
  }

  /**
   * Load configuration from file
   */
  public async loadConfig(): Promise<void> {
    try {
      const configContent = await fs.readFile(PATHS.CONFIG_FILE, 'utf-8');
      this.config = this.normalizeConfig(
        JSON.parse(configContent) as Partial<IModelGridConfig> & {
          models?: {
            greenlistUrl?: string;
            autoPull?: boolean;
          } & Partial<IModelGridConfig['models']>;
        },
      );
      logger.dim(`Configuration loaded from ${PATHS.CONFIG_FILE}`);
    } catch (error) {
      if ((error as NodeJS.ErrnoException).code === 'ENOENT') {
        throw new Error(`Configuration file not found: ${PATHS.CONFIG_FILE}`);
      }
      throw error;
    }
  }

  /**
   * Save configuration to file
   */
  public async saveConfig(): Promise<void> {
    if (!this.config) {
      throw new Error('No configuration to save');
    }

    await fs.mkdir(PATHS.CONFIG_DIR, { recursive: true });
    await fs.writeFile(PATHS.CONFIG_FILE, JSON.stringify(this.config, null, 2));
    logger.dim(`Configuration saved to ${PATHS.CONFIG_FILE}`);
  }

  /**
   * Get current configuration
   */
  public getConfig(): IModelGridConfig | undefined {
    return this.config;
  }

  /**
   * Set configuration
   */
  public setConfig(config: IModelGridConfig): void {
    this.config = config;
  }

  /**
   * Get version string
   */
  public getVersion(): string {
    return VERSION;
  }

  /**
   * Get Systemd instance
   */
  public getSystemd(): Systemd {
    return this.systemd;
  }

  /**
   * Get Daemon instance
   */
  public getDaemon(): Daemon {
    return this.daemon;
  }

  /**
   * Get GPU Detector instance
   */
  public getGpuDetector(): GpuDetector {
    return this.gpuDetector;
  }

  /**
   * Get System Info instance
   */
  public getSystemInfo(): SystemInfo {
    return this.systemInfo;
  }

  /**
   * Get Driver Manager instance
   */
  public getDriverManager(): DriverManager {
    return this.driverManager;
  }

  /**
   * Get Docker Manager instance
   */
  public getDockerManager(): DockerManager {
    return this.dockerManager;
  }

  /**
   * Get Container Manager instance
   */
  public getContainerManager(): ContainerManager {
    return this.containerManager;
  }

  public getClusterManager(): ClusterManager {
    return this.clusterManager;
  }

  public getClusterCoordinator(): ClusterCoordinator {
    if (!this.clusterCoordinator) {
      this.clusterCoordinator = new ClusterCoordinator(
        this.clusterManager,
        this.containerManager,
        this.modelRegistry,
        this.getModelLoader(),
      );
    }

    return this.clusterCoordinator;
  }

  /**
   * Get Model Registry instance
   */
  public getModelRegistry(): ModelRegistry {
    return this.modelRegistry;
  }

  /**
   * Get Model Loader instance
   */
  public getModelLoader(): ModelLoader {
    if (!this.modelLoader) {
      this.modelLoader = new ModelLoader(this.modelRegistry, this.containerManager);
    }
    return this.modelLoader;
  }

  // CLI Handlers

  /**
   * Get GPU Handler
   */
  public getGpuHandler(): GpuHandler {
    return this.gpuHandler;
  }

  /**
   * Get Container Handler
   */
  public getContainerHandler(): ContainerHandler {
    return this.containerHandler;
  }

  /**
   * Get Model Handler
   */
  public getModelHandler(): ModelHandler {
    return this.modelHandler;
  }

  public getClusterHandler(): ClusterHandler {
    return this.clusterHandler;
  }

  /**
   * Get Config Handler
   */
  public getConfigHandler(): ConfigHandler {
    return this.configHandler;
  }

  /**
   * Get Service Handler
   */
  public getServiceHandler(): ServiceHandler {
    return this.serviceHandler;
  }

  /**
   * Initialize the ModelGrid system
   */
  public async initialize(): Promise<void> {
    // Load configuration
    await this.loadConfig();

    if (!this.config) {
      throw new Error('Failed to load configuration');
    }

    // Initialize containers from config
    for (const containerConfig of this.config.containers) {
      await this.containerManager.addContainer(containerConfig);
    }

    // Initialize model registry
    this.modelRegistry.setCatalogUrl(this.config.models.registryUrl);
    this.clusterManager.configure(this.config.cluster);
    await this.clusterManager.initialize();

    // Create model loader
    this.modelLoader = new ModelLoader(
      this.modelRegistry,
      this.containerManager,
      this.config.models.autoDeploy,
    );
    this.clusterCoordinator = new ClusterCoordinator(
      this.clusterManager,
      this.containerManager,
      this.modelRegistry,
      this.modelLoader,
    );

    logger.success('ModelGrid initialized');
  }

  private normalizeConfig(
    config: Partial<IModelGridConfig> & {
      models?: {
        greenlistUrl?: string;
        autoPull?: boolean;
      } & Partial<IModelGridConfig['models']>;
    },
  ): IModelGridConfig {
    const filteredContainers = (config.containers || []).filter(
      (container) => (container as { type?: string }).type !== 'ollama',
    );

    return {
      version: config.version || VERSION,
      api: {
        port: config.api?.port || 8080,
        host: config.api?.host || '0.0.0.0',
        apiKeys: config.api?.apiKeys || [],
        rateLimit: config.api?.rateLimit,
        cors: config.api?.cors ?? true,
        corsOrigins: config.api?.corsOrigins || ['*'],
      },
      docker: {
        networkName: config.docker?.networkName || 'modelgrid',
        runtime: config.docker?.runtime || 'docker',
        socketPath: config.docker?.socketPath,
      },
      gpus: {
        autoDetect: config.gpus?.autoDetect ?? true,
        assignments: config.gpus?.assignments || {},
      },
      containers: filteredContainers,
      models: {
        registryUrl: config.models?.registryUrl || config.models?.greenlistUrl ||
          'https://list.modelgrid.com/catalog/models.json',
        autoDeploy: config.models?.autoDeploy ?? config.models?.autoPull ?? true,
        defaultEngine: 'vllm',
        autoLoad: config.models?.autoLoad || [],
      },
      cluster: {
        enabled: config.cluster?.enabled ?? false,
        nodeName: config.cluster?.nodeName || 'modelgrid-local',
        role: config.cluster?.role || 'standalone',
        bindHost: config.cluster?.bindHost || '0.0.0.0',
        gossipPort: config.cluster?.gossipPort || 7946,
        sharedSecret: config.cluster?.sharedSecret,
        advertiseUrl: config.cluster?.advertiseUrl,
        controlPlaneUrl: config.cluster?.controlPlaneUrl,
        heartbeatIntervalMs: config.cluster?.heartbeatIntervalMs || 5000,
        seedNodes: config.cluster?.seedNodes || [],
      },
      checkInterval: config.checkInterval || 30000,
    };
  }

  /**
   * Shutdown the ModelGrid system
   */
  public async shutdown(): Promise<void> {
    logger.info('Shutting down ModelGrid...');

    // Stop all containers
    await this.containerManager.stopAll();

    logger.success('ModelGrid shutdown complete');
  }
}