@toloka-platform/sdk 0.11.3 → 0.12.0

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,9 @@ 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, waitUntilDone } from "./slices";
+import { waitForSplitColumn } from "./splits";
 import { formatTable, type TableIo as TrainingIo, writeOutput } from "./table";
 
 const ESTIMATE_FILE_HINT: CliErrorHint = {
@@ -22,19 +28,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 +52,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 +64,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)],
@@ -155,6 +164,7 @@ const formatCostEstimate = (estimate: TrainingCostEstimate): string => {
   return formatTable([
     ["Training", estimate.trainingId],
     ["Training rows", String(estimate.inputs.training.report?.validCount ?? 0)],
+    ["Held-out test rows", String(estimate.inputs.training.heldOutRows ?? 0)],
     ["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 +197,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 +266,158 @@ 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, splitColumn }: { datasetId: string; datasetName: string; splitColumn: string | undefined },
+): Promise<DatasetField> => {
+  const fields: DatasetField[] = [];
+  for await (const field of client.datasets.listAllFields({ datasetId })) fields.push(field);
+  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;
+
+  return waitForSplitColumn(client, {
+    datasetId,
+    field: await client.datasets.split({ datasetId, column: existing?.name ?? splitColumn ?? "split" }),
+    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,
+    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,
+  });
+  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 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 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 +428,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 +455,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 +493,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 +508,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 +522,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 +550,8 @@ const buildTrainingConfiguration = async ({
         selections,
         current: current?.method === "gisting" ? current : undefined,
         options,
-      });
+      }));
+  return { configuration, notes };
 };
 
 export const listTrainings = async (
@@ -521,18 +600,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 +611,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 +624,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/client.ts

@@ -8,6 +8,7 @@ import type {
   CopyPipelineBody,
   CreateDatasetBody,
   CreateDatasetExportBody,
+  CreateDatasetSliceBody,
   CreateTrainingRunBody,
   CreateProjectModelBody,
   CreateDatasetFieldBodyAnyOf,
@@ -20,6 +21,7 @@ import type {
   Dataset as GeneratedDataset,
   DatasetField as GeneratedDatasetField,
   DatasetItem as GeneratedDatasetItem,
+  DatasetSlice as GeneratedDatasetSlice,
   ExportJob as GeneratedExportJob,
   EstimateTrainingCostBody,
   EstimatePipelineCostBody,
@@ -28,6 +30,7 @@ import type {
   GetTrainingRunLogsParams,
   GetTrainingRunMetricHistoryParams,
   ListDatasetFieldsParams,
+  ListDatasetSlicesParams,
   ListModelVersionsParams,
   ListDatasetItemsParams,
   ListPipelineRunItemsParams,
@@ -45,6 +48,7 @@ import type {
   ModelVersion as GeneratedModelVersion,
   PrepareTrainingCostEstimateInputsBody,
   PatchPipelineNodeBody,
+  SplitDatasetBody,
   Training as GeneratedTraining,
   TrainingCostEstimate as GeneratedTrainingCostEstimate,
   TrainingCostEstimateInputSet as GeneratedTrainingCostEstimateInputSet,
@@ -84,6 +88,7 @@ export type CheckRunId = Branded<string, "CheckRunId">;
 export type DatasetFieldId = Branded<string, "DatasetFieldId">;
 export type DatasetId = Branded<string, "DatasetId">;
 export type DatasetItemId = Branded<string, "DatasetItemId">;
+export type DatasetSliceId = Branded<string, "DatasetSliceId">;
 export type ExportId = Branded<string, "ExportId">;
 export type ModelId = Branded<string, "ModelId">;
 export type ModelVersionId = Branded<string, "ModelVersionId">;
@@ -116,11 +121,14 @@ type ResourceIdsByProperty = {
   trainingRunId: TrainingRunId;
   goldenDatasetId: DatasetId;
   itemId: RunItemId;
+  fromSliceId: DatasetSliceId;
   pipelineId: PipelineId;
   projectFileId: ProjectFileId;
   projectId: ProjectId;
   runId: RunId;
   runItemId: RunItemId;
+  sliceId: DatasetSliceId;
+  splitFieldId: DatasetFieldId;
   uploadId: UploadId;
 };
 
@@ -159,6 +167,7 @@ export type CheckJob = BrandedResource<GeneratedCheckJob, CheckRunId>;
 export type Dataset = BrandedResource<GeneratedDataset, DatasetId>;
 export type DatasetField = BrandedResource<GeneratedDatasetField, DatasetFieldId>;
 export type DatasetItem = BrandedResource<GeneratedDatasetItem, DatasetItemId>;
+export type DatasetSlice = BrandedResource<GeneratedDatasetSlice, DatasetSliceId>;
 export type ExportJob = BrandedResource<GeneratedExportJob, ExportId>;
 export type InferredField = GeneratedInferredField;
 export type Model = BrandedResource<GeneratedModel, ModelId>;
@@ -222,11 +231,24 @@ export type PipelineRunWaitTarget = {
   readonly runId: ResourceIdInput<RunId>;
 };
 
+export type DatasetSliceWaitTarget = {
+  readonly type: "datasetSlice";
+  readonly sliceId: ResourceIdInput<DatasetSliceId>;
+};
+
+export type DatasetSplitWaitTarget = {
+  readonly type: "datasetSplit";
+  readonly datasetId: ResourceIdInput<DatasetId>;
+  readonly column: string;
+};
+
 export type WaitForTarget =
   | DatasetUploadWaitTarget
   | DatasetExportWaitTarget
   | PipelineCheckWaitTarget
-  | PipelineRunWaitTarget;
+  | PipelineRunWaitTarget
+  | DatasetSliceWaitTarget
+  | DatasetSplitWaitTarget;
 
 const createClientImplementation = (config: ClientConfig) => {
   const requestOptions = createRequestOptions(config.baseUrl, config, config.fetch);
@@ -247,10 +269,19 @@ const createClientImplementation = (config: ClientConfig) => {
     target: PipelineRunWaitTarget,
     options?: WaitForOptions,
   ): Promise<TerminalJob<RunJob>>;
+  function waitFor(target: DatasetSliceWaitTarget, options?: WaitForOptions): Promise<DatasetSlice>;
+  function waitFor(target: DatasetSplitWaitTarget, options?: WaitForOptions): Promise<DatasetField>;
   async function waitFor(
     target: WaitForTarget,
     options?: WaitForOptions,
-  ): Promise<GeneratedUploadJob | GeneratedExportJob | GeneratedCheckJob | GeneratedRunJob> {
+  ): Promise<
+    | GeneratedUploadJob
+    | GeneratedExportJob
+    | GeneratedCheckJob
+    | GeneratedRunJob
+    | GeneratedDatasetSlice
+    | GeneratedDatasetField
+  > {
     switch (target.type) {
       case "datasetUpload":
         return pollUntilTerminal(
@@ -276,6 +307,27 @@ const createClientImplementation = (config: ClientConfig) => {
           (job) => job.status === "succeeded" || job.status === "failed",
           options,
         );
+      case "datasetSlice":
+        return pollUntilTerminal(
+          (signal) => wire.getDatasetSlice(target.sliceId, { ...requestOptions, signal }),
+          (slice) => slice.status !== "creating",
+          options,
+        );
+      case "datasetSplit":
+        return pollUntilTerminal(
+          async (signal) => {
+            const page = await wire.listDatasetFields(
+              target.datasetId,
+              { names: target.column },
+              { ...requestOptions, signal },
+            );
+            const field = page.fields.find((candidate) => candidate.name === target.column);
+            assert(field?.split, `Column ${target.column} is not a split column`);
+            return field;
+          },
+          (field) => field.split?.status !== "filling",
+          options,
+        );
     }
   }
 
@@ -584,6 +636,14 @@ const createClientImplementation = (config: ClientConfig) => {
       wire.createDatasetExport(datasetId, body, requestOptions),
     getExport: ({ exportId }: { exportId: ExportId }) =>
       wire.getDatasetExport(exportId, requestOptions),
+    createSlice: ({ datasetId, ...body }: { datasetId: DatasetId } & CreateDatasetSliceBody) =>
+      wire.createDatasetSlice(datasetId, body, requestOptions),
+    getSlice: ({ sliceId }: { sliceId: DatasetSliceId }) =>
+      wire.getDatasetSlice(sliceId, requestOptions),
+    listSlicesPage: ({ datasetId, ...params }: { datasetId: DatasetId } & ListDatasetSlicesParams) =>
+      wire.listDatasetSlices(datasetId, params, requestOptions),
+    split: ({ datasetId, ...body }: { datasetId: DatasetId } & SplitDatasetBody) =>
+      wire.splitDataset(datasetId, body, requestOptions),
   };
 
   const runs = {
@@ -714,7 +774,9 @@ type PublicOutput<T> = T extends GeneratedCheckJob
                       ? Project
                       : T extends GeneratedPipeline
                         ? Pipeline
-                        : T extends GeneratedDataset
+                        : T extends GeneratedDatasetSlice
+                          ? DatasetSlice
+                          : T extends GeneratedDataset
                           ? Dataset
                           : T extends GeneratedDatasetField
                             ? DatasetField

package/src/index.ts

@@ -11,6 +11,10 @@ export type {
   DatasetId,
   DatasetItem,
   DatasetItemId,
+  DatasetSlice,
+  DatasetSliceId,
+  DatasetSliceWaitTarget,
+  DatasetSplitWaitTarget,
   DatasetUploadWaitTarget,
   ExportId,
   ExportJob,