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

refactor(data-recipe): centralize json+stage constants, tighten parser/errors, sync seed ui

Browse source
This commit is contained in:
Shine1i 2026-02-26 14:06:53 +01:00
commit 00a869f837
8 changed files with 139 additions and 135 deletions
Download patch file Download diff file Expand all files Collapse all files

30
studio/backend/core/data_recipe/jobs/constants.py Normal file
View file

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

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

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

49
studio/backend/core/data_recipe/jobs/parse.py
View file

@ -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<rows>\d+) records across (?P<cols>\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:

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

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

30
studio/backend/core/data_recipe/jsonable.py Normal file
View file

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

33
studio/backend/core/data_recipe/service.py
View file

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

16
studio/frontend/src/features/recipe-studio/components/block-sheet.tsx
View file

@ -173,7 +173,8 @@ function BlockSheetButton({
) : trailing === "drag" ? (
<HugeiconsIcon
icon={DragDropVerticalIcon}
className="size-5 text-primary/80"
strokeWidth={3.5}
className="size-5 text-foreground"
/>
) : null}
</button>
@ -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={() => {

1
studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx
View file

@ -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({