alwaysai/src/subcommands/model/get-package.ts

The alwaysAI command-line interface (CLI)

import { CliLeaf, CliStringInput, CliNumberInput } from '@alwaysai/alwayscli';
import { createWriteStream } from 'fs';
import pump = require('pump');
import { yesCliInput } from '../../cli-inputs';
import { checkUserIsLoggedInComponent } from '../../components/user';
import {
  ModelId,
  modelPackageCache,
  downloadModelPackageToCache
} from '../../core/model';
import { ALWAYSAI_SHOW_HIDDEN } from '../../environment';
import { CliRpcClient } from '../../infrastructure';
import { runWithSpinner } from '../../util';

export const modelGetPackage = CliLeaf({
  name: 'get-package',
  description: 'Get a model package',
  hidden: !ALWAYSAI_SHOW_HIDDEN,
  namedInputs: {
    id: CliStringInput({
      description: 'An alwaysAI model ID, e.g. alwaysai/mobilenet_ssd',
      required: true
    }),
    version: CliNumberInput({
      description: 'The version number of the model'
    }),
    yes: yesCliInput
  },
  async action(_, { yes, id, version: maybeVersion }) {
    const { publisher, name } = ModelId.parse(id);
    await checkUserIsLoggedInComponent({ yes });
    const version =
      maybeVersion || (await CliRpcClient().getModelVersion({ id })).version;
    const fileName = `${publisher}-${name}-${version}.tar.gz`;
    await runWithSpinner(
      async () => {
        if (!modelPackageCache.has(id, version)) {
          await downloadModelPackageToCache(id, version);
        }
        await new Promise<void>((resolve, reject) => {
          pump(
            modelPackageCache.read(id, version),
            createWriteStream(fileName),
            (err) => {
              if (err) {
                reject(err);
              } else {
                resolve();
              }
            }
          );
        });
      },
      [],
      `Getting ${fileName}`
    );
  }
});