@toloka-platform/sdk 0.11.4 → 0.12.1

package/src/cli/trainings.ts

@@ -2,8 +2,11 @@ import { basename } from "node:path";
 import { readFile, stat } from "node:fs/promises";
 
 import { randomUUID } from "node:crypto";
+import { omit } from "es-toolkit";
 
 import type {
+  DatasetField,
+  DatasetSlice,
   PotatoClient,
   Training,
   TrainingCostEstimate,
@@ -14,6 +17,10 @@ import type {
 import type { BaseModelId, TrainingConfigurationInput } from "../../openapi/gen/models";
 import { assertCli, RETRY_OR_REPORT_HINT, TolokaCliError } from "./cli-error";
 import type { CliErrorHint } from "./cli-error";
+import { resolveLiveDataset } from "./data";
+import { formatCount, waitForFilledFields, waitUntilDone } from "./slices";
+import { DEFAULT_SPLIT_COLUMN, waitForSplitColumn } from "./splits";
+import { uploadedFieldHint } from "./uploaded-fields";
 import { formatTable, type TableIo as TrainingIo, writeOutput } from "./table";
 
 const ESTIMATE_FILE_HINT: CliErrorHint = {
@@ -22,19 +29,19 @@ const ESTIMATE_FILE_HINT: CliErrorHint = {
 
 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">;
+type ResolvedSelections = Pick<SftConfigurationInput, "sliceId" | "splitFieldId" | "trainParts">;
 
 export type TrainingDraftOptions = {
   method?: TrainingMethod;
   baseModelId?: BaseModelId;
   modelVersionId?: string;
-  datasetId?: string;
-  validationDatasetId?: string;
-  trainingFilter?: string;
-  validationFilter?: string;
+  sourceRoot?: string;
+  datasetName?: string;
+  sliceId?: string;
+  splitColumn?: string;
+  trainParts?: string[];
   loraRank?: SftConfigurationInput["loraRank"];
   gistTokens?: GistingConfigurationInput["gistTokens"];
   epochs?: number;
@@ -46,7 +53,7 @@ export type TrainingDraftOptions = {
 };
 
 type TrainingApi = Pick<PotatoClient, "trainings">;
-type TrainingDraftApi = Pick<PotatoClient, "trainings" | "files" | "datasets">;
+type TrainingDraftApi = PotatoClient;
 type TrainingCostApi = Pick<PotatoClient, "trainings" | "waitFor">;
 type TrainingRunApi = Pick<PotatoClient, "trainings" | "waitFor">;
 const GISTING_PROMPT_MAX_BYTES = 256 * 1024;
@@ -58,6 +65,9 @@ const formatTraining = (training: Training): string =>
       "Method",
       training.configuration.method === "sft" ? "Supervised fine-tuning (SFT)" : "Prompt gisting",
     ],
+    ["Slice", training.configuration.sliceId ?? "-"],
+    ["Split column", training.configuration.splitFieldId ?? "-"],
+    ["Train parts", training.configuration.trainParts?.join(", ") ?? "every train part"],
     training.configuration.method === "sft"
       ? ["LoRA rank", String(training.configuration.loraRank)]
       : ["Gist tokens", String(training.configuration.gistTokens)],
@@ -152,9 +162,14 @@ const formatCostEstimate = (estimate: TrainingCostEstimate): string => {
     ].join("\n");
   }
 
+  const leftOutAsTest = estimate.inputs.training.heldOutRows ?? 0;
+
   return formatTable([
     ["Training", estimate.trainingId],
     ["Training rows", String(estimate.inputs.training.report?.validCount ?? 0)],
+    ...(leftOutAsTest === 0
+      ? []
+      : [["Left out as test", `${formatCount(leftOutAsTest)} train rows that a split column now marks as test`]]),
     ["Training tokens", String(estimate.inputs.training.report?.tokens.total ?? 0)],
     ["Maximum GPU seconds", String(estimate.budget.maximumGpuSeconds)],
     ["Maximum charge", `$${estimate.maximumCharge} ${estimate.currency}`],
@@ -187,63 +202,6 @@ const resolveTraining = async (
   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> => {
-  assertCli(datasetId, "A Dataset is required when specifying a filter.", {
-    hint: {
-      description:
-        "Pass --dataset together with --train-filter, and --validation-dataset together with --validation-filter.",
-    },
-  });
-  if (filter === undefined) return { datasetId };
-  const separator = filter.indexOf("=");
-  assertCli(
-    separator > 0 && separator < filter.length - 1,
-    "Dataset filters must use FIELD=VALUE syntax.",
-    { hint: { description: "Write the filter as field=value, for example split=train." } },
-  );
-  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;
-  assertCli(
-    field && matches.length === 1,
-    `Dataset field ${fieldName} was not found or is ambiguous.`,
-    { hint: { description: "Run toloka dataset fields <dataset> to see the field names." } },
-  );
-  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,
@@ -313,44 +271,177 @@ const uploadPromptFile = async (
   }
 };
 
-const assertDatasetIdOption = (
-  flag: "--dataset" | "--validation-dataset",
-  value: string | undefined,
-): void =>
+const ROWS_HINT: CliErrorHint = {
+  description:
+    "Pass --dataset <dataset> to split it and freeze a slice for this Training, or --slice <slice-id> for a slice that freezes a split column.",
+};
+
+const AUTO_WAIT = { timeoutMs: 10 * 60_000 };
+
+type ResolvedRows = ResolvedSelections & { readonly notes: readonly string[] };
+
+const splitColumnId = async (
+  client: TrainingDraftApi,
+  datasetId: string,
+  column: string,
+): Promise<string> => {
+  for await (const field of client.datasets.listAllFields({ datasetId })) {
+    if (field.id === column || field.name === column) {
+      assertCli(field.split !== null, `Column ${field.name} is not a split column.`, {
+        hint: { description: `Make it one with toloka dataset split ${datasetId} --column ${field.name}.` },
+      });
+      return field.id;
+    }
+  }
+  assertCli(false, `Dataset ${datasetId} has no column ${column}.`, {
+    hint: { description: `List its columns with toloka dataset columns ${datasetId}.` },
+  });
+};
+
+const fillTrainingSplitColumn = async (
+  client: TrainingDraftApi,
+  {
+    datasetId,
+    datasetName,
+    datasetRef,
+    splitColumn,
+  }: { datasetId: string; datasetName: string; datasetRef: string; splitColumn: string | undefined },
+): Promise<DatasetField> => {
+  const fields = await waitForFilledFields(client, { datasetId, waitFor: AUTO_WAIT });
+  const splitColumns = fields.filter((field) => field.split !== null);
   assertCli(
-    value === undefined || /^ds\./i.test(value),
-    `${flag} takes a Dataset id (ds.…), got ${JSON.stringify(value)}.`,
-    { hint: { description: "Run toloka dataset list and pass the Dataset id (ds.…)." } },
+    splitColumn !== undefined || splitColumns.length <= 1,
+    `Dataset ${datasetName} has the split columns ${splitColumns.map((field) => field.name).join(", ")}.`,
+    { hint: { description: "Pass --split-column <name> to choose one." } },
   );
+  const existing =
+    splitColumn === undefined
+      ? splitColumns[0]
+      : fields.find((field) => field.id === splitColumn || field.name === splitColumn);
+
+  if (existing?.split?.recipe === null) return existing;
+
+  const column = existing?.name ?? splitColumn ?? DEFAULT_SPLIT_COLUMN;
+  if (existing === undefined) {
+    const uploaded = await uploadedFieldHint(client, { datasetId, datasetRef, name: column });
+    assertCli(
+      uploaded === undefined,
+      `Dataset ${datasetName} has no field ${column}, but its uploaded rows have one.`,
+      {
+        hint: {
+          description: `Add it as a dataset field with toloka dataset fields add ${datasetRef} ${column} to train on its parts, or pass --split-column <name> to sample a new split column.`,
+        },
+      },
+    );
+  }
+
+  return waitForSplitColumn(client, {
+    datasetId,
+    field: await client.datasets.split({ datasetId, column }),
+    waitFor: AUTO_WAIT,
+  });
+};
+
+const freezeTrainingSlice = async (
+  client: TrainingDraftApi,
+  { datasetId, name }: { datasetId: string; name: string },
+): Promise<DatasetSlice> => {
+  const created = await client.datasets.createSlice({ datasetId, name });
+  const slice = await waitUntilDone(() => client.waitFor({ type: "datasetSlice", sliceId: created.id }, AUTO_WAIT), {
+    what: `slice ${created.id}`,
+    checkCommand: `dataset slice show ${created.id}`,
+  });
+  assertCli(slice.status === "ready", `Slice ${slice.id} failed: ${slice.error?.message ?? "no reason given"}`, {
+    hint: { description: `Check it with toloka dataset slice show ${slice.id}.` },
+  });
+  return slice;
+};
+
+const splitAndSlice = async (
+  client: TrainingDraftApi,
+  {
+    sourceRoot,
+    projectId,
+    datasetName,
+    splitColumn,
+    trainingName,
+  }: { sourceRoot: string; projectId: string; datasetName: string; splitColumn: string | undefined; trainingName: string },
+): Promise<ResolvedRows> => {
+  const dataset = await resolveLiveDataset(client, { sourceRoot, projectId, datasetName }, false);
+  const column = await fillTrainingSplitColumn(client, {
+    datasetId: dataset.datasetId,
+    datasetName: dataset.name,
+    datasetRef: dataset.nameIsUnique ? dataset.name : dataset.datasetId,
+    splitColumn,
+  });
+  const slice = await freezeTrainingSlice(client, { datasetId: dataset.datasetId, name: trainingName });
+  const lastFill = column.split?.lastFill;
+
+  return {
+    sliceId: slice.id,
+    splitFieldId: column.id,
+    notes: [
+      ...(lastFill == null
+        ? []
+        : [
+            `Split ${dataset.name} by column ${column.name}: placed ${formatCount(lastFill.placed)} rows, ${formatCount(lastFill.kept)} kept their part.`,
+          ]),
+      `Froze slice ${slice.name} (${slice.id}) with ${formatCount(slice.rows)} rows.`,
+    ],
+  };
+};
 
 const resolveSelections = async ({
   client,
+  projectId,
+  trainingName,
   current,
   options,
 }: {
   client: TrainingDraftApi;
+  projectId: string;
+  trainingName: string;
   current: Training["configuration"] | undefined;
   options: TrainingDraftOptions;
-}): Promise<ResolvedSelections> => {
-  assertDatasetIdOption("--dataset", options.datasetId);
-  assertDatasetIdOption("--validation-dataset", options.validationDatasetId);
-  const trainingDatasetId = options.datasetId ?? current?.training.datasetId;
-  assertCli(trainingDatasetId, "--dataset is required.", {
-    hint: { description: "Pass --dataset <dataset-id>; toloka dataset list shows the ids." },
+}): Promise<ResolvedRows> => {
+  assertCli(options.datasetName === undefined || options.sliceId === undefined, "Pass --dataset or --slice, not both.", {
+    hint: ROWS_HINT,
   });
-  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 };
+  if (options.datasetName !== undefined) {
+    assertCli(options.sourceRoot !== undefined, "The CLI needs the project source root to resolve --dataset.", {
+      hint: RETRY_OR_REPORT_HINT,
+    });
+    return {
+      ...(await splitAndSlice(client, {
+        sourceRoot: options.sourceRoot,
+        projectId,
+        datasetName: options.datasetName,
+        splitColumn: options.splitColumn,
+        trainingName,
+      })),
+      ...(options.trainParts === undefined ? {} : { trainParts: options.trainParts }),
+    };
+  }
+
+  const inherited = options.sliceId === undefined ? current : undefined;
+  const sliceId = options.sliceId ?? inherited?.sliceId;
+  assertCli(sliceId !== undefined, "--dataset or --slice is required.", { hint: ROWS_HINT });
+  const slice = await client.datasets.getSlice({ sliceId });
+  assertCli(slice.status === "ready", `Slice ${slice.id} is ${slice.status}, not ready.`, {
+    hint: { description: `Check it with toloka dataset slice show ${slice.id}.` },
+  });
+  const splitFieldId =
+    options.splitColumn === undefined
+      ? inherited?.splitFieldId
+      : await splitColumnId(client, slice.datasetId, options.splitColumn);
+  const trainParts = options.trainParts ?? inherited?.trainParts;
+
+  return {
+    sliceId: slice.id,
+    ...(splitFieldId === undefined ? {} : { splitFieldId }),
+    ...(trainParts === undefined ? {} : { trainParts }),
+    notes: [],
+  };
 };
 
 const sharedHyperparameters = (options: TrainingDraftOptions) => ({
@@ -361,17 +452,19 @@ const sharedHyperparameters = (options: TrainingDraftOptions) => ({
     : { maxSequenceLength: options.maxSequenceLength }),
 });
 
-const buildSftConfiguration = ({
-  source,
-  selections,
-  current,
-  options,
-}: {
-  source: TrainingSource;
-  selections: ResolvedSelections;
-  current: SftConfigurationInput | undefined;
-  options: TrainingDraftOptions;
-}): SftConfigurationInput => {
+const assertMethodOptions = (method: TrainingMethod, options: TrainingDraftOptions): void => {
+  if (method === "gisting") {
+    assertCli(
+      options.loraRank === undefined,
+      "--lora-rank applies to SFT; gisting takes --gist-tokens.",
+      {
+        hint: {
+          description: "Use --gist-tokens for gisting, or omit --method gisting to create SFT.",
+        },
+      },
+    );
+    return;
+  }
   assertCli(options.promptFilePath === undefined, "--prompt-file applies to gisting only.", {
     hint: {
       description: "Pass --method gisting to use --prompt-file, or omit --prompt-file for SFT.",
@@ -386,8 +479,23 @@ const buildSftConfiguration = ({
       },
     },
   );
+};
+
+const SELECTION_KEYS = ["training", "validation", "sliceId", "splitFieldId", "trainParts"] as const;
+
+const buildSftConfiguration = ({
+  source,
+  selections,
+  current,
+  options,
+}: {
+  source: TrainingSource;
+  selections: ResolvedSelections;
+  current: Extract<Training["configuration"], { method: "sft" }> | undefined;
+  options: TrainingDraftOptions;
+}): SftConfigurationInput => {
   return {
-    ...current,
+    ...(current === undefined ? {} : omit(current, SELECTION_KEYS)),
     method: "sft",
     source,
     ...selections,
@@ -409,18 +517,9 @@ const buildGistingConfiguration = async ({
   projectId: string;
   source: TrainingSource;
   selections: ResolvedSelections;
-  current: GistingConfigurationInput | undefined;
+  current: Extract<Training["configuration"], { method: "gisting" }> | undefined;
   options: TrainingDraftOptions;
 }): Promise<GistingConfigurationInput> => {
-  assertCli(
-    options.loraRank === undefined,
-    "--lora-rank applies to SFT; gisting takes --gist-tokens.",
-    {
-      hint: {
-        description: "Use --gist-tokens for gisting, or omit --method gisting to create SFT.",
-      },
-    },
-  );
   const promptFileId = options.promptFilePath
     ? await uploadPromptFile(
         client,
@@ -433,7 +532,7 @@ const buildGistingConfiguration = async ({
     hint: { description: "Pass --prompt-file <absolute-path> with the prompt text." },
   });
   return {
-    ...current,
+    ...(current === undefined ? {} : omit(current, SELECTION_KEYS)),
     method: "gisting",
     source,
     ...selections,
@@ -447,18 +546,21 @@ const buildGistingConfiguration = async ({
 const buildTrainingConfiguration = async ({
   client,
   projectId,
+  trainingName,
   current,
   options,
 }: {
   client: TrainingDraftApi;
   projectId: string;
+  trainingName: string;
   current?: Training["configuration"];
   options: TrainingDraftOptions;
-}): Promise<TrainingConfigurationInput> => {
+}): Promise<{ configuration: TrainingConfigurationInput; notes: readonly string[] }> => {
   const source = sourceFromOptions(options, current?.source);
-  const selections = await resolveSelections({ client, current, options });
   const method = options.method ?? current?.method ?? "sft";
-  return method === "sft"
+  assertMethodOptions(method, options);
+  const { notes, ...selections } = await resolveSelections({ client, projectId, trainingName, current, options });
+  const configuration = await (method === "sft"
     ? buildSftConfiguration({
         source,
         selections,
@@ -472,7 +574,8 @@ const buildTrainingConfiguration = async ({
         selections,
         current: current?.method === "gisting" ? current : undefined,
         options,
-      });
+      }));
+  return { configuration, notes };
 };
 
 export const listTrainings = async (
@@ -521,18 +624,10 @@ export const createTraining = async (
   },
   io: TrainingIo,
 ): Promise<void> => {
-  assertCli(input.datasetId, "--dataset is required.", {
-    hint: { description: "Pass --dataset <dataset-id>; toloka dataset list shows the ids." },
-  });
-  assertCli(
-    !input.validationFilter || input.validationDatasetId,
-    "--validation-filter requires --validation-dataset.",
-    { hint: { disabled: true } },
-  );
-
-  const configuration = await buildTrainingConfiguration({
+  const { configuration, notes } = await buildTrainingConfiguration({
     client,
     projectId: input.projectId,
+    trainingName: input.name,
     options: input,
   });
   const training = await client.trainings.create({
@@ -540,7 +635,7 @@ export const createTraining = async (
     name: input.name,
     configuration,
   });
-  writeOutput(io, input.json, training, `Created training ${training.id}`);
+  writeOutput(io, input.json, training, [...notes, `Created training ${training.id}`].join("\n"));
 };
 
 export const updateTraining = async (
@@ -553,16 +648,15 @@ export const updateTraining = async (
   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,
-    }),
+  const { configuration, notes } = await buildTrainingConfiguration({
+    client,
+    projectId: input.projectId,
+    trainingName: current.name,
+    current: current.configuration,
+    options: input,
   });
-  writeOutput(io, input.json, updated, `Updated training ${updated.id}`);
+  const updated = await client.trainings.replaceConfiguration({ trainingId: current.id, configuration });
+  writeOutput(io, input.json, updated, [...notes, `Updated training ${updated.id}`].join("\n"));
 };
 
 const prepareTrainingCostEstimate = async (

package/src/cli/uploaded-fields.ts

@@ -0,0 +1,20 @@
+import type { PotatoClient } from "../client";
+import type { CliErrorHint } from "./cli-error";
+
+export const uploadedRowFieldNames = async (
+  client: Pick<PotatoClient, "datasets">,
+  datasetId: string,
+): Promise<string[]> =>
+  (await client.datasets.getInferredFields({ datasetId })).fields
+    .filter((field) => field.depth === 1)
+    .map((field) => field.name);
+
+export const uploadedFieldHint = async (
+  client: Pick<PotatoClient, "datasets">,
+  { datasetId, datasetRef, name }: { readonly datasetId: string; readonly datasetRef: string; readonly name: string },
+): Promise<CliErrorHint | undefined> =>
+  (await uploadedRowFieldNames(client, datasetId)).includes(name)
+    ? {
+        description: `${name} is a field of the uploaded rows but not a dataset field yet. Add it with toloka dataset fields add ${datasetRef} ${name}, then run this again.`,
+      }
+    : undefined;