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 { 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 { 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 modelRegistry: ModelRegistry;
  private modelLoader?: ModelLoader;

  // CLI Handlers
  private gpuHandler: GpuHandler;
  private containerHandler: ContainerHandler;
  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.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.modelHandler = new ModelHandler(this.containerManager, 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 = JSON.parse(configContent) as IModelGridConfig;
      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;
  }

  /**
   * 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;
  }

  /**
   * 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.setGreenlistUrl(this.config.models.greenlistUrl);

    // Create model loader
    this.modelLoader = new ModelLoader(
      this.modelRegistry,
      this.containerManager,
      this.config.models.autoPull,
    );

    logger.success('ModelGrid initialized');
  }

  /**
   * 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');
  }
}