From 1323e0af53ada80faf439b4648ab3cdccf89181c Mon Sep 17 00:00:00 2001 From: Shine1i Date: Mon, 23 Feb 2026 21:32:20 +0100 Subject: [PATCH] refactor: add batch processing support with configuration options and execution enhancements --- .../backend/core/data_recipe/jobs/manager.py | 87 ++- .../backend/core/data_recipe/jobs/worker.py | 37 +- studio/backend/routes/data_recipe.py | 2 + .../frontend/src/components/ui/combobox.tsx | 43 +- studio/frontend/src/components/ui/dialog.tsx | 70 ++- studio/frontend/src/components/ui/select.tsx | 31 +- .../executions/execution-sidebar.tsx | 8 + .../components/executions/executions-view.tsx | 18 + .../recipe-studio/dialogs/preview-dialog.tsx | 585 +++++++++++------- .../features/recipe-studio/execution-types.ts | 6 + .../executions/execution-helpers.ts | 1 + .../recipe-studio/executions/run-settings.ts | 16 +- .../recipe-studio/executions/runtime.ts | 15 +- .../recipe-studio/stores/recipe-executions.ts | 8 +- .../recipe-studio/utils/payload/types.ts | 2 + 15 files changed, 634 insertions(+), 295 deletions(-) diff --git a/studio/backend/core/data_recipe/jobs/manager.py b/studio/backend/core/data_recipe/jobs/manager.py index a71e8eac8e..eb8c10bb81 100644 --- a/studio/backend/core/data_recipe/jobs/manager.py +++ b/studio/backend/core/data_recipe/jobs/manager.py @@ -108,10 +108,12 @@ class JobManager: self._events.clear() self._seq = 0 + run_payload = dict(run) + run_payload["_job_id"] = job_id mp_q = _CTX.Queue() proc = _CTX.Process( target=run_job_process, - kwargs={"event_queue": mp_q, "recipe": recipe, "run": run}, + kwargs={"event_queue": mp_q, "recipe": recipe, "run": run_payload}, daemon=True, ) proc.start() @@ -246,15 +248,86 @@ class JobManager: if not parquet_dir.exists(): return {"error": f"dataset path missing: {parquet_dir}"} - from data_designer.config.utils.io_helpers import read_parquet_dataset - - dataframe = read_parquet_dataset(parquet_dir) - total = int(len(dataframe.index)) - rows = dataframe.iloc[offset:offset + limit].to_dict(orient="records") - return {"dataset": _to_jsonable(rows), "total": total} + return self._load_dataset_page(parquet_dir=parquet_dir, limit=limit, offset=offset) except Exception as exc: return {"error": f"dataset load failed: {exc}"} + @staticmethod + def _load_dataset_page( + *, + parquet_dir: Path, + limit: int, + offset: int, + ) -> dict[str, Any]: + dataset_page = JobManager._load_dataset_page_with_duckdb( + parquet_dir=parquet_dir, + limit=limit, + offset=offset, + ) + if dataset_page is not None: + return dataset_page + return JobManager._load_dataset_page_with_data_designer( + parquet_dir=parquet_dir, + limit=limit, + offset=offset, + ) + + @staticmethod + def _load_dataset_page_with_duckdb( + *, + parquet_dir: Path, + limit: int, + offset: int, + ) -> dict[str, Any] | None: + parquet_glob = str((parquet_dir / "*.parquet").resolve()) + try: + import duckdb # type: ignore + except Exception: + return None + + try: + conn = duckdb.connect(":memory:") + try: + total_row = conn.execute( + "SELECT COUNT(*) FROM read_parquet(?)", + [parquet_glob], + ).fetchone() + total = int(total_row[0] if total_row else 0) + dataframe = conn.execute( + ( + "SELECT *, row_number() OVER (PARTITION BY filename) AS __row_num__ " + "FROM read_parquet(?, filename=true) " + "ORDER BY filename, __row_num__ " + "LIMIT ? OFFSET ?" + ), + [parquet_glob, int(limit), int(offset)], + ).fetchdf() + finally: + conn.close() + except Exception: + return None + + for helper_col in ("filename", "__row_num__"): + if helper_col in dataframe.columns: + dataframe = dataframe.drop(columns=[helper_col]) + + rows = dataframe.to_dict(orient="records") + return {"dataset": _to_jsonable(rows), "total": total} + + @staticmethod + def _load_dataset_page_with_data_designer( + *, + parquet_dir: Path, + limit: int, + offset: int, + ) -> dict[str, Any]: + from data_designer.config.utils.io_helpers import read_parquet_dataset + + dataframe = read_parquet_dataset(parquet_dir) + total = int(len(dataframe.index)) + rows = dataframe.iloc[offset:offset + limit].to_dict(orient="records") + return {"dataset": _to_jsonable(rows), "total": total} + def subscribe(self, job_id: str, *, after_seq: int | None = None) -> Subscription | None: """SSE subscribe: get replay buffer + live events stream.""" with self._lock: diff --git a/studio/backend/core/data_recipe/jobs/worker.py b/studio/backend/core/data_recipe/jobs/worker.py index c39ef286c4..2d33a1c632 100644 --- a/studio/backend/core/data_recipe/jobs/worker.py +++ b/studio/backend/core/data_recipe/jobs/worker.py @@ -1,12 +1,17 @@ from __future__ import annotations import logging +import shutil import time import traceback +from pathlib import Path from typing import Any from ..service import build_config_builder, create_data_designer +_PROJECT_ROOT = Path(__file__).resolve().parents[5] +_ARTIFACT_ROOT = _PROJECT_ROOT / "datasets" / "recipes" + class _QueueLogHandler(logging.Handler): def __init__(self, event_queue): @@ -69,15 +74,16 @@ def run_job_process( from data_designer.config.run_config import RunConfig rows = int(run.get("rows") or 1000) - dataset_name = str(run.get("dataset_name") or "dataset") - artifact_path_raw = run.get("artifact_path") - artifact_path = None - if isinstance(artifact_path_raw, str) and artifact_path_raw.strip(): - artifact_path = artifact_path_raw.strip() + job_id = str(run.get("_job_id") or "").strip() + if not job_id: + job_id = f"{int(time.time())}" + dataset_name = f"recipe_{job_id}" + merge_batches = bool(run.get("merge_batches")) + _ARTIFACT_ROOT.mkdir(parents=True, exist_ok=True) run_config_raw = run.get("run_config") or {} builder = build_config_builder(recipe) - designer = create_data_designer(recipe, artifact_path=artifact_path) + designer = create_data_designer(recipe, artifact_path=str(_ARTIFACT_ROOT)) # DataDesigner configures root logging in DataDesigner.__init__. # Attach queue logger directly to `data_designer` so parser events survive root resets. @@ -123,6 +129,8 @@ def run_job_process( else: results = designer.create(builder, num_records=rows, dataset_name=dataset_name) analysis = _to_jsonable(results.load_analysis().model_dump(mode="json")) + if merge_batches: + _merge_batches_to_single_parquet(results.artifact_storage.base_dataset_path) artifact_path = str(results.artifact_storage.base_dataset_path) event_queue.put( { @@ -142,3 +150,20 @@ def run_job_process( "stack": traceback.format_exc(limit=20), } ) + + +def _merge_batches_to_single_parquet(base_dataset_path: Path) -> None: + parquet_dir = base_dataset_path / "parquet-files" + parquet_files = sorted(parquet_dir.glob("*.parquet")) + if len(parquet_files) <= 1: + return + + try: + from data_designer.config.utils.io_helpers import read_parquet_dataset + except Exception: + return + + dataframe = read_parquet_dataset(parquet_dir) + shutil.rmtree(parquet_dir) + parquet_dir.mkdir(parents=True, exist_ok=True) + dataframe.to_parquet(parquet_dir / "batch_00000.parquet", index=False) diff --git a/studio/backend/routes/data_recipe.py b/studio/backend/routes/data_recipe.py index e05347b3d4..8a713f2cec 100644 --- a/studio/backend/routes/data_recipe.py +++ b/studio/backend/routes/data_recipe.py @@ -397,6 +397,8 @@ def create_job(payload: RecipePayload): raise HTTPException(status_code=400, detail="Recipe must include columns.") run: dict[str, Any] = payload.run or {} + run.pop("artifact_path", None) + run.pop("dataset_name", None) execution_type = str(run.get("execution_type") or "full").strip().lower() if execution_type not in {"preview", "full"}: raise HTTPException(status_code=400, detail="invalid execution_type: must be 'preview' or 'full'") diff --git a/studio/frontend/src/components/ui/combobox.tsx b/studio/frontend/src/components/ui/combobox.tsx index 0893e9b218..9c1e970c57 100644 --- a/studio/frontend/src/components/ui/combobox.tsx +++ b/studio/frontend/src/components/ui/combobox.tsx @@ -6,7 +6,8 @@ import { Combobox as ComboboxPrimitive } from "@base-ui/react"; import * as React from "react"; import { createContext, useContext, useState } from "react"; -import { Button } from "@/components/ui/button"; +import { Button } from "@/components/ui/button"; +import { useDialogPortalContainer } from "@/components/ui/dialog"; import { InputGroup, InputGroupAddon, @@ -139,24 +140,28 @@ function ComboboxInput({ ); } -function ComboboxContent({ - className, - side = "bottom", - sideOffset = 6, - align = "start", - alignOffset = 0, - anchor, - ...props -}: ComboboxPrimitive.Popup.Props & - Pick< - ComboboxPrimitive.Positioner.Props, - "side" | "align" | "sideOffset" | "alignOffset" | "anchor" - >): React.ReactElement { - return ( - - & { + container?: HTMLElement | null; + }): React.ReactElement { + const dialogContainer = useDialogPortalContainer(); + return ( + + (null); + +export function useDialogPortalContainer(): HTMLElement | null { + return useContext(DialogPortalContainerContext); +} + function Dialog({ ...props }: React.ComponentProps) { @@ -68,36 +75,39 @@ function DialogContent({ overlayClassName?: string; overlayPosition?: "fixed" | "absolute"; }) { + const resolvedContainer = container ?? null; return ( - - - - {children} - {showCloseButton && ( - - - - )} - - + + + + + {children} + {showCloseButton && ( + + + + )} + + + ); } @@ -123,7 +133,7 @@ function DialogFooter({
) { - return ( - - & { + container?: HTMLElement | null; +}) { + const dialogContainer = useDialogPortalContainer(); + return ( + +

{execution.rows} rows

+ {isExecutionInProgress(execution.status) && + typeof execution.batch?.total === "number" && + execution.batch.total > 1 && ( +

+ Batch {execution.batch.idx ?? "--"}/{execution.batch.total} +

+ )}

{formatTimestamp(execution.createdAt)}

diff --git a/studio/frontend/src/features/recipe-studio/components/executions/executions-view.tsx b/studio/frontend/src/features/recipe-studio/components/executions/executions-view.tsx index 285017a624..2daa8b7e11 100644 --- a/studio/frontend/src/features/recipe-studio/components/executions/executions-view.tsx +++ b/studio/frontend/src/features/recipe-studio/components/executions/executions-view.tsx @@ -264,6 +264,9 @@ export function ExecutionsView({ (selectedExecution ? isExecutionInProgress(selectedExecution.status) : false); const progressComplete = selectedExecution?.status === "completed"; const progressPercent = selectedExecution?.progress?.percent ?? (progressComplete ? 100 : 0); + const batchTotal = selectedExecution?.batch?.total ?? null; + const batchIdx = selectedExecution?.batch?.idx ?? null; + const showBatchProgress = typeof batchTotal === "number" && batchTotal > 1; const terminalLines = selectedExecution?.log_lines ?? []; const rawExecution = useMemo(() => { if (!selectedExecution) { @@ -328,6 +331,11 @@ export function ExecutionsView({ : ""} )} + {showBatchProgress && ( + + Batch {batchIdx ?? "--"}/{batchTotal} + + )} {isStale && Recipe changed since this run} @@ -392,6 +400,16 @@ export function ExecutionsView({ {formatPercent(selectedExecution.column_progress.percent)})

)} + {showBatchProgress && ( +

+ Processed batch: {batchIdx ?? "--"}/{batchTotal} +

+ )} )} diff --git a/studio/frontend/src/features/recipe-studio/dialogs/preview-dialog.tsx b/studio/frontend/src/features/recipe-studio/dialogs/preview-dialog.tsx index 39563e2278..296172777f 100644 --- a/studio/frontend/src/features/recipe-studio/dialogs/preview-dialog.tsx +++ b/studio/frontend/src/features/recipe-studio/dialogs/preview-dialog.tsx @@ -1,9 +1,3 @@ -import { type ReactElement, useMemo, useState } from "react"; -import { - CookBookIcon, - TestTube01Icon, -} from "@hugeicons/core-free-icons"; -import { HugeiconsIcon } from "@hugeicons/react"; import { Button } from "@/components/ui/button"; import { Collapsible, @@ -19,6 +13,9 @@ import { } from "@/components/ui/dialog"; import { Input } from "@/components/ui/input"; import { Switch } from "@/components/ui/switch"; +import { CookBookIcon, TestTube01Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { type ReactElement, useEffect, useState } from "react"; import type { RecipeExecutionKind } from "../execution-types"; import type { RecipeRunSettings } from "../stores/recipe-executions"; import { FieldLabel } from "./shared/field-label"; @@ -45,6 +42,13 @@ type RunDialogProps = { container?: HTMLDivElement | null; }; +type ValidationResult = RunDialogProps["validateResult"]; + +const MAX_RECORDS = 200_000; +const MAX_WORKERS = 2_048; +const MAX_SHUTDOWN_WINDOW = 10_000; +const MAX_RETRY_STEPS = 100; + function clampInt(value: number, min: number, max: number): number { if (!Number.isFinite(value)) { return min; @@ -72,6 +76,131 @@ function clampFloat(value: number, min: number, max: number): number { return value; } +function commitInt( + raw: string, + current: number, + min: number, + max: number, + apply: (value: number) => void, + setDraft: (value: string) => void, +): void { + const trimmed = raw.trim(); + if (!trimmed) { + setDraft(String(current)); + return; + } + const parsed = Number(trimmed); + if (!Number.isFinite(parsed)) { + setDraft(String(current)); + return; + } + const next = clampInt(parsed, min, max); + apply(next); + setDraft(String(next)); +} + +function commitFloat( + raw: string, + current: number, + min: number, + max: number, + apply: (value: number) => void, + setDraft: (value: string) => void, +): void { + const trimmed = raw.trim(); + if (!trimmed) { + setDraft(String(current)); + return; + } + const parsed = Number(trimmed); + if (!Number.isFinite(parsed)) { + setDraft(String(current)); + return; + } + const next = clampFloat(parsed, min, max); + apply(next); + setDraft(String(next)); +} + +type DraftInputFieldProps = { + id: string; + label: string; + hint: string; + inputMode: "numeric" | "decimal"; + value: string; + onChange: (value: string) => void; + onBlur: () => void; + placeholder?: string; +}; + +function DraftInputField({ + id, + label, + hint, + inputMode, + value, + onChange, + onBlur, + placeholder, +}: DraftInputFieldProps): ReactElement { + return ( +
+ + onChange(event.target.value)} + onBlur={onBlur} + placeholder={placeholder} + /> +
+ ); +} + +function ValidationResultPanel({ + validateResult, +}: { + validateResult: ValidationResult; +}): ReactElement | null { + if (!validateResult) { + return null; + } + + return ( +
+

+ {validateResult.valid ? "Validation passed" : "Validation failed"} +

+ {!validateResult.valid && validateResult.errors.length > 0 && ( +
+ {validateResult.errors.map((error) => ( +

+ {error} +

+ ))} +
+ )} + {!validateResult.valid && validateResult.rawDetail && ( +

{validateResult.rawDetail}

+ )} +
+ ); +} + export function RunDialog({ open, onOpenChange, @@ -91,13 +220,63 @@ export function RunDialog({ }: RunDialogProps): ReactElement { const [advancedOpen, setAdvancedOpen] = useState(false); const kindLabel = kind === "preview" ? "Preview" : "Full run"; - const rowHint = useMemo( - () => - kind === "preview" - ? "How many sample rows to generate for a quick check." - : "How many rows to generate in total.", - [kind], + const rowHint = + kind === "preview" + ? "How many sample rows to generate for a quick check." + : "How many rows to generate in total."; + + const [rowsDraft, setRowsDraft] = useState(String(rows)); + const [batchSizeDraft, setBatchSizeDraft] = useState( + String(settings.batchSize), ); + const [llmParallelDraft, setLlmParallelDraft] = useState( + settings.llmParallelRequests === null + ? "" + : String(settings.llmParallelRequests), + ); + const [workersDraft, setWorkersDraft] = useState( + String(settings.nonInferenceWorkers), + ); + const [windowDraft, setWindowDraft] = useState( + String(settings.shutdownErrorWindow), + ); + const [restartsDraft, setRestartsDraft] = useState( + String(settings.maxConversationRestarts), + ); + const [correctionsDraft, setCorrectionsDraft] = useState( + String(settings.maxConversationCorrectionSteps), + ); + const [shutdownRateDraft, setShutdownRateDraft] = useState( + String(settings.shutdownErrorRate), + ); + + useEffect(() => { + if (!open) { + return; + } + setRowsDraft(String(rows)); + setBatchSizeDraft(String(settings.batchSize)); + setLlmParallelDraft( + settings.llmParallelRequests === null + ? "" + : String(settings.llmParallelRequests), + ); + setWorkersDraft(String(settings.nonInferenceWorkers)); + setWindowDraft(String(settings.shutdownErrorWindow)); + setRestartsDraft(String(settings.maxConversationRestarts)); + setCorrectionsDraft(String(settings.maxConversationCorrectionSteps)); + setShutdownRateDraft(String(settings.shutdownErrorRate)); + }, [ + rows, + settings.batchSize, + settings.llmParallelRequests, + settings.nonInferenceWorkers, + settings.shutdownErrorWindow, + settings.maxConversationRestarts, + settings.maxConversationCorrectionSteps, + settings.shutdownErrorRate, + open, + ]); return ( @@ -115,7 +294,7 @@ export function RunDialog({

-