modelgrid/ts/cli/cluster-handler.ts

import * as fs from 'node:fs/promises';
import { CLUSTER, PATHS } from '../constants.ts';
import type { IModelGridConfig } from '../interfaces/config.ts';
import { logger } from '../logger.ts';

export class ClusterHandler {
  public async status(): Promise<void> {
    const response = await this.request('/_cluster/status');
    if (!response) {
      return;
    }

    logger.log(JSON.stringify(response, null, 2));
  }

  public async nodes(): Promise<void> {
    const response = await this.request('/_cluster/nodes');
    if (!response) {
      return;
    }

    logger.log(JSON.stringify(response, null, 2));
  }

  public async models(): Promise<void> {
    const response = await this.request('/_cluster/status');
    if (!response || typeof response !== 'object' || !('models' in response)) {
      return;
    }

    logger.log(JSON.stringify((response as { models: unknown }).models, null, 2));
  }

  public async desired(): Promise<void> {
    const response = await this.request('/_cluster/desired');
    if (!response) {
      return;
    }

    logger.log(JSON.stringify(response, null, 2));
  }

  public async ensure(model: string): Promise<void> {
    if (!model) {
      logger.error('Model ID is required');
      return;
    }

    const response = await this.request('/_cluster/models/ensure', {
      method: 'POST',
      body: { model },
    });
    if (!response) {
      return;
    }

    logger.log(JSON.stringify(response, null, 2));
  }

  public async scale(model: string, desiredReplicas: number): Promise<void> {
    if (!model || Number.isNaN(desiredReplicas)) {
      logger.error('Model ID and desired replica count are required');
      return;
    }

    const response = await this.request('/_cluster/models/desired', {
      method: 'POST',
      body: { model, desiredReplicas },
    });
    if (!response) {
      return;
    }

    logger.log(JSON.stringify(response, null, 2));
  }

  public async clear(model: string): Promise<void> {
    if (!model) {
      logger.error('Model ID is required');
      return;
    }

    const response = await this.request('/_cluster/models/desired/remove', {
      method: 'POST',
      body: { model },
    });
    if (!response) {
      return;
    }

    logger.log(JSON.stringify(response, null, 2));
  }

  public async cordon(nodeName: string): Promise<void> {
    await this.updateNodeState('/_cluster/nodes/cordon', nodeName);
  }

  public async uncordon(nodeName: string): Promise<void> {
    await this.updateNodeState('/_cluster/nodes/uncordon', nodeName);
  }

  public async drain(nodeName: string): Promise<void> {
    await this.updateNodeState('/_cluster/nodes/drain', nodeName);
  }

  public async activate(nodeName: string): Promise<void> {
    await this.updateNodeState('/_cluster/nodes/activate', nodeName);
  }

  private async request(
    path: string,
    options: {
      method?: 'GET' | 'POST';
      body?: unknown;
    } = {},
  ): Promise<unknown | null> {
    const config = await this.readConfig();
    if (!config) {
      return null;
    }

    const endpoint = this.resolveEndpoint(config);
    const headers: Record<string, string> = {
      'Content-Type': 'application/json',
    };

    if (config.cluster.sharedSecret) {
      headers[CLUSTER.AUTH_HEADER_NAME] = config.cluster.sharedSecret;
    }

    try {
      const response = await fetch(`${endpoint}${path}`, {
        method: options.method || 'GET',
        headers,
        body: options.body ? JSON.stringify(options.body) : undefined,
      });

      if (!response.ok) {
        logger.error(`Cluster request failed: ${response.status} ${await response.text()}`);
        return null;
      }

      return await response.json();
    } catch (error) {
      logger.error(
        `Cluster request failed: ${error instanceof Error ? error.message : String(error)}`,
      );
      return null;
    }
  }

  private async readConfig(): Promise<IModelGridConfig | null> {
    try {
      return JSON.parse(await fs.readFile(PATHS.CONFIG_FILE, 'utf-8')) as IModelGridConfig;
    } catch (error) {
      logger.error(
        `Failed to read config: ${error instanceof Error ? error.message : String(error)}`,
      );
      return null;
    }
  }

  private resolveEndpoint(config: IModelGridConfig): string {
    if (config.cluster.controlPlaneUrl) {
      return config.cluster.controlPlaneUrl;
    }

    if (config.cluster.advertiseUrl) {
      return config.cluster.advertiseUrl;
    }

    const host = config.api.host === '0.0.0.0' ? '127.0.0.1' : config.api.host;
    return `http://${host}:${config.api.port}`;
  }

  private async updateNodeState(path: string, nodeName: string): Promise<void> {
    if (!nodeName) {
      logger.error('Node name is required');
      return;
    }

    const response = await this.request(path, {
      method: 'POST',
      body: { nodeName },
    });
    if (!response) {
      return;
    }

    logger.log(JSON.stringify(response, null, 2));
  }
}