From 00a869f837f250d62f0fce705c9f9cfe45d9424b Mon Sep 17 00:00:00 2001 From: Shine1i Date: Thu, 26 Feb 2026 14:06:53 +0100 Subject: [PATCH] refactor(data-recipe): centralize json+stage constants, tighten parser/errors, sync seed ui --- .../core/data_recipe/jobs/constants.py | 30 +++++++++ .../backend/core/data_recipe/jobs/manager.py | 67 +++++++------------ studio/backend/core/data_recipe/jobs/parse.py | 49 +++++++------- .../backend/core/data_recipe/jobs/worker.py | 48 ++++--------- studio/backend/core/data_recipe/jsonable.py | 30 +++++++++ studio/backend/core/data_recipe/service.py | 33 ++------- .../recipe-studio/components/block-sheet.tsx | 16 +++-- .../dialogs/seed/seed-dialog.tsx | 1 - 8 files changed, 139 insertions(+), 135 deletions(-) create mode 100644 studio/backend/core/data_recipe/jobs/constants.py create mode 100644 studio/backend/core/data_recipe/jsonable.py diff --git a/studio/backend/core/data_recipe/jobs/constants.py b/studio/backend/core/data_recipe/jobs/constants.py new file mode 100644 index 0000000000..5d4081abbc --- /dev/null +++ b/studio/backend/core/data_recipe/jobs/constants.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +# stages parsed from data-designer logs +STAGE_CREATE = "create" +STAGE_PREVIEW = "preview" +STAGE_DAG = "dag" +STAGE_HEALTHCHECK = "healthcheck" +STAGE_SAMPLING = "sampling" +STAGE_COLUMN_CONFIG = "column_config" +STAGE_GENERATING = "generating" +STAGE_BATCH = "batch" +STAGE_PROFILING = "profiling" + +USAGE_RESET_STAGES = { + STAGE_CREATE, + STAGE_PREVIEW, + STAGE_DAG, + STAGE_HEALTHCHECK, + STAGE_SAMPLING, + STAGE_GENERATING, + STAGE_PROFILING, +} + +# job event types emitted by worker/manager +EVENT_JOB_ENQUEUED = "job.enqueued" +EVENT_JOB_STARTED = "job.started" +EVENT_JOB_CANCELLING = "job.cancelling" +EVENT_JOB_CANCELLED = "job.cancelled" +EVENT_JOB_COMPLETED = "job.completed" +EVENT_JOB_ERROR = "job.error" diff --git a/studio/backend/core/data_recipe/jobs/manager.py b/studio/backend/core/data_recipe/jobs/manager.py index eb8c10bb81..dbaa620004 100644 --- a/studio/backend/core/data_recipe/jobs/manager.py +++ b/studio/backend/core/data_recipe/jobs/manager.py @@ -13,6 +13,15 @@ from typing import Any import multiprocessing as mp +from ..jsonable import to_jsonable +from .constants import ( + EVENT_JOB_CANCELLING, + EVENT_JOB_CANCELLED, + EVENT_JOB_COMPLETED, + EVENT_JOB_ENQUEUED, + EVENT_JOB_ERROR, + EVENT_JOB_STARTED, +) from .parse import apply_update, coerce_event, parse_log_message from .types import Job from .worker import run_job_process @@ -20,29 +29,6 @@ from .worker import run_job_process _CTX = mp.get_context("spawn") -def _to_jsonable(value: Any) -> Any: - try: - import numpy as np # type: ignore - except Exception: # pragma: no cover - np = None # type: ignore - - if np is not None: - if isinstance(value, np.ndarray): - return value.tolist() - if isinstance(value, np.generic): - return value.item() - - if isinstance(value, dict): - return {str(k): _to_jsonable(v) for k, v in value.items()} - if isinstance(value, (list, tuple, set)): - return [_to_jsonable(v) for v in value] - if hasattr(value, "isoformat") and callable(value.isoformat): - try: - return value.isoformat() - except Exception: - pass - return value - @dataclass class Subscription: @@ -123,7 +109,7 @@ class JobManager: self._pump_thread = threading.Thread(target=self._pump_loop, daemon=True) self._pump_thread.start() - self._emit({"type": "job.enqueued", "ts": time.time(), "job_id": job_id}) + self._emit({"type": EVENT_JOB_ENQUEUED, "ts": time.time(), "job_id": job_id}) return job_id def cancel(self, job_id: str) -> bool: @@ -134,15 +120,15 @@ class JobManager: if self._proc is None or not self._proc.is_alive(): return True self._job.status = "cancelling" - self._emit({"type": "job.cancelling", "ts": time.time(), "job_id": job_id}) + self._emit({"type": EVENT_JOB_CANCELLING, "ts": time.time(), "job_id": job_id}) try: self._proc.terminate() - except Exception: + except (AttributeError, OSError): pass return True def get_status(self, job_id: str) -> dict | None: - """UI-friendly snapshot. Poll this if you don't want SSE.""" + """UI friendly snapshot that we need. Alternative to sse kinda of and structured""" with self._lock: if self._job is None or self._job.job_id != job_id: return None @@ -304,7 +290,7 @@ class JobManager: ).fetchdf() finally: conn.close() - except Exception: + except (RuntimeError, ValueError, duckdb.Error): return None for helper_col in ("filename", "__row_num__"): @@ -312,7 +298,7 @@ class JobManager: dataframe = dataframe.drop(columns=[helper_col]) rows = dataframe.to_dict(orient="records") - return {"dataset": _to_jsonable(rows), "total": total} + return {"dataset": to_jsonable(rows), "total": total} @staticmethod def _load_dataset_page_with_data_designer( @@ -326,7 +312,7 @@ class JobManager: 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 {"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.""" @@ -355,7 +341,7 @@ class JobManager: for q in self._subs: try: q.put_nowait(event) - except Exception: + except queue.Full: stale.append(q) if stale: self._subs = [q for q in self._subs if q not in stale] @@ -374,7 +360,7 @@ class JobManager: return coerce_event(q.get(timeout=timeout_sec)) except queue.Empty: return None - except Exception: + except (EOFError, OSError, ValueError): return None @staticmethod @@ -386,7 +372,7 @@ class JobManager: events.append(coerce_event(q.get_nowait())) except queue.Empty: return events - except Exception: + except (EOFError, OSError, ValueError): return events def _pump_loop(self) -> None: @@ -416,13 +402,10 @@ class JobManager: self._job.status = "error" self._job.error = self._job.error or "process exited" self._job.finished_at = time.time() - self._emit( - { - "type": f"job.{self._job.status}", - "ts": time.time(), - "job_id": self._job.job_id, - } + event_type = ( + EVENT_JOB_CANCELLED if self._job.status == "cancelled" else EVENT_JOB_ERROR ) + self._emit({"type": event_type, "ts": time.time(), "job_id": self._job.job_id}) return def _handle_event(self, job: Job, event: dict) -> None: @@ -433,9 +416,9 @@ class JobManager: with self._lock: if self._job is None or self._job.job_id != job.job_id: return - if et == "job.started": + if et == EVENT_JOB_STARTED: self._job.status = "active" - if et == "job.completed": + if et == EVENT_JOB_COMPLETED: self._job.status = "completed" self._job.finished_at = time.time() self._job.analysis = event.get("analysis") @@ -445,7 +428,7 @@ class JobManager: if self._job.progress.total and self._job.progress.total > 0: self._job.progress.done = self._job.progress.total self._job.progress.percent = 100.0 - if et == "job.error": + if et == EVENT_JOB_ERROR: self._job.status = "error" self._job.finished_at = time.time() self._job.error = event.get("error") or "error" diff --git a/studio/backend/core/data_recipe/jobs/parse.py b/studio/backend/core/data_recipe/jobs/parse.py index 99d1a85a79..6e2142adf2 100644 --- a/studio/backend/core/data_recipe/jobs/parse.py +++ b/studio/backend/core/data_recipe/jobs/parse.py @@ -4,6 +4,18 @@ import re from dataclasses import dataclass from typing import Any +from .constants import ( + STAGE_BATCH, + STAGE_COLUMN_CONFIG, + STAGE_CREATE, + STAGE_DAG, + STAGE_GENERATING, + STAGE_HEALTHCHECK, + STAGE_PREVIEW, + STAGE_PROFILING, + STAGE_SAMPLING, + USAGE_RESET_STAGES, +) from .types import Job, ModelUsage, Progress @@ -27,8 +39,7 @@ class ParsedUpdate: usage_rpm: float | None = None usage_section_start: bool | None = None -# welp, best effort to parse the logs and convert them to structured information so we can access it and read it properly on the client -# i couldnt find what datadesigner do for progress tracking besides the logs so il probably raise a pr in their repo to add some sort of progress monitoring +# kinda of a bummber but currently only option, Best effort parser from data-designer logs -> structured status for UI. _RE_SAMPLERS = re.compile( r"Preparing samplers to generate (?P\d+) records across (?P\d+) columns" ) @@ -52,31 +63,31 @@ def parse_log_message(msg: str) -> ParsedUpdate | None: m = _RE_SAMPLERS.search(msg) if m: return ParsedUpdate( - stage="sampling", + stage=STAGE_SAMPLING, rows=int(m.group("rows")), cols=int(m.group("cols")), ) if "Sorting column configs into a Directed Acyclic Graph" in msg: - return ParsedUpdate(stage="dag") + return ParsedUpdate(stage=STAGE_DAG) if "Running health checks for models" in msg: - return ParsedUpdate(stage="healthcheck") + return ParsedUpdate(stage=STAGE_HEALTHCHECK) if "Preview generation in progress" in msg: - return ParsedUpdate(stage="preview") + return ParsedUpdate(stage=STAGE_PREVIEW) if "Creating Data Designer dataset" in msg: - return ParsedUpdate(stage="create") + return ParsedUpdate(stage=STAGE_CREATE) if "Measuring dataset column statistics" in msg: - return ParsedUpdate(stage="profiling") + return ParsedUpdate(stage=STAGE_PROFILING) m = _RE_COLCFG.search(msg) if m: col = m.group("col") - return ParsedUpdate(stage="column_config", current_column=col) + return ParsedUpdate(stage=STAGE_COLUMN_CONFIG, current_column=col) m = _RE_PROCESSING_COL.search(msg) if m: col = m.group("col") - return ParsedUpdate(stage="generating", current_column=col) + return ParsedUpdate(stage=STAGE_GENERATING, current_column=col) m = _RE_PROGRESS.search(msg) if m: @@ -89,12 +100,12 @@ def parse_log_message(msg: str) -> ParsedUpdate | None: rate=float(m.group("rate")), eta_sec=float(m.group("eta")), ) - return ParsedUpdate(stage="generating", progress=p) + return ParsedUpdate(stage=STAGE_GENERATING, progress=p) m = _RE_BATCH.search(msg) if m: return ParsedUpdate( - stage="batch", + stage=STAGE_BATCH, batch_idx=int(m.group("idx")), batch_total=int(m.group("total")), ) @@ -132,7 +143,7 @@ def apply_update(job: Job, update: ParsedUpdate) -> None: job.stage = update.stage if update.current_column is not None: job.current_column = update.current_column - if update.stage == "generating" and update.current_column not in job._seen_generation_columns: + if update.stage == STAGE_GENERATING and update.current_column not in job._seen_generation_columns: job._seen_generation_columns.append(update.current_column) if update.rows is not None: job.rows = update.rows @@ -146,16 +157,8 @@ def apply_update(job: Job, update: ParsedUpdate) -> None: if update.batch_total is not None: job.batch.total = update.batch_total - if update.stage in { - "profiling", - "generating", - "sampling", - "healthcheck", - "dag", - "create", - "preview", - }: - # usage summary is a short block; reset once we move into the next stage. + if update.stage in USAGE_RESET_STAGES: + # usage summary is a short block so we reset once we move into the next stage. job._in_usage_summary = False if update.usage_section_start is not None: diff --git a/studio/backend/core/data_recipe/jobs/worker.py b/studio/backend/core/data_recipe/jobs/worker.py index ac27cd0d0b..8c0996b140 100644 --- a/studio/backend/core/data_recipe/jobs/worker.py +++ b/studio/backend/core/data_recipe/jobs/worker.py @@ -7,6 +7,8 @@ import traceback from pathlib import Path from typing import Any +from ..jsonable import to_jsonable +from .constants import EVENT_JOB_COMPLETED, EVENT_JOB_ERROR, EVENT_JOB_STARTED from ..service import build_config_builder, create_data_designer _PROJECT_ROOT = Path(__file__).resolve().parents[5] @@ -28,36 +30,10 @@ class _QueueLogHandler(logging.Handler): "message": record.getMessage(), } self._q.put(event) - except Exception: + except (OSError, RuntimeError, ValueError): pass -def _to_jsonable(value: Any) -> Any: - try: - import numpy as np # type: ignore - except Exception: # pragma: no cover - np = None # type: ignore - - if np is not None: - if isinstance(value, np.ndarray): - return value.tolist() - if isinstance(value, np.generic): - return value.item() - - if isinstance(value, dict): - return {str(k): _to_jsonable(v) for k, v in value.items()} - if isinstance(value, (list, tuple, set)): - return [_to_jsonable(v) for v in value] - - if hasattr(value, "isoformat") and callable(value.isoformat): - try: - return value.isoformat() - except Exception: - pass - - return value - - def run_job_process( *, event_queue, @@ -68,7 +44,7 @@ def run_job_process( Subprocess entrypoint. Sends events to `event_queue`. """ - event_queue.put({"type": "job.started", "ts": time.time()}) + event_queue.put({"type": EVENT_JOB_STARTED, "ts": time.time()}) try: from data_designer.config.run_config import RunConfig @@ -103,21 +79,21 @@ def run_job_process( analysis = ( None if results.analysis is None - else _to_jsonable(results.analysis.model_dump(mode="json")) + else to_jsonable(results.analysis.model_dump(mode="json")) ) dataset = ( [] if results.dataset is None - else _to_jsonable(results.dataset.to_dict(orient="records")) + else to_jsonable(results.dataset.to_dict(orient="records")) ) processor_artifacts = ( None if results.processor_artifacts is None - else _to_jsonable(results.processor_artifacts) + else to_jsonable(results.processor_artifacts) ) event_queue.put( { - "type": "job.completed", + "type": EVENT_JOB_COMPLETED, "ts": time.time(), "analysis": analysis, "dataset": dataset, @@ -128,13 +104,13 @@ 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")) + 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( { - "type": "job.completed", + "type": EVENT_JOB_COMPLETED, "ts": time.time(), "analysis": analysis, "artifact_path": artifact_path, @@ -144,7 +120,7 @@ def run_job_process( except Exception as exc: event_queue.put( { - "type": "job.error", + "type": EVENT_JOB_ERROR, "ts": time.time(), "error": str(exc), "stack": traceback.format_exc(limit=20), @@ -160,7 +136,7 @@ def _merge_batches_to_single_parquet(base_dataset_path: Path) -> None: try: from data_designer.config.utils.io_helpers import read_parquet_dataset - except Exception: + except ImportError: return dataframe = read_parquet_dataset(parquet_dir) diff --git a/studio/backend/core/data_recipe/jsonable.py b/studio/backend/core/data_recipe/jsonable.py new file mode 100644 index 0000000000..aa6e1d6b2e --- /dev/null +++ b/studio/backend/core/data_recipe/jsonable.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +from typing import Any + + +def to_jsonable(value: Any) -> Any: + """Convert numpy/pandas-ish values into plain JSON-safe values.""" + try: + import numpy as np # type: ignore + except ImportError: # pragma: no cover + np = None # type: ignore + + if np is not None: + if isinstance(value, np.ndarray): + return value.tolist() + if isinstance(value, np.generic): + return value.item() + + if isinstance(value, dict): + return {str(k): to_jsonable(v) for k, v in value.items()} + if isinstance(value, (list, tuple, set)): + return [to_jsonable(v) for v in value] + + if hasattr(value, "isoformat") and callable(value.isoformat): + try: + return value.isoformat() + except (TypeError, ValueError): + return value + + return value diff --git a/studio/backend/core/data_recipe/service.py b/studio/backend/core/data_recipe/service.py index b04c6bbf00..2d11cb2845 100644 --- a/studio/backend/core/data_recipe/service.py +++ b/studio/backend/core/data_recipe/service.py @@ -3,32 +3,7 @@ from __future__ import annotations import os from typing import Any -def _to_jsonable(value: Any) -> Any: - # pydantic/fastapi can't serialize numpy arrays/scalars. - try: - import numpy as np # type: ignore - except Exception: # pragma: no cover - np = None # type: ignore - - if np is not None: - if isinstance(value, np.ndarray): - return value.tolist() - if isinstance(value, np.generic): - return value.item() - - if isinstance(value, dict): - return {str(k): _to_jsonable(v) for k, v in value.items()} - if isinstance(value, (list, tuple, set)): - return [_to_jsonable(v) for v in value] - - # pandas Timestamp/date-like - if hasattr(value, "isoformat") and callable(value.isoformat): - try: - return value.isoformat() - except Exception: - pass - - return value +from .jsonable import to_jsonable def build_model_providers(recipe: dict[str, Any]): @@ -158,17 +133,17 @@ def preview_recipe( dataset: list[dict[str, Any]] = [] if results.dataset is not None: raw_rows = results.dataset.to_dict(orient="records") - dataset = [_to_jsonable(row) for row in raw_rows] + dataset = [to_jsonable(row) for row in raw_rows] artifacts = ( None if results.processor_artifacts is None - else _to_jsonable(results.processor_artifacts) + else to_jsonable(results.processor_artifacts) ) analysis = ( None if results.analysis is None - else _to_jsonable(results.analysis.model_dump(mode="json")) + else to_jsonable(results.analysis.model_dump(mode="json")) ) return dataset, artifacts, analysis diff --git a/studio/frontend/src/features/recipe-studio/components/block-sheet.tsx b/studio/frontend/src/features/recipe-studio/components/block-sheet.tsx index 976de9f9a4..be301a36fd 100644 --- a/studio/frontend/src/features/recipe-studio/components/block-sheet.tsx +++ b/studio/frontend/src/features/recipe-studio/components/block-sheet.tsx @@ -173,7 +173,8 @@ function BlockSheetButton({ ) : trailing === "drag" ? ( ) : null} @@ -278,8 +279,7 @@ export function BlockSheet({ event.dataTransfer.setData("text/plain", serialized); event.dataTransfer.effectAllowed = "copy"; }; - const getTrailing = (kind: SheetKind): "drag" | "none" => - kind === "sampler" || kind === "seed" || kind === "llm" ? "drag" : "none"; + const getTrailing = (_kind: SheetKind): "drag" => "drag"; const onBlockClick = (kind: SheetKind, type: BlockType) => { setSheetOpen(false); if (kind === "sampler") { @@ -388,9 +388,17 @@ export function BlockSheet({ title={item.title} description={item.description} isActive={index === 0} + draggable={item.kind === "expression" || item.kind === "note"} + onDragStart={ + item.kind === "expression" && expressionBlocks[0] + ? buildDragStart("expression", expressionBlocks[0].type) + : item.kind === "note" && noteBlocks[0] + ? buildDragStart("note", noteBlocks[0].type) + : undefined + } trailing={ item.kind === "expression" || item.kind === "note" - ? "none" + ? "drag" : "chevron" } onClick={() => { diff --git a/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx b/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx index 1959c28c45..187f48b331 100644 --- a/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx +++ b/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx @@ -253,7 +253,6 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl dataset_name: datasetName, hf_token: config.hf_token?.trim() || undefined, subset: undefined, - split: "train", preview_size: 10, }); onUpdate({