alwaysai/src/subcommands/model/prune.ts

The alwaysAI command-line interface (CLI)

import * as path from 'path';
import * as rimraf from 'rimraf';
import * as logSymbols from 'log-symbols';
import {
  CliLeaf,
  CliStringArrayInput,
  CliStringInput,
  CliTerseError
} from '@alwaysai/alwayscli';
import { echo, logger } from '../../util';
import { ModelPackageDirectory } from '../../core/model';
import { existsSync } from 'fs';

function removeCache(modelPath: string, ignore: string[] | undefined) {
  let glob = '*';
  if (ignore) {
    glob = ignore.join('|');
    glob = `!(${glob})*`;
  }
  const removePath = path.join(modelPath, glob);
  logger.debug(removePath);
  rimraf.sync(removePath);
}

function errorIfIdNotInCache(modelPath: string, id: string) {
  if (!existsSync(modelPath)) {
    throw new CliTerseError(`Model ${id} not found in cache`);
  }
}

function errorIfIgnoredNotInCache(
  ignoreVersions: string[],
  modelPath: string,
  id: string
) {
  ignoreVersions.forEach((version) => {
    if (!existsSync(path.join(modelPath, `${version}.tar.gz`))) {
      throw new CliTerseError(
        `Ignore version ${version} for model ${id} not found in cache.`
      );
    }
  });
}

export const modelPrune = CliLeaf({
  name: 'prune',
  description: 'Prune model from your model cache',
  positionalInput: CliStringInput({
    placeholder: 'Model id ex: alwaysai/human-pose',
    required: true
  }),
  namedInputs: {
    ignore: CliStringArrayInput({
      description: 'Versions to ignore ex: 0.4 0.2'
    })
  },
  async action(id, opts) {
    const modelPath = path.join(ModelPackageDirectory(), id);
    errorIfIdNotInCache(modelPath, id);

    const ignoreVersions = opts['ignore'];
    if (ignoreVersions) {
      errorIfIgnoredNotInCache(ignoreVersions, modelPath, id);
    }

    removeCache(path.join(ModelPackageDirectory(), id), ignoreVersions);
    echo(`${logSymbols.success} Pruned ${id} from model cache`);
  }
});

export const testSuiteFuncs = {
  errorIfIdNotInCache,
  errorIfIgnoredNotInCache,
  removeCache
};