@toloka-platform/sdk 0.1.0 → 0.2.0-preview.2

package/src/cli/trainings.ts

@@ -0,0 +1,627 @@
+import { basename } from "node:path";
+import { readFile, stat } from "node:fs/promises";
+
+import { randomUUID } from "node:crypto";
+
+import { assert } from "es-toolkit";
+
+import type { PotatoClient, Training, TrainingCostEstimate, TrainingRun } from "../client";
+import type { BaseModelId, TrainingConfigurationInput } from "../../openapi/gen/models";
+
+type TrainingMethod = TrainingConfigurationInput["method"];
+type TrainingSource = TrainingConfigurationInput["source"];
+type TrainingSelection = Extract<TrainingConfigurationInput, { method: "sft" }>["training"];
+type SftConfigurationInput = Extract<TrainingConfigurationInput, { method: "sft" }>;
+type GistingConfigurationInput = Extract<TrainingConfigurationInput, { method: "gisting" }>;
+type ResolvedSelections = Pick<SftConfigurationInput, "training" | "validation">;
+
+export type TrainingDraftOptions = {
+  method?: TrainingMethod;
+  baseModelId?: BaseModelId;
+  modelVersionId?: string;
+  datasetId?: string;
+  validationDatasetId?: string;
+  trainingFilter?: string;
+  validationFilter?: string;
+  regime?: "low" | "medium" | "high";
+  epochs?: number;
+  evalSteps?: number;
+  promptFilePath?: string;
+  fetch?: typeof globalThis.fetch;
+};
+
+type TrainingApi = Pick<PotatoClient, "trainings">;
+type TrainingDraftApi = Pick<PotatoClient, "trainings" | "files" | "datasets">;
+type TrainingCostApi = Pick<PotatoClient, "trainings" | "waitFor">;
+const GISTING_PROMPT_MAX_BYTES = 256 * 1024;
+type TrainingIo = { readonly stdout: (text: string) => void };
+type TableRow = readonly string[];
+
+const formatTable = (rows: readonly TableRow[]): string => {
+  const widths = rows.reduce<number[]>(
+    (current, row) => row.map((cell, index) => Math.max(current[index] ?? 0, cell.length)),
+    [],
+  );
+  return rows
+    .map((row) =>
+      row
+        .map((cell, index) => cell.padEnd(widths[index] ?? cell.length))
+        .join("  ")
+        .trimEnd(),
+    )
+    .join("\n");
+};
+
+const writeOutput = (io: TrainingIo, json: boolean, value: unknown, formatted: string): void => {
+  io.stdout(json ? JSON.stringify(value, null, 2) : formatted);
+};
+
+const formatTraining = (training: Training): string =>
+  formatTable([
+    ["ID", training.id],
+    ["Project", training.projectId],
+    [
+      "Method",
+      training.configuration.method === "sft" ? "Supervised fine-tuning (SFT)" : "Prompt gisting",
+    ],
+    ["Output Model", training.modelId ?? "-"],
+    ["Archived at", training.archivedAt ?? "-"],
+    ["Created at", training.createdAt],
+    ["Updated at", training.updatedAt],
+  ]);
+
+const formatTrainingRun = (run: TrainingRun): string =>
+  formatTable([
+    ["ID", run.id],
+    ["Training", run.trainingId],
+    ["Status", run.status],
+    ["Blocked reason", run.blockedReason ?? "-"],
+    ["Message", run.message ?? "-"],
+    ["Error code", run.error?.code ?? "-"],
+    ["Billing status", run.billingStatus],
+    [
+      "Charged amount",
+      run.chargedAmount === null ? "-" : `$${run.chargedAmount} ${run.quote.currency}`,
+    ],
+    [
+      "Progress",
+      run.progress
+        ? `${run.progress.phase}${run.progress.step === null ? "" : ` ${run.progress.step}${run.progress.total === null ? "" : `/${run.progress.total}`}`}`
+        : "-",
+    ],
+    ["Model Version name", run.outputModelVersionName],
+    [
+      "Model Version revision",
+      run.outputModelVersionRevision === null ? "-" : String(run.outputModelVersionRevision),
+    ],
+    ["Future Model Version", run.outputModelVersionId],
+    ["Maximum charge", `$${run.quote.maximumCharge} ${run.quote.currency}`],
+    ["Created at", run.createdAt],
+  ]);
+
+type ReadyEstimateFile = {
+  status: "ready";
+  trainingId: string;
+  inputSetId: string;
+  fingerprint: string;
+};
+
+const parseReadyEstimateFile = (value: unknown): ReadyEstimateFile => {
+  assert(
+    typeof value === "object" && value !== null && !Array.isArray(value),
+    "The estimate file must contain a JSON object.",
+  );
+  assert(
+    "status" in value && value.status === "ready",
+    "The estimate file is not a ready Training Cost Estimate.",
+  );
+  assert(
+    "trainingId" in value && typeof value.trainingId === "string",
+    "The estimate file has no Training ID.",
+  );
+  assert(
+    "inputSetId" in value && typeof value.inputSetId === "string",
+    "The estimate file has no Input Set ID.",
+  );
+  assert(
+    "fingerprint" in value &&
+      typeof value.fingerprint === "string" &&
+      /^[0-9a-f]{64}$/.test(value.fingerprint),
+    "The estimate file has no valid fingerprint.",
+  );
+
+  return {
+    status: "ready",
+    trainingId: value.trainingId,
+    inputSetId: value.inputSetId,
+    fingerprint: value.fingerprint,
+  };
+};
+
+const formatCostEstimate = (estimate: TrainingCostEstimate): string => {
+  if (estimate.status === "blocked") {
+    return [
+      "Training cannot run with these inputs:",
+      ...estimate.issues.map((item) => `- ${item.message}`),
+    ].join("\n");
+  }
+
+  return formatTable([
+    ["Training", estimate.trainingId],
+    ["Training rows", String(estimate.inputs.training.report?.validCount ?? 0)],
+    ["Training tokens", String(estimate.inputs.training.report?.tokens.total ?? 0)],
+    ["Maximum GPU seconds", String(estimate.budget.maximumGpuSeconds)],
+    ["Maximum charge", `$${estimate.maximumCharge} ${estimate.currency}`],
+    ["Fingerprint", estimate.fingerprint],
+  ]);
+};
+
+const resolveTraining = async (
+  client: TrainingApi,
+  projectId: string,
+  nameOrId: string,
+): Promise<Training> => {
+  if (nameOrId.startsWith("trg.")) {
+    const training = await client.trainings.get({ trainingId: nameOrId });
+    assert(
+      training.projectId === projectId,
+      `Training ${nameOrId} was not found in project ${projectId}.`,
+    );
+    return training;
+  }
+
+  const matches: Training[] = [];
+  for await (const training of client.trainings.listAll({ projectId, limit: 100 })) {
+    if (training.name === nameOrId) matches.push(training);
+  }
+  assert(matches.length === 1, `Training ${nameOrId} was not found or is ambiguous.`);
+  return matches[0]!;
+};
+
+const parseFilterValue = (raw: string): string | number | boolean | null => {
+  if (raw === "true") return true;
+  if (raw === "false") return false;
+  if (raw === "null") return null;
+  const numeric = Number(raw);
+  return raw.trim() !== "" && Number.isFinite(numeric) ? numeric : raw;
+};
+
+const resolveSelection = async (
+  client: Pick<PotatoClient, "datasets">,
+  datasetId: string,
+  filter: string | undefined,
+): Promise<TrainingSelection> => {
+  assert(datasetId, "A Dataset is required when specifying a filter.");
+  if (filter === undefined) return { datasetId };
+  const separator = filter.indexOf("=");
+  assert(
+    separator > 0 && separator < filter.length - 1,
+    "Dataset filters must use FIELD=VALUE syntax.",
+  );
+  const fieldName = filter.slice(0, separator);
+  const rawValue = filter.slice(separator + 1);
+  const matches = [];
+  for await (const field of client.datasets.listAllFields({ datasetId })) {
+    if (field.name === fieldName) matches.push(field);
+  }
+  const [field] = matches;
+  assert(
+    field && matches.length === 1,
+    `Dataset field ${fieldName} was not found or is ambiguous.`,
+  );
+  return {
+    datasetId,
+    itemFilter: {
+      fieldValues: [
+        {
+          fieldId: field.id,
+          jsonPath: "$",
+          operator: "text_eq",
+          value: parseFilterValue(rawValue),
+          latestOnly: true,
+          sources: [],
+        },
+      ],
+      fieldSources: [],
+      missingRevisionFields: [],
+    },
+  };
+};
+
+const sourceFromOptions = (
+  options: Pick<TrainingDraftOptions, "baseModelId" | "modelVersionId">,
+  fallback?: TrainingSource,
+): TrainingSource => {
+  assert(
+    !(options.baseModelId && options.modelVersionId),
+    "Choose exactly one of --base-model and --model-version.",
+  );
+  if (options.baseModelId) return { type: "base_model", baseModelId: options.baseModelId };
+  if (options.modelVersionId)
+    return { type: "model_version", modelVersionId: options.modelVersionId };
+  assert(fallback, "Choose exactly one of --base-model and --model-version.");
+  return fallback;
+};
+
+const uploadPromptFile = async (
+  client: Pick<PotatoClient, "files">,
+  projectId: string,
+  filePath: string,
+  fetchImplementation: typeof globalThis.fetch,
+): Promise<string> => {
+  const metadata = await stat(filePath);
+  assert(metadata.isFile(), "--prompt-file must point to a file.");
+  assert(metadata.size > 0, "The gisting prompt file is empty.");
+  assert(
+    metadata.size <= GISTING_PROMPT_MAX_BYTES,
+    "The gisting prompt file exceeds the 256 KiB limit.",
+  );
+  const payload = await readFile(filePath);
+  let prompt: string;
+  try {
+    prompt = new TextDecoder("utf-8", { fatal: true }).decode(payload);
+  } catch {
+    throw new Error("The gisting prompt file must contain valid UTF-8 text.");
+  }
+  assert(prompt.trim().length > 0, "The gisting prompt file is empty.");
+  const started = await client.files.startUpload({
+    projectId,
+    filename: basename(filePath),
+    declaredSizeBytes: payload.byteLength,
+  });
+  try {
+    const response = await fetchImplementation(started.uploadUrl, {
+      method: "PUT",
+      headers: { "x-ms-blob-type": "BlockBlob", "Content-Type": "text/plain; charset=utf-8" },
+      body: payload,
+    });
+    assert(response.ok, `Failed to upload gisting prompt: HTTP ${response.status}.`);
+    return (await client.files.finalizeUpload({ projectId, fileId: started.file.id })).id;
+  } catch (error) {
+    await client.files.delete({ projectId, fileId: started.file.id }).catch(() => undefined);
+    throw error;
+  }
+};
+
+const resolveSelections = async ({
+  client,
+  current,
+  options,
+}: {
+  client: TrainingDraftApi;
+  current: Training["configuration"] | undefined;
+  options: TrainingDraftOptions;
+}): Promise<ResolvedSelections> => {
+  const trainingDatasetId = options.datasetId ?? current?.training.datasetId;
+  assert(trainingDatasetId, "--dataset is required.");
+  const training =
+    options.datasetId === undefined && options.trainingFilter === undefined && current
+      ? current.training
+      : await resolveSelection(client, trainingDatasetId, options.trainingFilter);
+  const validation =
+    options.validationDatasetId === undefined && options.validationFilter === undefined
+      ? current?.validation
+      : await resolveSelection(
+          client,
+          options.validationDatasetId ?? current?.validation?.datasetId ?? "",
+          options.validationFilter,
+        );
+  return { training, validation };
+};
+
+const buildSftConfiguration = ({
+  source,
+  selections,
+  current,
+  options,
+}: {
+  source: TrainingSource;
+  selections: ResolvedSelections;
+  current: SftConfigurationInput | undefined;
+  options: TrainingDraftOptions;
+}): SftConfigurationInput => {
+  assert(options.promptFilePath === undefined, "Gisting options cannot be used with SFT.");
+  return {
+    ...current,
+    method: "sft",
+    source,
+    ...selections,
+    ...(options.regime === undefined ? {} : { regime: options.regime }),
+    ...(options.epochs === undefined ? {} : { epochs: options.epochs }),
+    ...(options.evalSteps === undefined ? {} : { evalSteps: options.evalSteps }),
+  };
+};
+
+const buildGistingConfiguration = async ({
+  client,
+  projectId,
+  source,
+  selections,
+  current,
+  options,
+}: {
+  client: TrainingDraftApi;
+  projectId: string;
+  source: TrainingSource;
+  selections: ResolvedSelections;
+  current: GistingConfigurationInput | undefined;
+  options: TrainingDraftOptions;
+}): Promise<GistingConfigurationInput> => {
+  assert(options.evalSteps === undefined, "SFT options cannot be used with gisting.");
+  const promptFileId = options.promptFilePath
+    ? await uploadPromptFile(
+        client,
+        projectId,
+        options.promptFilePath,
+        options.fetch ?? globalThis.fetch,
+      )
+    : current?.promptFileId;
+  assert(promptFileId, "--prompt-file is required for gisting.");
+  return {
+    ...current,
+    method: "gisting",
+    source,
+    ...selections,
+    promptFileId,
+    ...(options.regime === undefined ? {} : { regime: options.regime }),
+    ...(options.epochs === undefined ? {} : { epochs: options.epochs }),
+  };
+};
+
+const buildTrainingConfiguration = async ({
+  client,
+  projectId,
+  current,
+  options,
+}: {
+  client: TrainingDraftApi;
+  projectId: string;
+  current?: Training["configuration"];
+  options: TrainingDraftOptions;
+}): Promise<TrainingConfigurationInput> => {
+  const source = sourceFromOptions(options, current?.source);
+  const selections = await resolveSelections({ client, current, options });
+  const method = options.method ?? current?.method ?? "sft";
+  return method === "sft"
+    ? buildSftConfiguration({
+        source,
+        selections,
+        current: current?.method === "sft" ? current : undefined,
+        options,
+      })
+    : buildGistingConfiguration({
+        client,
+        projectId,
+        source,
+        selections,
+        current: current?.method === "gisting" ? current : undefined,
+        options,
+      });
+};
+
+export const listTrainings = async (
+  client: TrainingApi,
+  input: { projectId: string; json: boolean },
+  io: TrainingIo,
+): Promise<void> => {
+  const page = await client.trainings.listPage({
+    projectId: input.projectId,
+    limit: 100,
+    offset: 0,
+  });
+  writeOutput(
+    io,
+    input.json,
+    page.trainings,
+    page.trainings.length === 0
+      ? "No Trainings found."
+      : formatTable([
+          ["ID", "NAME", "METHOD", "OUTPUT MODEL"],
+          ...page.trainings.map((training) => [
+            training.id,
+            training.name,
+            training.configuration.method === "sft" ? "SFT" : "GISTING",
+            training.modelId ?? "-",
+          ]),
+        ]),
+  );
+};
+
+export const getTraining = async (
+  client: TrainingApi,
+  input: { projectId: string; nameOrId: string; json: boolean },
+  io: TrainingIo,
+): Promise<void> => {
+  const training = await resolveTraining(client, input.projectId, input.nameOrId);
+  writeOutput(io, input.json, training, formatTraining(training));
+};
+
+export const createTraining = async (
+  client: TrainingDraftApi,
+  input: TrainingDraftOptions & {
+    projectId: string;
+    name: string;
+    json: boolean;
+  },
+  io: TrainingIo,
+): Promise<void> => {
+  assert(input.datasetId, "--dataset is required.");
+  assert(!input.validationFilter || input.validationDatasetId, "--validation-filter requires --validation-dataset.");
+
+  const configuration = await buildTrainingConfiguration({
+    client,
+    projectId: input.projectId,
+    options: input,
+  });
+  const training = await client.trainings.create({
+    projectId: input.projectId,
+    name: input.name,
+    configuration,
+  });
+  writeOutput(
+    io,
+    input.json,
+    training,
+    `Created training ${training.id}`,
+  );
+};
+
+export const updateTraining = async (
+  client: TrainingDraftApi,
+  input: TrainingDraftOptions & {
+    projectId: string;
+    nameOrId: string;
+    json: boolean;
+  },
+  io: TrainingIo,
+): Promise<void> => {
+  const current = await resolveTraining(client, input.projectId, input.nameOrId);
+  const updated = await client.trainings.replaceConfiguration({
+    trainingId: current.id,
+    configuration: await buildTrainingConfiguration({
+      client,
+      projectId: input.projectId,
+      current: current.configuration,
+      options: input,
+    }),
+  });
+  writeOutput(
+    io,
+    input.json,
+    updated,
+    `Updated training ${updated.id}`,
+  );
+};
+
+export const estimateTrainingCost = async (
+  client: TrainingCostApi,
+  input: { projectId: string; nameOrId: string; idempotencyKey?: string; json: boolean },
+  io: TrainingIo,
+): Promise<void> => {
+  const training = await resolveTraining(client, input.projectId, input.nameOrId);
+  const inputSet = await client.trainings.prepareCostEstimateInputs({
+    trainingId: training.id,
+    idempotencyKey: input.idempotencyKey ?? randomUUID(),
+  });
+  const exportIds = [inputSet.exports.training, inputSet.exports.validation].filter(
+    (exportId): exportId is string => exportId !== null,
+  );
+
+  await Promise.all(
+    exportIds.map((exportId) => client.waitFor({ type: "datasetExport", exportId })),
+  );
+
+  const estimate = await client.trainings.estimateCost({
+    trainingId: training.id,
+    inputSetId: inputSet.id,
+  });
+  writeOutput(io, input.json, estimate, formatCostEstimate(estimate));
+};
+
+export const createTrainingRun = async (
+  client: TrainingApi,
+  input: {
+    projectId: string;
+    nameOrId: string;
+    estimateFile: string;
+    versionName: string;
+    idempotencyKey?: string;
+    json: boolean;
+  },
+  io: TrainingIo,
+): Promise<void> => {
+  const training = await resolveTraining(client, input.projectId, input.nameOrId);
+  const estimate = parseReadyEstimateFile(JSON.parse(await readFile(input.estimateFile, "utf8")));
+  assert(estimate.trainingId === training.id, "The estimate file belongs to another Training.");
+  const run = await client.trainings.createRun({
+    trainingId: training.id,
+    inputSetId: estimate.inputSetId,
+    estimateFingerprint: estimate.fingerprint,
+    versionName: input.versionName,
+    idempotencyKey: input.idempotencyKey ?? randomUUID(),
+  });
+  writeOutput(io, input.json, run, formatTrainingRun(run));
+};
+
+export const listTrainingRuns = async (
+  client: TrainingApi,
+  input: { projectId: string; nameOrId: string; json: boolean },
+  io: TrainingIo,
+): Promise<void> => {
+  const training = await resolveTraining(client, input.projectId, input.nameOrId);
+  const page = await client.trainings.listRunsPage({
+    trainingId: training.id,
+    limit: 100,
+    offset: 0,
+  });
+  writeOutput(
+    io,
+    input.json,
+    page.trainingRuns,
+    page.trainingRuns.length === 0
+      ? "No Training Runs found."
+      : formatTable([
+          ["ID", "STATUS", "VERSION", "MAXIMUM CHARGE", "OUTPUT MODEL VERSION"],
+          ...page.trainingRuns.map((run) => [
+            run.id,
+            run.status,
+            run.outputModelVersionName,
+            `$${run.quote.maximumCharge}`,
+            run.outputModelVersionId,
+          ]),
+        ]),
+  );
+};
+
+export const getTrainingRun = async (
+  client: TrainingApi,
+  input: { trainingRunId: string; json: boolean },
+  io: TrainingIo,
+): Promise<void> => {
+  const run = await client.trainings.getRun({ trainingRunId: input.trainingRunId });
+  writeOutput(io, input.json, run, formatTrainingRun(run));
+};
+
+export const resumeTrainingRun = async (
+  client: TrainingApi,
+  input: { trainingRunId: string; json: boolean },
+  io: TrainingIo,
+): Promise<void> => {
+  const run = await client.trainings.resumeRun({ trainingRunId: input.trainingRunId });
+  writeOutput(io, input.json, run, formatTrainingRun(run));
+};
+
+export const cancelTrainingRun = async (
+  client: TrainingApi,
+  input: { trainingRunId: string; json: boolean },
+  io: TrainingIo,
+): Promise<void> => {
+  const run = await client.trainings.cancelRun({ trainingRunId: input.trainingRunId });
+  writeOutput(io, input.json, run, formatTrainingRun(run));
+};
+
+export const readTrainingRunLogs = async (
+  client: TrainingApi,
+  input: { trainingRunId: string; follow: boolean; pollIntervalMs?: number },
+  io: TrainingIo,
+): Promise<void> => {
+  let cursor: string | null = null;
+  for (;;) {
+    const page = await client.trainings.getRunLogs({
+      trainingRunId: input.trainingRunId,
+      cursor: cursor ?? undefined,
+      limit: 100,
+    });
+    for (const entry of page.entries) {
+      io.stdout(`${entry.timestamp} ${entry.level.toUpperCase()} ${entry.message}`);
+    }
+    if (!input.follow || page.complete) {
+      if (page.truncated) io.stdout("Training logs were truncated by the server.");
+      return;
+    }
+    cursor = page.nextCursor;
+    await new Promise<void>((resolve) => {
+      setTimeout(resolve, input.pollIntervalMs ?? 1_000);
+    });
+  }
+};