@toloka-platform/sdk 0.12.2 → 0.13.1

package/src/cli/commands/batch-inference.ts

@@ -0,0 +1,116 @@
+import type { Command } from "commander";
+
+import {
+  changeBatchInferenceRun,
+  getBatchInferenceRun,
+  listBatchInferenceRuns,
+  startBatchInferenceRun,
+  type StartBatchInferenceRunInput,
+} from "../batch-inference";
+import {
+  collectValues,
+  nonNegativeInteger,
+  partNames,
+  positiveInteger,
+  resolveContext,
+  samplingTemperature,
+  samplingTopP,
+  type CommandHost,
+} from "../context";
+import { parseBaseModelId } from "../models";
+import { assertProjectSelected } from "../project-context";
+
+type RunOptions = Omit<StartBatchInferenceRunInput, "projectId" | "datasetId"> & {
+  dataset: string;
+};
+
+export const registerBatchInferenceCommands = (program: Command, host: CommandHost): void => {
+  const { context, globals } = host;
+  const batchInference = program
+    .command("batch-inference")
+    .description("answer dataset rows with a base model or Model Version");
+  batchInference
+    .command("run")
+    .requiredOption("--dataset <id>", "Dataset whose chat messages are the prompts")
+    .option("--filter <field=value>", "only rows whose field equals the value")
+    .option("--slice <id>", "only rows of this ready Dataset Slice, with the prompts it froze")
+    .option(
+      "--parts <names>",
+      "only rows in these split parts, comma-separated, for example test",
+      partNames,
+    )
+    .option(
+      "--split-column <name-or-id>",
+      "split column that --parts names, when the dataset has more than one",
+    )
+    .option("--base-model <id>", "catalog Base Model that answers", parseBaseModelId)
+    .option("--model-version <id>", "Model Version that answers: mdv.<id> or Model@name")
+    .option("--output-field <name>", "text column on the dataset that receives the answers")
+    .option("--output-dataset <name>", "new dataset that receives the prompts and answers")
+    .requiredOption(
+      "--max-output-tokens <count>",
+      "token limit per answer, at most 8192",
+      positiveInteger,
+    )
+    .option(
+      "--temperature <value>",
+      "sampling temperature from 0 to 2 (default 0)",
+      samplingTemperature,
+    )
+    .option("--top-p <value>", "nucleus sampling threshold above 0 and at most 1", samplingTopP)
+    .option("--stop <text>", "stop sequence; repeat for up to four", collectValues, [])
+    .option("--seed <number>", "sampling seed", nonNegativeInteger)
+    .option(
+      "--samples-per-row <count>",
+      "answers per row, at most 16; above 1 needs --output-dataset and a temperature above 0",
+      positiveInteger,
+    )
+    .option("--idempotency-key <key>", "reuse the run of an earlier identical request")
+    .option("--yes", "confirm the quoted maximum charge and start right away", false)
+    .option("--json", "print JSON", false)
+    .description("quote a Batch Inference Run, then confirm it to start")
+    .action(async (options: RunOptions) => {
+      const { client, projectId } = await resolveContext(globals(), context);
+      assertProjectSelected(projectId);
+      const { dataset, ...rest } = options;
+      await startBatchInferenceRun(client, { ...rest, projectId, datasetId: dataset }, context.io);
+    });
+  batchInference
+    .command("list")
+    .option("--json", "print JSON", false)
+    .description("list the selected project's Batch Inference Runs")
+    .action(async (options: { json: boolean }) => {
+      const { client, projectId } = await resolveContext(globals(), context);
+      assertProjectSelected(projectId);
+      await listBatchInferenceRuns(client, { projectId, json: options.json }, context.io);
+    });
+  batchInference
+    .command("status")
+    .argument("<batch-inference-run-id>", "Batch Inference Run ID")
+    .option("--wait", "wait until the run finishes or needs you", false)
+    .option("--json", "print JSON", false)
+    .description("show one Batch Inference Run")
+    .action(async (batchInferenceRunId: string, options: { wait: boolean; json: boolean }) => {
+      const { client } = await resolveContext(globals(), context);
+      await getBatchInferenceRun(client, { batchInferenceRunId, ...options }, context.io);
+    });
+  for (const [action, description] of [
+    ["confirm", "accept the quoted maximum charge and start the run"],
+    ["resume", "resume a run blocked on funds or payment"],
+    ["cancel", "cancel a run; answers already generated are kept and the GPU time used is charged"],
+  ] as const) {
+    batchInference
+      .command(action)
+      .argument("<batch-inference-run-id>", "Batch Inference Run ID")
+      .option("--json", "print JSON", false)
+      .description(description)
+      .action(async (batchInferenceRunId: string, options: { json: boolean }) => {
+        const { client } = await resolveContext(globals(), context);
+        await changeBatchInferenceRun(
+          client,
+          { batchInferenceRunId, action, json: options.json },
+          context.io,
+        );
+      });
+  }
+};

package/src/cli/context.ts

@@ -97,6 +97,20 @@ export const nonNegativeInteger = (raw: string): number => {
   return value;
 };
 
+export const samplingTemperature = (raw: string): number => {
+  const value = Number(raw);
+  if (raw.trim() === "" || !Number.isFinite(value) || value < 0 || value > 2)
+    throw new InvalidArgumentError("Expected a temperature from 0 to 2.");
+  return value;
+};
+
+export const samplingTopP = (raw: string): number => {
+  const value = Number(raw);
+  if (raw.trim() === "" || !Number.isFinite(value) || value <= 0 || value > 1)
+    throw new InvalidArgumentError("Expected a top-p above 0 and at most 1.");
+  return value;
+};
+
 export const trainingMethod = (raw: string): "sft" | "gisting" => {
   if (raw !== "sft" && raw !== "gisting") {
     throw new InvalidArgumentError('Expected training method "sft" or "gisting".');

package/src/cli/models.ts

@@ -69,7 +69,7 @@ const resolveModel = async (
   return matches[0]!;
 };
 
-const resolveModelVersion = async (
+export const resolveModelVersion = async (
   client: ModelApi,
   projectId: string | undefined,
   nameOrId: string,

package/src/cli/program.ts

@@ -23,6 +23,7 @@ import { registerSourceCommands } from "./commands/source";
 import { registerStatusCommand } from "./commands/status";
 import { registerTrainingCommands } from "./commands/training";
 import { registerUiCommands } from "./commands/ui";
+import { registerBatchInferenceCommands } from "./commands/batch-inference";
 import { registerVersionCommand } from "./commands/version";
 import type { TolokaEnvironment } from "./config";
 import {
@@ -102,6 +103,7 @@ export function createProgram(context: ProgramContext): CreatedProgram {
   registerSecretCommands(program, host);
   registerModelCommands(program, host);
   registerTrainingCommands(program, host);
+  registerBatchInferenceCommands(program, host);
   registerDatasetCommands(program, host);
   registerPipelineCommands(program, host);
   registerStatusCommand(program, host);

package/src/cli/splits.ts

@@ -111,6 +111,46 @@ export const formatSplitColumn = (field: DatasetField, fieldNames: ReadonlyMap<s
   return formatTable(rows);
 };
 
+export const splitColumnId = async (
+  client: Pick<PotatoClient, "datasets">,
+  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}.` },
+  });
+};
+
+export const resolveSplitParts = async (
+  client: Pick<PotatoClient, "datasets">,
+  datasetId: string,
+  { splitColumn, parts }: { splitColumn: string | undefined; parts: string[] },
+): Promise<{ fieldId: string; names: string[] }> => {
+  if (splitColumn !== undefined) return { fieldId: await splitColumnId(client, datasetId, splitColumn), names: parts };
+  const splitColumns: DatasetField[] = [];
+  for await (const field of client.datasets.listAllFields({ datasetId })) {
+    if (field.split !== null) splitColumns.push(field);
+  }
+  const [only, ...others] = splitColumns;
+  assertCli(only !== undefined, `Dataset ${datasetId} has no split column.`, {
+    hint: { description: `Make one with toloka dataset split ${datasetId}.` },
+  });
+  assertCli(
+    others.length === 0,
+    `Dataset ${datasetId} has the split columns ${splitColumns.map((field) => field.name).join(", ")}.`,
+    { hint: { description: "Pass --split-column <name> to choose one." } },
+  );
+  return { fieldId: only.id, names: parts };
+};
+
 export const waitForSplitColumn = async (
   client: SplitApi,
   { datasetId, field, waitFor }: { datasetId: string; field: DatasetField; waitFor?: WaitForOptions },

package/src/cli/trainings.ts

@@ -19,7 +19,7 @@ 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 { DEFAULT_SPLIT_COLUMN, splitColumnId, waitForSplitColumn } from "./splits";
 import { uploadedFieldHint } from "./uploaded-fields";
 import { formatTable, type TableIo as TrainingIo, writeOutput } from "./table";
 
@@ -280,24 +280,6 @@ 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,
   {

package/src/client.ts

@@ -4,8 +4,10 @@ import * as wire from "../openapi/gen/api";
 import type {
   AddDatasetItemsBody,
   AttachPipelineDatasetBody,
+  BatchInferenceRun as GeneratedBatchInferenceRun,
   CheckJob as GeneratedCheckJob,
   CopyPipelineBody,
+  CreateBatchInferenceRunBody,
   CreateDatasetBody,
   CreateDatasetExportBody,
   CreateDatasetFieldsFromUploadedRowsBody,
@@ -39,6 +41,7 @@ import type {
   ListProjectDatasetsParams,
   ListProjectPipelinesParams,
   ListAudienceSpecializationsParams,
+  ListBatchInferenceRunsParams,
   ListProjectModelsParams,
   ListProjectsParams,
   ListProjectTokensParams,
@@ -85,6 +88,7 @@ declare const __brand: unique symbol;
 type Brand<B> = { [__brand]: B };
 export type Branded<T extends string, B> = T & Brand<B>;
 
+export type BatchInferenceRunId = Branded<string, "BatchInferenceRunId">;
 export type CheckRunId = Branded<string, "CheckRunId">;
 export type DatasetFieldId = Branded<string, "DatasetFieldId">;
 export type DatasetId = Branded<string, "DatasetId">;
@@ -108,6 +112,7 @@ type ResourceIdInput<T extends Branded<string, unknown>> = string & {
 };
 
 type ResourceIdsByProperty = {
+  batchInferenceRunId: BatchInferenceRunId;
   checkRunId: CheckRunId;
   datasetFieldId: DatasetFieldId;
   datasetId: DatasetId;
@@ -164,6 +169,7 @@ type PublicInput<T> = T extends readonly unknown[]
 
 type BrandResourceUnion<T> = T extends object ? BrandResourceReferences<T> : T;
 
+export type BatchInferenceRun = BrandedResource<GeneratedBatchInferenceRun, BatchInferenceRunId>;
 export type CheckJob = BrandedResource<GeneratedCheckJob, CheckRunId>;
 export type Dataset = BrandedResource<GeneratedDataset, DatasetId>;
 export type DatasetField = BrandedResource<GeneratedDatasetField, DatasetFieldId>;
@@ -592,6 +598,46 @@ const createClientImplementation = (config: ClientConfig) => {
       wire.getTrainingRunMetricHistory(trainingRunId, metricKey, params, requestOptions),
   };
 
+  const batchInference = {
+    createRun: ({
+      projectId,
+      idempotencyKey,
+      ...body
+    }: { projectId: ProjectId; idempotencyKey: string } & CreateBatchInferenceRunBody) => {
+      const headers = new Headers(requestOptions.headers);
+      headers.set("Idempotency-Key", idempotencyKey);
+      return wire.createBatchInferenceRun(projectId, body, {
+        ...requestOptions,
+        headers: Object.fromEntries(headers.entries()),
+      });
+    },
+    listRunsPage: ({
+      projectId,
+      ...params
+    }: { projectId: ProjectId } & ListBatchInferenceRunsParams) =>
+      wire.listBatchInferenceRuns(projectId, params, requestOptions),
+    listAllRuns: ({
+      projectId,
+      ...params
+    }: { projectId: ProjectId } & ListBatchInferenceRunsParams) =>
+      listAllPages(async (pageOptions) => {
+        const page = await wire.listBatchInferenceRuns(
+          projectId,
+          { ...params, ...pageOptions },
+          requestOptions,
+        );
+        return { items: page.batchInferenceRuns, pagination: page.pagination };
+      }, params),
+    getRun: ({ batchInferenceRunId }: { batchInferenceRunId: BatchInferenceRunId }) =>
+      wire.getBatchInferenceRun(batchInferenceRunId, requestOptions),
+    confirmRun: ({ batchInferenceRunId }: { batchInferenceRunId: BatchInferenceRunId }) =>
+      wire.confirmBatchInferenceRun(batchInferenceRunId, requestOptions),
+    resumeRun: ({ batchInferenceRunId }: { batchInferenceRunId: BatchInferenceRunId }) =>
+      wire.resumeBatchInferenceRun(batchInferenceRunId, requestOptions),
+    cancelRun: ({ batchInferenceRunId }: { batchInferenceRunId: BatchInferenceRunId }) =>
+      wire.cancelBatchInferenceRun(batchInferenceRunId, requestOptions),
+  };
+
   const datasets = {
     create: ({ projectId, ...body }: { projectId: ProjectId } & CreateDatasetBody) =>
       wire.createDataset(projectId, body, requestOptions),
@@ -778,6 +824,7 @@ const createClientImplementation = (config: ClientConfig) => {
     generationModels,
     models,
     trainings,
+    batchInference,
     runs,
     secrets,
     files,
@@ -807,6 +854,8 @@ type PublicOutput<T> = T extends GeneratedCheckJob
                   ? TrainingCostEstimate
                   : T extends GeneratedTrainingRun
                     ? TrainingRun
+                    : T extends GeneratedBatchInferenceRun
+                      ? BatchInferenceRun
                     : T extends GeneratedProject
                       ? Project
                       : T extends GeneratedPipeline

package/src/pipeline-graph-v1.contract.json

@@ -1535,7 +1535,7 @@
                           "values": []
                         }
                       },
-                      "pricePerTask": 1,
+                      "pricePerTask": 6,
                       "timeForTask": 1800,
                       "perExpertCapEnabled": false,
                       "perExpertCapValue": 1
@@ -2663,7 +2663,7 @@
                           "values": []
                         }
                       },
-                      "pricePerTask": 1,
+                      "pricePerTask": 6,
                       "timeForTask": 1800,
                       "perExpertCapEnabled": false,
                       "perExpertCapValue": 1