@toloka-platform/sdk 0.11.4 → 0.12.0

package/src/cli/slices.ts

@@ -0,0 +1,147 @@
+import type { DatasetSlice, PotatoClient } from "../client";
+import { PotatoApiError } from "../errors";
+import type { WaitForOptions } from "../wait-for";
+import { assertCli } from "./cli-error";
+import { resolveLiveDataset } from "./data";
+import { resolveFieldValueFilters } from "./filters";
+import { formatTable } from "./table";
+
+type SliceApi = Pick<PotatoClient, "datasets" | "waitFor">;
+
+export type DatasetSliceCommandResult = { readonly output: string };
+
+export type CreateDatasetSliceOptions = {
+  readonly sourceRoot: string;
+  readonly datasetName: string;
+  readonly projectId?: string;
+  readonly name?: string;
+  readonly filters?: readonly string[];
+  readonly from?: string;
+  readonly freeze?: readonly string[];
+  readonly wait: boolean;
+  readonly waitFor?: WaitForOptions;
+  readonly json: boolean;
+};
+
+export const formatCount = (value: number): string => value.toLocaleString("en-US");
+
+export const fieldNamesOf = async (
+  client: Pick<PotatoClient, "datasets">,
+  datasetId: string,
+): Promise<ReadonlyMap<string, string>> => {
+  const names = new Map<string, string>();
+  for await (const field of client.datasets.listAllFields({ datasetId })) names.set(field.id, field.name);
+  return names;
+};
+
+const describeFilter = (slice: DatasetSlice, fieldNames: ReadonlyMap<string, string>): string => {
+  if (slice.itemFilter === null) return "every row";
+  const { fieldValues = [], fieldSources = [], missingRevisionFields = [] } = slice.itemFilter;
+  return [
+    ...fieldValues.map((filter) => `${fieldNames.get(filter.fieldId) ?? filter.fieldId} = ${JSON.stringify(filter.value)}`),
+    ...fieldSources.map((filter) => `${fieldNames.get(filter.fieldId) ?? filter.fieldId} written by given sources`),
+    ...missingRevisionFields.map((fieldId) => `${fieldNames.get(fieldId) ?? fieldId} is empty`),
+  ].join(", ");
+};
+
+export const formatSlice = (slice: DatasetSlice, fieldNames: ReadonlyMap<string, string>): string => {
+  const rows: string[][] = [
+    ["Slice", slice.id],
+    ["Name", slice.name],
+    ["Dataset", slice.datasetId],
+    ["Status", slice.status],
+    ["Rows", formatCount(slice.rows)],
+    ["Rows matching", describeFilter(slice, fieldNames)],
+    ["Frozen", slice.frozenFieldIds.map((fieldId) => fieldNames.get(fieldId) ?? fieldId).join(", ") || "-"],
+  ];
+  if (slice.fromSliceId !== null) rows.push(["Taken from", slice.fromSliceId]);
+  if (slice.error !== null) rows.push(["Error", `${slice.error.code}: ${slice.error.message}`]);
+  rows.push(["Frozen at", slice.cutoff], ["Created", slice.createdAt]);
+  return formatTable(rows);
+};
+
+export const waitUntilDone = async <Resource>(
+  wait: () => Promise<Resource>,
+  { what, checkCommand }: { readonly what: string; readonly checkCommand: string },
+): Promise<Resource> => {
+  try {
+    return await wait();
+  } catch (error) {
+    assertCli(
+      !(error instanceof PotatoApiError && error.code === "wait_timeout"),
+      `Stopped waiting for ${what}; it keeps running. Check it with toloka ${checkCommand}.`,
+      { hint: { disabled: true } },
+    );
+    throw error;
+  }
+};
+
+const render = async (client: SliceApi, slice: DatasetSlice, json: boolean): Promise<DatasetSliceCommandResult> => ({
+  output: json ? JSON.stringify(slice, null, 2) : formatSlice(slice, await fieldNamesOf(client, slice.datasetId)),
+});
+
+export async function createDatasetSlice(
+  client: PotatoClient,
+  options: CreateDatasetSliceOptions,
+): Promise<DatasetSliceCommandResult> {
+  assertCli(
+    options.from === undefined || options.filters === undefined,
+    "--from takes the source slice again with its own filter, so it takes no --filter.",
+    { hint: { description: "Pass either --from <slice> or --filter field=value." } },
+  );
+  const dataset = await resolveLiveDataset(client, options, false);
+  const frozenColumns = options.freeze && [...options.freeze];
+  const created = await client.datasets.createSlice(
+    options.from === undefined
+      ? {
+          datasetId: dataset.datasetId,
+          name: options.name,
+          ...(options.filters === undefined
+            ? {}
+            : { itemFilter: await resolveFieldValueFilters(client, dataset.datasetId, options.filters) }),
+          frozenColumns,
+        }
+      : { datasetId: dataset.datasetId, name: options.name, fromSlice: options.from, frozenColumns },
+  );
+  const slice = options.wait
+    ? await waitUntilDone(
+        () => client.waitFor({ type: "datasetSlice", sliceId: created.id }, options.waitFor ?? {}),
+        { what: `slice ${created.id}`, checkCommand: `dataset slice show ${created.id}` },
+      )
+    : created;
+  assertCli(slice.status !== "failed", `Slice ${slice.id} failed: ${slice.error?.message ?? "no reason given"}`, {
+    hint: { description: `Check it with toloka dataset slice show ${slice.id}.` },
+  });
+  return render(client, slice, options.json);
+}
+
+export async function listDatasetSlices(
+  client: PotatoClient,
+  options: { readonly sourceRoot: string; readonly datasetName: string; readonly projectId?: string },
+): Promise<DatasetSliceCommandResult> {
+  const dataset = await resolveLiveDataset(client, options, false);
+  const page = await client.datasets.listSlicesPage({ datasetId: dataset.datasetId, limit: 100 });
+  if (page.slices.length === 0) {
+    return { output: `Dataset ${dataset.name} has no slices. Create one with toloka dataset slice create ${dataset.datasetId}.` };
+  }
+  return {
+    output: formatTable([
+      ["SLICE", "NAME", "STATUS", "ROWS", "TAKEN FROM", "CREATED"],
+      ...page.slices.map((slice) => [
+        slice.id,
+        slice.name,
+        slice.status,
+        formatCount(slice.rows),
+        slice.fromSliceId ?? "-",
+        slice.createdAt,
+      ]),
+    ]),
+  };
+}
+
+export async function showDatasetSlice(
+  client: PotatoClient,
+  options: { readonly sliceId: string; readonly json: boolean },
+): Promise<DatasetSliceCommandResult> {
+  return render(client, await client.datasets.getSlice({ sliceId: options.sliceId }), options.json);
+}

package/src/cli/splits.ts

@@ -0,0 +1,168 @@
+import type { DatasetField, PotatoClient } from "../client";
+import type { WaitForOptions } from "../wait-for";
+import { assertCli } from "./cli-error";
+import type { SplitPartSpec } from "./context";
+import { resolveLiveDataset } from "./data";
+import { resolveFieldValueFilters } from "./filters";
+import { fieldNamesOf, formatCount, waitUntilDone } from "./slices";
+import { formatTable } from "./table";
+
+type SplitApi = Pick<PotatoClient, "datasets" | "waitFor">;
+
+export type SplitRatios = { readonly train: number; readonly validate: number; readonly test: number };
+
+export type DatasetSplitCommandResult = { readonly output: string };
+
+export type SplitDatasetOptions = {
+  readonly sourceRoot: string;
+  readonly datasetName: string;
+  readonly projectId?: string;
+  readonly column?: string;
+  readonly ratios?: SplitRatios;
+  readonly parts?: readonly SplitPartSpec[];
+  readonly validateOnTest: boolean;
+  readonly balanceBy?: string;
+  readonly seed?: number;
+  readonly filters?: readonly string[];
+  readonly slice?: string;
+  readonly wait: boolean;
+  readonly waitFor?: WaitForOptions;
+  readonly json: boolean;
+};
+
+const formatPercent = (share: number): string => `${Math.round(share * 100)}%`;
+
+const presetParts = (ratios: SplitRatios, validateOnTest: boolean): SplitPartSpec[] => {
+  const parts: SplitPartSpec[] = [
+    { name: "train", roles: ["train"], share: ratios.train },
+    { name: "validate", roles: ["validate"], share: ratios.validate },
+    { name: "test", roles: validateOnTest ? ["validate", "test"] : ["test"], share: ratios.test },
+  ];
+  return parts.filter((part) => (part.share ?? 0) > 0);
+};
+
+const requestedParts = (options: SplitDatasetOptions): SplitPartSpec[] | undefined => {
+  assertCli(options.ratios === undefined || options.parts === undefined, "Pass --ratios or --parts, not both.", {
+    hint: { description: "--ratios 80/10/10 is the same as --parts train:80,validate:10,test:10." },
+  });
+  assertCli(
+    !options.validateOnTest || options.ratios?.validate === 0,
+    "--validate-on-test makes the test part the validation set too, so it needs a validate ratio of 0.",
+    {
+      hint: {
+        description: "Pass --ratios with validate 0, for example --ratios 80/0/20, or give a part both roles with --parts test=validate+test:20.",
+      },
+    },
+  );
+  return options.ratios === undefined
+    ? options.parts && [...options.parts]
+    : presetParts(options.ratios, options.validateOnTest);
+};
+
+const describeRoles = (roles: readonly string[]): string => roles.join("+") || "no role";
+
+type FieldSplit = NonNullable<DatasetField["split"]>;
+
+const describeParts = (split: FieldSplit): string => {
+  const shareOf = (name: string) => split.recipe?.shares.find((share) => share.name === name)?.share;
+  return split.parts
+    .map((part) => {
+      const share = shareOf(part.name);
+      return `${part.name} ${share === undefined ? "" : `${formatPercent(share)} `}(${describeRoles(part.roles)})`;
+    })
+    .join(", ");
+};
+
+export const describeSplitColumn = (field: DatasetField): string =>
+  field.split === null ? "-" : `${field.split.status} · ${describeParts(field.split)}`;
+
+export const formatSplitColumn = (field: DatasetField, fieldNames: ReadonlyMap<string, string>): string => {
+  const { split } = field;
+  assertCli(split !== null, `Column ${field.name} is not a split column.`, { hint: { disabled: true } });
+  const rows: string[][] = [
+    ["Column", `${field.name} (${field.id})`],
+    ["Status", split.status],
+    ["Parts", describeParts(split)],
+  ];
+  if (split.recipe !== null) {
+    const balanceBy = split.recipe.balanceByFieldId;
+    rows.push(["Balanced by", balanceBy === null ? "-" : (fieldNames.get(balanceBy) ?? balanceBy)]);
+    rows.push(["Seed", String(split.recipe.seed)]);
+  }
+  if (split.progress !== null) {
+    rows.push(["Placed", `${formatCount(split.progress.placed)} of ${formatCount(split.progress.total)}`]);
+  }
+  if (split.lastFill !== null) {
+    rows.push([
+      "Last fill",
+      `placed ${formatCount(split.lastFill.placed)}, ${formatCount(split.lastFill.kept)} already had a part`,
+    ]);
+    rows.push([
+      "Rows",
+      Object.entries(split.lastFill.parts)
+        .map(([name, count]) => `${name} ${formatCount(count)}`)
+        .join(" · "),
+    ]);
+  }
+  if (split.error !== null) rows.push(["Error", `${split.error.code}: ${split.error.message}`]);
+  return formatTable(rows);
+};
+
+export const waitForSplitColumn = async (
+  client: SplitApi,
+  { datasetId, field, waitFor }: { datasetId: string; field: DatasetField; waitFor?: WaitForOptions },
+): Promise<DatasetField> => {
+  const done =
+    field.split?.status === "filling"
+      ? await waitUntilDone(
+          () => client.waitFor({ type: "datasetSplit", datasetId, column: field.name }, waitFor ?? {}),
+          { what: `column ${field.name}`, checkCommand: `dataset columns ${datasetId}` },
+        )
+      : field;
+  assertCli(
+    done.split?.status !== "failed",
+    `Splitting column ${done.name} failed: ${done.split?.error?.message ?? "no reason given"}`,
+    { hint: { description: "Fix the cause, then run toloka dataset split again; rows that already have a part keep it." } },
+  );
+  return done;
+};
+
+export async function splitDataset(client: PotatoClient, options: SplitDatasetOptions): Promise<DatasetSplitCommandResult> {
+  assertCli(options.slice === undefined || options.filters === undefined, "Pass --slice or --filter, not both.", {
+    hint: { description: "--slice fills the rows of a slice; --filter fills the rows that match now." },
+  });
+  const dataset = await resolveLiveDataset(client, options, false);
+  const field = await client.datasets.split({
+    datasetId: dataset.datasetId,
+    column: options.column,
+    parts: requestedParts(options),
+    balanceBy: options.balanceBy,
+    seed: options.seed,
+    ...(options.slice === undefined ? {} : { slice: options.slice }),
+    ...(options.filters === undefined
+      ? {}
+      : { itemFilter: await resolveFieldValueFilters(client, dataset.datasetId, options.filters) }),
+  });
+  const done = options.wait ? await waitForSplitColumn(client, { datasetId: dataset.datasetId, field, waitFor: options.waitFor }) : field;
+  return {
+    output: options.json
+      ? JSON.stringify(done, null, 2)
+      : formatSplitColumn(done, await fieldNamesOf(client, dataset.datasetId)),
+  };
+}
+
+export async function listDatasetColumns(
+  client: PotatoClient,
+  options: { readonly sourceRoot: string; readonly datasetName: string; readonly projectId?: string; readonly json: boolean },
+): Promise<DatasetSplitCommandResult> {
+  const dataset = await resolveLiveDataset(client, options, false);
+  const fields: DatasetField[] = [];
+  for await (const field of client.datasets.listAllFields({ datasetId: dataset.datasetId })) fields.push(field);
+  if (options.json) return { output: JSON.stringify(fields, null, 2) };
+  return {
+    output: formatTable([
+      ["COLUMN", "ID", "SPLIT"],
+      ...fields.filter((field) => !field.internal).map((field) => [field.name, field.id, describeSplitColumn(field)]),
+    ]),
+  };
+}