@toloka-platform/sdk 0.12.2 → 0.13.0

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