alwaysai/src/core/app/app-install-model.ts

The alwaysAI command-line interface (CLI)

import { CliTerseError } from '@alwaysai/alwayscli';
import { dirname, posix } from 'path';
import {
  ModelId,
  modelPackageCache,
  downloadModelPackageToCache
} from '../model';
import { Spawner, RandomString, logger, stringifyError } from '../../util';
import { MODEL_JSON_FILE_NAME } from '../model/model-package-json-file';
import { ALWAYSAI_CLI_EXECUTABLE_NAME } from '../../constants';
import { APP_MODELS_DIRECTORY_NAME } from '../../paths';

export async function appInstallModel(
  target: Spawner,
  id: string,
  version: number
) {
  const { publisher, name } = ModelId.parse(id);
  const destinationDir = posix.join(APP_MODELS_DIRECTORY_NAME, publisher, name);
  let installedVersion: number | undefined = undefined;
  try {
    const parsed = await readModelJson(destinationDir);
    installedVersion = parsed.version;
  } catch (exception) {
    logger.warn(
      `Failed to read existing model config: ${stringifyError(exception)}`
    );
    // TODO finer-grained error handling
  }

  if (installedVersion !== version) {
    const tmpId = RandomString();
    const tmpDir = `${destinationDir}.${tmpId}.tmp`;
    await target.mkdirp(tmpDir);
    try {
      if (!modelPackageCache.has(id, version)) {
        await downloadModelPackageToCache(id, version);
      }
      try {
        const modelPackageStream = modelPackageCache.read(id, version);
        await target.untar(modelPackageStream, tmpDir);
      } catch {
        try {
          await modelPackageCache.remove(id, version);
        } catch {
          throw new CliTerseError(
            `Failed to install model due to corrupt cache. If the command continues to fail you may need to run ${ALWAYSAI_CLI_EXECUTABLE_NAME} model prune ${id} to remove the corrupt file`
          );
        }
      }
      const fileNames = await target.readdir(tmpDir);

      // Sanity check
      if (fileNames.length !== 1 || !fileNames[0]) {
        throw new Error('Expected package to contain single directory');
      }

      // The model json file in the package does not have the version number,
      // so we write it into the json file during install
      await updateModelJson(posix.join(tmpDir, fileNames[0]), (modelJson) => ({
        ...modelJson,
        version
      }));
      await target.rimraf(destinationDir);
      await target.mkdirp(dirname(destinationDir));
      await target.rename(
        target.resolvePath(tmpDir, fileNames[0]),
        destinationDir
      );
    } catch (exception) {
      try {
        // Attempt to delete new files since they may be corrupted
        await target.rimraf(tmpDir);
        await target.rimraf(destinationDir);
      } finally {
        // Do nothing
      }
      throw exception;
    } finally {
      await target.rimraf(tmpDir);
    }
  }

  async function readModelJson(dir: string) {
    const filePath = target.resolvePath(dir, MODEL_JSON_FILE_NAME);
    const output = await target.readFile(filePath);
    const parsed = JSON.parse(output);
    return parsed;
  }

  async function updateModelJson(dir: string, updater: (current: any) => any) {
    const parsed = await readModelJson(dir);
    const updated = updater(parsed);
    const filePath = target.resolvePath(dir, MODEL_JSON_FILE_NAME);
    const serialized = JSON.stringify(updated, null, 2);
    await target.writeFile(filePath, serialized);
  }
}