ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor: add batch processing support with configuration options and execution enhancements

Browse source
This commit is contained in:
Shine1i 2026-02-23 21:32:20 +01:00
commit 1323e0af53
15 changed files with 634 additions and 295 deletions
Download patch file Download diff file Expand all files Collapse all files

87
studio/backend/core/data_recipe/jobs/manager.py
View file

@ -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:

37
studio/backend/core/data_recipe/jobs/worker.py
View file

@ -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)

2
studio/backend/routes/data_recipe.py
View file

@ -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'")

43
studio/frontend/src/components/ui/combobox.tsx
View file

@ -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 (
<ComboboxPrimitive.Portal>
<ComboboxPrimitive.Positioner
side={side}
sideOffset={sideOffset}
function ComboboxContent({
className,
side = "bottom",
sideOffset = 6,
align = "start",
alignOffset = 0,
anchor,
container,
...props
}: ComboboxPrimitive.Popup.Props &
Pick<
ComboboxPrimitive.Positioner.Props,
"side" | "align" | "sideOffset" | "alignOffset" | "anchor"
> & {
container?: HTMLElement | null;
}): React.ReactElement {
const dialogContainer = useDialogPortalContainer();
return (
<ComboboxPrimitive.Portal container={container ?? dialogContainer ?? undefined}>
<ComboboxPrimitive.Positioner
side={side}
sideOffset={sideOffset}
align={align}
alignOffset={alignOffset}
anchor={anchor}

70
studio/frontend/src/components/ui/dialog.tsx
View file

@ -2,12 +2,19 @@
import { Dialog as DialogPrimitive } from "radix-ui";
import type * as React from "react";
import { createContext, useContext } from "react";
import { Button } from "@/components/ui/button";
import { cn } from "@/lib/utils";
import { Cancel01Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
const DialogPortalContainerContext = createContext<HTMLElement | null>(null);
export function useDialogPortalContainer(): HTMLElement | null {
return useContext(DialogPortalContainerContext);
}
function Dialog({
...props
}: React.ComponentProps<typeof DialogPrimitive.Root>) {
@ -68,36 +75,39 @@ function DialogContent({
overlayClassName?: string;
overlayPosition?: "fixed" | "absolute";
}) {
const resolvedContainer = container ?? null;
return (
<DialogPortal container={container ?? undefined}>
<DialogOverlay
className={overlayClassName}
position={overlayPosition ?? position}
/>
<DialogPrimitive.Content
data-slot="dialog-content"
className={cn(
"bg-background data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 data-closed:zoom-out-95 data-open:zoom-in-95 ring-foreground/5 grid max-w-[calc(100%-2rem)] gap-6 rounded-4xl p-6 text-sm ring-1 duration-100 sm:max-w-md top-1/2 left-1/2 z-50 w-full -translate-x-1/2 -translate-y-1/2",
position === "fixed" ? "fixed" : "absolute",
className,
)}
{...props}
>
{children}
{showCloseButton && (
<DialogPrimitive.Close data-slot="dialog-close" asChild>
<Button
variant="ghost"
className="absolute top-4 right-4"
size="icon-sm"
>
<HugeiconsIcon icon={Cancel01Icon} strokeWidth={2} />
<span className="sr-only">Close</span>
</Button>
</DialogPrimitive.Close>
)}
</DialogPrimitive.Content>
</DialogPortal>
<DialogPortalContainerContext.Provider value={resolvedContainer}>
<DialogPortal container={resolvedContainer ?? undefined}>
<DialogOverlay
className={overlayClassName}
position={overlayPosition ?? position}
/>
<DialogPrimitive.Content
data-slot="dialog-content"
className={cn(
"bg-background data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 data-closed:zoom-out-95 data-open:zoom-in-95 ring-foreground/5 grid max-w-[calc(100%-2rem)] gap-6 rounded-4xl p-6 text-sm ring-1 duration-100 sm:max-w-md top-1/2 left-1/2 z-50 w-full -translate-x-1/2 -translate-y-1/2",
position === "fixed" ? "fixed" : "absolute",
className,
)}
{...props}
>
{children}
{showCloseButton && (
<DialogPrimitive.Close data-slot="dialog-close" asChild>
<Button
variant="ghost"
className="absolute top-4 right-4"
size="icon-sm"
>
<HugeiconsIcon icon={Cancel01Icon} strokeWidth={2} />
<span className="sr-only">Close</span>
</Button>
</DialogPrimitive.Close>
)}
</DialogPrimitive.Content>
</DialogPortal>
</DialogPortalContainerContext.Provider>
);
}
@ -123,7 +133,7 @@ function DialogFooter({
<div
data-slot="dialog-footer"
className={cn(
"gap-2 flex flex-col-reverse gap-2 sm:flex-row sm:justify-end",
"flex flex-col-reverse gap-2 sm:flex-row sm:justify-end",
className,
)}
{...props}

31
studio/frontend/src/components/ui/select.tsx
View file

@ -4,7 +4,8 @@ import { Select as SelectPrimitive } from "radix-ui";
import type * as React from "react";
import { createContext, useContext, useState } from "react";
import { cn } from "@/lib/utils";
import { cn } from "@/lib/utils";
import { useDialogPortalContainer } from "@/components/ui/dialog";
import {
ArrowDown01Icon,
ArrowUp01Icon,
@ -91,18 +92,22 @@ function SelectTrigger({
);
}
function SelectContent({
className,
children,
position = "item-aligned",
align = "center",
...props
}: React.ComponentProps<typeof SelectPrimitive.Content>) {
return (
<SelectPrimitive.Portal>
<SelectPrimitive.Content
data-slot="select-content"
data-align-trigger={position === "item-aligned"}
function SelectContent({
className,
children,
position = "item-aligned",
align = "center",
container,
...props
}: React.ComponentProps<typeof SelectPrimitive.Content> & {
container?: HTMLElement | null;
}) {
const dialogContainer = useDialogPortalContainer();
return (
<SelectPrimitive.Portal container={container ?? dialogContainer ?? undefined}>
<SelectPrimitive.Content
data-slot="select-content"
data-align-trigger={position === "item-aligned"}
className={cn(
"bg-popover text-popover-foreground data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 data-closed:zoom-out-95 data-open:zoom-in-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 shadow-border ring-1 ring-border min-w-36 rounded-xl p-1 corner-squircle duration-100 relative z-50 max-h-(--radix-select-content-available-height) origin-(--radix-select-content-transform-origin) overflow-x-hidden overflow-y-auto ",
position === "popper" &&

8
studio/frontend/src/features/recipe-studio/components/executions/execution-sidebar.tsx
View file

@ -2,6 +2,7 @@ import type { ReactElement } from "react";
import { Badge } from "@/components/ui/badge";
import { cn } from "@/lib/utils";
import type { RecipeExecutionRecord } from "../../execution-types";
import { isExecutionInProgress } from "../../executions/execution-helpers";
import {
formatStatus,
formatTimestamp,
@ -58,6 +59,13 @@ export function ExecutionSidebar({
</Badge>
</div>
<p className="text-xs text-muted-foreground">{execution.rows} rows</p>
{isExecutionInProgress(execution.status) &&
typeof execution.batch?.total === "number" &&
execution.batch.total > 1 && (
<p className="text-xs text-muted-foreground">
Batch {execution.batch.idx ?? "--"}/{execution.batch.total}
</p>
)}
<p className="text-xs text-muted-foreground">
{formatTimestamp(execution.createdAt)}
</p>

18
studio/frontend/src/features/recipe-studio/components/executions/executions-view.tsx
View file

@ -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({
: ""}
</span>
)}
{showBatchProgress && (
<span>
Batch {batchIdx ?? "--"}/{batchTotal}
</span>
)}
{isStale && <Badge variant="outline">Recipe changed since this run</Badge>}
</div>
@ -392,6 +400,16 @@ export function ExecutionsView({
{formatPercent(selectedExecution.column_progress.percent)})
</p>
)}
{showBatchProgress && (
<p
className={cn(
"text-xs",
progressComplete ? "text-emerald-900" : "text-amber-900",
)}
>
Processed batch: {batchIdx ?? "--"}/{batchTotal}
</p>
)}
</div>
)}

585
studio/frontend/src/features/recipe-studio/dialogs/preview-dialog.tsx
View file

@ -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 (
<div className="grid gap-2">
<FieldLabel label={label} htmlFor={id} hint={hint} />
<Input
id={id}
type="text"
inputMode={inputMode}
value={value}
onChange={(event) => onChange(event.target.value)}
onBlur={onBlur}
placeholder={placeholder}
/>
</div>
);
}
function ValidationResultPanel({
validateResult,
}: {
validateResult: ValidationResult;
}): ReactElement | null {
if (!validateResult) {
return null;
}
return (
<div
className={
validateResult.valid
? "space-y-1 rounded-xl border border-emerald-300 bg-emerald-50 p-3"
: "space-y-1 rounded-xl border border-destructive/30 bg-destructive/5 p-3"
}
>
<p
className={
validateResult.valid
? "text-xs font-semibold uppercase text-emerald-700"
: "text-xs font-semibold uppercase text-destructive"
}
>
{validateResult.valid ? "Validation passed" : "Validation failed"}
</p>
{!validateResult.valid && validateResult.errors.length > 0 && (
<div className="space-y-1">
{validateResult.errors.map((error) => (
<p key={error} className="text-xs text-destructive">
{error}
</p>
))}
</div>
)}
{!validateResult.valid && validateResult.rawDetail && (
<p className="text-xs text-destructive">{validateResult.rawDetail}</p>
)}
</div>
);
}
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 (
<Dialog open={open} onOpenChange={onOpenChange}>
@ -115,7 +294,7 @@ export function RunDialog({
</p>
</DialogHeader>
<label className="flex items-center justify-between rounded-xl border bg-muted/20 px-3 py-2 text-sm">
<div className="flex items-center justify-between rounded-xl border bg-muted/20 px-3 py-2 text-sm">
<span className="font-medium text-foreground">Preview mode</span>
<Switch
checked={kind === "preview"}
@ -123,51 +302,27 @@ export function RunDialog({
onKindChange(checked ? "preview" : "full")
}
/>
</label>
</div>
<div className="grid gap-4 md:grid-cols-3">
<div className="grid gap-4 md:grid-cols-2">
<div className="grid gap-2">
<FieldLabel
label="Records"
htmlFor="run-rows"
hint={rowHint}
/>
<FieldLabel label="Records" htmlFor="run-rows" hint={rowHint} />
<Input
id="run-rows"
type="number"
min={1}
max={200000}
value={String(rows)}
onChange={(event) => {
const parsed = Number(event.target.value);
if (!Number.isFinite(parsed)) {
return;
}
onRowsChange(clampInt(parsed, 1, 200000));
}}
/>
</div>
<div className="grid gap-2">
<FieldLabel
label="Batch size"
htmlFor="run-buffer-size"
hint="Rows handled per batch. Bigger can be faster; smaller uses less memory."
/>
<Input
id="run-buffer-size"
type="number"
min={1}
max={200000}
value={String(settings.bufferSize)}
onChange={(event) => {
const parsed = Number(event.target.value);
if (!Number.isFinite(parsed)) {
return;
}
onSettingsChange({
bufferSize: clampInt(parsed, 1, 200000),
});
}}
type="text"
inputMode="numeric"
value={rowsDraft}
onChange={(event) => setRowsDraft(event.target.value)}
onBlur={() =>
commitInt(
rowsDraft,
rows,
1,
MAX_RECORDS,
onRowsChange,
setRowsDraft,
)
}
/>
</div>
<div className="grid gap-2">
@ -178,29 +333,83 @@ export function RunDialog({
/>
<Input
id="run-llm-parallel"
type="number"
min={1}
max={2048}
type="text"
inputMode="numeric"
placeholder="Use model config"
value={settings.llmParallelRequests ?? ""}
onChange={(event) => {
const value = event.target.value.trim();
if (!value) {
value={llmParallelDraft}
onChange={(event) => setLlmParallelDraft(event.target.value)}
onBlur={() => {
const trimmed = llmParallelDraft.trim();
if (!trimmed) {
onSettingsChange({ llmParallelRequests: null });
setLlmParallelDraft("");
return;
}
const parsed = Number(value);
const parsed = Number(trimmed);
if (!Number.isFinite(parsed)) {
setLlmParallelDraft(
settings.llmParallelRequests === null
? ""
: String(settings.llmParallelRequests),
);
return;
}
onSettingsChange({
llmParallelRequests: clampInt(parsed, 1, 2048),
});
const next = clampInt(parsed, 1, MAX_WORKERS);
onSettingsChange({ llmParallelRequests: next });
setLlmParallelDraft(String(next));
}}
/>
</div>
</div>
{kind === "full" && (
<div className="space-y-3">
<div className="flex items-center justify-between gap-3 text-sm">
<span className="font-medium">Enable batching</span>
<Switch
checked={settings.batchEnabled}
onCheckedChange={(checked) =>
onSettingsChange({ batchEnabled: Boolean(checked) })
}
/>
</div>
{settings.batchEnabled && (
<div className="space-y-3">
<DraftInputField
id="run-batch-size"
label="Batch size"
hint="Rows handled per batch during generation."
inputMode="numeric"
value={batchSizeDraft}
onChange={setBatchSizeDraft}
onBlur={() =>
commitInt(
batchSizeDraft,
settings.batchSize,
1,
MAX_RECORDS,
(value) => onSettingsChange({ batchSize: value }),
setBatchSizeDraft,
)
}
/>
<div className="flex items-center justify-between gap-3 text-sm">
<span className="font-medium">
Merge batches to one parquet
</span>
<Switch
checked={settings.mergeBatches}
onCheckedChange={(checked) =>
onSettingsChange({ mergeBatches: Boolean(checked) })
}
/>
</div>
</div>
)}
</div>
)}
<Collapsible open={advancedOpen} onOpenChange={setAdvancedOpen}>
<CollapsibleTrigger asChild={true}>
<button
@ -212,123 +421,101 @@ export function RunDialog({
</CollapsibleTrigger>
<CollapsibleContent className="mt-3 space-y-4">
<div className="grid gap-4 md:grid-cols-2">
<div className="grid gap-2">
<FieldLabel
label="CPU workers"
htmlFor="run-non-inference-workers"
hint="Worker threads for non-LLM steps like samplers and expressions."
/>
<Input
id="run-non-inference-workers"
type="number"
min={1}
max={2048}
value={String(settings.nonInferenceWorkers)}
onChange={(event) => {
const parsed = Number(event.target.value);
if (!Number.isFinite(parsed)) {
return;
}
onSettingsChange({
nonInferenceWorkers: clampInt(parsed, 1, 2048),
});
}}
/>
</div>
<div className="grid gap-2">
<FieldLabel
label="Error window"
htmlFor="run-shutdown-window"
hint="How many attempts to observe before early-stop checks kick in."
/>
<Input
id="run-shutdown-window"
type="number"
min={1}
max={10000}
value={String(settings.shutdownErrorWindow)}
onChange={(event) => {
const parsed = Number(event.target.value);
if (!Number.isFinite(parsed)) {
return;
}
onSettingsChange({
shutdownErrorWindow: clampInt(parsed, 1, 10000),
});
}}
/>
</div>
<div className="grid gap-2">
<FieldLabel
label="Conversation restarts"
htmlFor="run-max-restarts"
hint="How many full retries to do if model output fails validation."
/>
<Input
id="run-max-restarts"
type="number"
min={0}
max={100}
value={String(settings.maxConversationRestarts)}
onChange={(event) => {
const parsed = Number(event.target.value);
if (!Number.isFinite(parsed)) {
return;
}
onSettingsChange({
maxConversationRestarts: clampInt(parsed, 0, 100),
});
}}
/>
</div>
<div className="grid gap-2">
<FieldLabel
label="Correction steps"
htmlFor="run-correction-steps"
hint="Extra in-chat fix attempts before a full retry."
/>
<Input
id="run-correction-steps"
type="number"
min={0}
max={100}
value={String(settings.maxConversationCorrectionSteps)}
onChange={(event) => {
const parsed = Number(event.target.value);
if (!Number.isFinite(parsed)) {
return;
}
onSettingsChange({
maxConversationCorrectionSteps: clampInt(parsed, 0, 100),
});
}}
/>
</div>
<div className="grid gap-2">
<FieldLabel
label="Shutdown error rate"
htmlFor="run-shutdown-rate"
hint="Stop early if failure rate passes this value. Example: 0.5 = 50%."
/>
<Input
id="run-shutdown-rate"
type="number"
min={0}
max={1}
step={0.05}
value={String(settings.shutdownErrorRate)}
onChange={(event) => {
const parsed = Number(event.target.value);
if (!Number.isFinite(parsed)) {
return;
}
onSettingsChange({
shutdownErrorRate: clampFloat(parsed, 0, 1),
});
}}
/>
</div>
<label className="flex items-center gap-3 text-sm text-foreground">
<DraftInputField
id="run-non-inference-workers"
label="CPU workers"
hint="Worker threads for non-LLM steps like samplers and expressions."
inputMode="numeric"
value={workersDraft}
onChange={setWorkersDraft}
onBlur={() =>
commitInt(
workersDraft,
settings.nonInferenceWorkers,
1,
MAX_WORKERS,
(value) => onSettingsChange({ nonInferenceWorkers: value }),
setWorkersDraft,
)
}
/>
<DraftInputField
id="run-shutdown-window"
label="Error window"
hint="How many attempts to observe before early-stop checks kick in."
inputMode="numeric"
value={windowDraft}
onChange={setWindowDraft}
onBlur={() =>
commitInt(
windowDraft,
settings.shutdownErrorWindow,
1,
MAX_SHUTDOWN_WINDOW,
(value) => onSettingsChange({ shutdownErrorWindow: value }),
setWindowDraft,
)
}
/>
<DraftInputField
id="run-max-restarts"
label="Conversation restarts"
hint="How many full retries to do if model output fails validation."
inputMode="numeric"
value={restartsDraft}
onChange={setRestartsDraft}
onBlur={() =>
commitInt(
restartsDraft,
settings.maxConversationRestarts,
0,
MAX_RETRY_STEPS,
(value) =>
onSettingsChange({ maxConversationRestarts: value }),
setRestartsDraft,
)
}
/>
<DraftInputField
id="run-correction-steps"
label="Correction steps"
hint="Extra in-chat fix attempts before a full retry."
inputMode="numeric"
value={correctionsDraft}
onChange={setCorrectionsDraft}
onBlur={() =>
commitInt(
correctionsDraft,
settings.maxConversationCorrectionSteps,
0,
MAX_RETRY_STEPS,
(value) =>
onSettingsChange({
maxConversationCorrectionSteps: value,
}),
setCorrectionsDraft,
)
}
/>
<DraftInputField
id="run-shutdown-rate"
label="Shutdown error rate"
hint="Stop early if failure rate passes this value. Example: 0.5 = 50%."
inputMode="decimal"
value={shutdownRateDraft}
onChange={setShutdownRateDraft}
onBlur={() =>
commitFloat(
shutdownRateDraft,
settings.shutdownErrorRate,
0,
1,
(value) => onSettingsChange({ shutdownErrorRate: value }),
setShutdownRateDraft,
)
}
/>
<div className="flex items-center gap-3 text-sm text-foreground">
<Switch
checked={settings.disableEarlyShutdown}
onCheckedChange={(checked) =>
@ -338,7 +525,7 @@ export function RunDialog({
}
/>
Disable early shutdown
</label>
</div>
</div>
</CollapsibleContent>
</Collapsible>
@ -356,37 +543,7 @@ export function RunDialog({
</div>
)}
{validateResult && (
<div
className={
validateResult.valid
? "space-y-1 rounded-xl border border-emerald-300 bg-emerald-50 p-3"
: "space-y-1 rounded-xl border border-destructive/30 bg-destructive/5 p-3"
}
>
<p
className={
validateResult.valid
? "text-xs font-semibold uppercase text-emerald-700"
: "text-xs font-semibold uppercase text-destructive"
}
>
{validateResult.valid ? "Validation passed" : "Validation failed"}
</p>
{!validateResult.valid && validateResult.errors.length > 0 && (
<div className="space-y-1">
{validateResult.errors.map((error) => (
<p key={error} className="text-xs text-destructive">
{error}
</p>
))}
</div>
)}
{!validateResult.valid && validateResult.rawDetail && (
<p className="text-xs text-destructive">{validateResult.rawDetail}</p>
)}
</div>
)}
<ValidationResultPanel validateResult={validateResult} />
<DialogFooter>
<Button

6
studio/frontend/src/features/recipe-studio/execution-types.ts
View file

@ -22,6 +22,11 @@ export type RecipeExecutionProgress = {
failed?: number | null;
};
export type RecipeExecutionBatch = {
idx?: number | null;
total?: number | null;
};
export type RecipeExecutionAnalysis = {
num_records?: number;
target_num_records?: number;
@ -50,6 +55,7 @@ export type RecipeExecutionRecord = {
progress: RecipeExecutionProgress | null;
// biome-ignore lint/style/useNamingConvention: backend schema
column_progress: RecipeExecutionProgress | null;
batch: RecipeExecutionBatch | null;
// biome-ignore lint/style/useNamingConvention: backend schema
model_usage: Record<string, unknown> | null;
// biome-ignore lint/style/useNamingConvention: backend schema

1
studio/frontend/src/features/recipe-studio/executions/execution-helpers.ts
View file

@ -139,6 +139,7 @@ export function withExecutionDefaults(
datasetPage,
datasetPageSize,
column_progress: record.column_progress ?? null,
batch: record.batch ?? null,
};
}

16
studio/frontend/src/features/recipe-studio/executions/run-settings.ts
View file

@ -61,7 +61,9 @@ export function sanitizeExecutionRows(
export function normalizeRunSettings(settings: RecipeRunSettings): RecipeRunSettings {
return {
bufferSize: toPositiveInt(settings.bufferSize, 1000, 1, 200_000),
batchSize: toPositiveInt(settings.batchSize, 1000, 1, 200_000),
batchEnabled: Boolean(settings.batchEnabled),
mergeBatches: Boolean(settings.mergeBatches),
llmParallelRequests:
typeof settings.llmParallelRequests === "number"
? toPositiveInt(settings.llmParallelRequests, 4, 1, 2048)
@ -90,10 +92,13 @@ export function normalizeRunSettings(settings: RecipeRunSettings): RecipeRunSett
function buildRunConfigPayload(
settings: RecipeRunSettings,
rows: number,
kind: RecipeExecutionKind,
): Record<string, unknown> {
const useBatching = kind === "full" && settings.batchEnabled;
return {
// biome-ignore lint/style/useNamingConvention: backend schema
buffer_size: settings.bufferSize,
buffer_size: useBatching ? settings.batchSize : toPositiveInt(rows, 1000, 1, 200_000),
// biome-ignore lint/style/useNamingConvention: backend schema
non_inference_max_parallel_workers: settings.nonInferenceWorkers,
// biome-ignore lint/style/useNamingConvention: backend schema
@ -163,7 +168,12 @@ export function buildExecutionPayload(input: {
// biome-ignore lint/style/useNamingConvention: backend schema
execution_type: input.kind,
// biome-ignore lint/style/useNamingConvention: backend schema
run_config: buildRunConfigPayload(normalizedSettings),
run_config: buildRunConfigPayload(normalizedSettings, input.rows, input.kind),
// biome-ignore lint/style/useNamingConvention: backend schema
merge_batches:
input.kind === "full" &&
normalizedSettings.batchEnabled &&
normalizedSettings.mergeBatches,
},
};
}

15
studio/frontend/src/features/recipe-studio/executions/runtime.ts
View file

@ -1,5 +1,9 @@
import type { JobEvent, JobStatusResponse } from "../api";
import type { RecipeExecutionKind, RecipeExecutionRecord } from "../execution-types";
import type {
RecipeExecutionBatch,
RecipeExecutionKind,
RecipeExecutionRecord,
} from "../execution-types";
import {
DATASET_PAGE_SIZE,
mapJobStatus,
@ -70,6 +74,13 @@ export function applyExecutionStatusSnapshot(
status: JobStatusResponse,
): RecipeExecutionRecord {
const mappedStatus = mapJobStatus(status.status);
const batchRaw = normalizeObject(status.batch);
const batch: RecipeExecutionBatch | null = batchRaw
? {
idx: typeof batchRaw.idx === "number" ? batchRaw.idx : null,
total: typeof batchRaw.total === "number" ? batchRaw.total : null,
}
: null;
return {
...execution,
status: mappedStatus,
@ -80,6 +91,7 @@ export function applyExecutionStatusSnapshot(
column_progress:
(normalizeObject(status.column_progress) as RecipeExecutionRecord["column_progress"]) ??
null,
batch,
model_usage: normalizeObject(status.model_usage),
artifact_path: status.artifact_path ?? execution.artifact_path,
error: status.error ?? null,
@ -113,6 +125,7 @@ export function createBaseExecutionRecord(input: {
current_column: null,
progress: null,
column_progress: null,
batch: null,
model_usage: null,
lastEventId: null,
artifact_path: null,

8
studio/frontend/src/features/recipe-studio/stores/recipe-executions.ts
View file

@ -4,7 +4,9 @@ import type { RecipeExecutionRecord } from "../execution-types";
import { sortExecutions, withExecutionDefaults } from "../executions/execution-helpers";
export type RecipeRunSettings = {
bufferSize: number;
batchSize: number;
batchEnabled: boolean;
mergeBatches: boolean;
llmParallelRequests: number | null;
nonInferenceWorkers: number;
maxConversationRestarts: number;
@ -15,7 +17,9 @@ export type RecipeRunSettings = {
};
const DEFAULT_RUN_SETTINGS: RecipeRunSettings = {
bufferSize: 1000,
batchSize: 1000,
batchEnabled: true,
mergeBatches: false,
llmParallelRequests: null,
nonInferenceWorkers: 4,
maxConversationRestarts: 5,

2
studio/frontend/src/features/recipe-studio/utils/payload/types.ts
View file

@ -26,6 +26,8 @@ export type RecipePayload = {
dataset_name?: string;
// biome-ignore lint/style/useNamingConvention: backend schema
artifact_path?: string;
// biome-ignore lint/style/useNamingConvention: backend schema
merge_batches?: boolean;
};
ui: {
nodes: { id: string; x: number; y: number }[];