diff --git a/studio/backend/core/data_recipe/jobs/manager.py b/studio/backend/core/data_recipe/jobs/manager.py index eae789ed50..6f86196bcc 100644 --- a/studio/backend/core/data_recipe/jobs/manager.py +++ b/studio/backend/core/data_recipe/jobs/manager.py @@ -35,11 +35,14 @@ class Subscription: def format_sse(self, event: dict) -> bytes: """Turn event dict into SSE bytes (id/event/data).""" - self._next_id += 1 + event_id = event.get("seq") + if event_id is None: + self._next_id += 1 + event_id = self._next_id body = json.dumps(event, separators=(",", ":"), ensure_ascii=False) event_type = event.get("type") or "message" return ( - f"id: {self._next_id}\n" + f"id: {event_id}\n" f"event: {event_type}\n" f"data: {body}\n\n" ).encode("utf-8") @@ -55,6 +58,7 @@ class JobManager: self._events: deque[dict] = deque(maxlen=5000) self._subs: list[queue.Queue] = [] self._pump_thread: threading.Thread | None = None + self._seq: int = 0 def start(self, *, recipe: dict, run: dict) -> str: """Spawn the job subprocess (one at a time, no cap).""" @@ -65,6 +69,7 @@ class JobManager: job_id = uuid.uuid4().hex self._job = Job(job_id=job_id, status="pending", started_at=time.time()) self._events.clear() + self._seq = 0 mp_q = _CTX.Queue() proc = _CTX.Process( @@ -162,14 +167,18 @@ class JobManager: return None return self._job.analysis - def subscribe(self, job_id: str) -> Subscription | None: + def subscribe(self, job_id: str, *, after_seq: int | None = None) -> Subscription | None: """SSE subscribe: get replay buffer + live events stream.""" with self._lock: if self._job is None or self._job.job_id != job_id: return None q: queue.Queue = queue.Queue(maxsize=2000) self._subs.append(q) - return Subscription(replay=list(self._events), _q=q) + if after_seq is None: + replay = list(self._events) + else: + replay = [e for e in self._events if int(e.get("seq") or 0) > after_seq] + return Subscription(replay=replay, _q=q) def unsubscribe(self, sub: Subscription) -> None: """Drop SSE subscriber (client disconnected).""" @@ -178,6 +187,8 @@ class JobManager: def _emit(self, event: dict) -> None: """Broadcast event to replay buffer + all subscribers.""" + self._seq += 1 + event["seq"] = self._seq self._events.append(event) stale: list[queue.Queue] = [] for q in self._subs: @@ -290,4 +301,3 @@ def get_job_manager() -> JobManager: if _JOB_MANAGER is None: _JOB_MANAGER = JobManager() return _JOB_MANAGER - diff --git a/studio/backend/core/data_recipe/jobs/parse.py b/studio/backend/core/data_recipe/jobs/parse.py index 612b491f41..00c05833f7 100644 --- a/studio/backend/core/data_recipe/jobs/parse.py +++ b/studio/backend/core/data_recipe/jobs/parse.py @@ -30,21 +30,21 @@ class ParsedUpdate: # 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 _RE_SAMPLERS = re.compile( - r"Preparing samplers to generate (?P\\d+) records across (?P\\d+) columns" + r"Preparing samplers to generate (?P\d+) records across (?P\d+) columns" ) _RE_COLCFG = re.compile(r"model config for column '(?P[^']+)'") _RE_PROCESSING_COL = re.compile(r"Processing .* column '(?P[^']+)'") _RE_PROGRESS = re.compile( - r"progress: (?P\\d+)/(?P\\d+) \\((?P\\d+)%\\) complete, " - r"(?P\\d+) ok, (?P\\d+) failed, (?P[0-9.]+) rec/s, eta (?P[0-9.]+)s" + r"progress: (?P\d+)/(?P\d+) \((?P\d+)%\) complete, " + r"(?P\d+) ok, (?P\d+) failed, (?P[0-9.]+) rec/s, eta (?P[0-9.]+)s" ) -_RE_BATCH = re.compile(r"Processing batch (?P\\d+) of (?P\\d+)") -_RE_USAGE_MODEL = re.compile(r"model:\\s*(?P.+)$") +_RE_BATCH = re.compile(r"Processing batch (?P\d+) of (?P\d+)") +_RE_USAGE_MODEL = re.compile(r"model:\s*(?P.+)$") _RE_USAGE_TOKENS = re.compile( - r"tokens:\\s*input=(?P\\d+),\\s*output=(?P\\d+),\\s*total=(?P\\d+),\\s*tps=(?P[0-9.]+)" + r"tokens:\s*input=(?P\d+),\s*output=(?P\d+),\s*total=(?P\d+),\s*tps=(?P[0-9.]+)" ) _RE_USAGE_REQUESTS = re.compile( - r"requests:\\s*success=(?P\\d+),\\s*failed=(?P\\d+),\\s*total=(?P\\d+),\\s*rpm=(?P[0-9.]+)" + r"requests:\s*success=(?P\d+),\s*failed=(?P\d+),\s*total=(?P\d+),\s*rpm=(?P[0-9.]+)" ) @@ -197,4 +197,3 @@ def apply_update(job: Job, update: ParsedUpdate) -> None: def coerce_event(obj: Any) -> dict: # worker sends dict already return obj if isinstance(obj, dict) else {"type": "log", "message": str(obj)} - diff --git a/studio/backend/core/data_recipe/jobs/worker.py b/studio/backend/core/data_recipe/jobs/worker.py index 1fb54454ab..f13e3a8b73 100644 --- a/studio/backend/core/data_recipe/jobs/worker.py +++ b/studio/backend/core/data_recipe/jobs/worker.py @@ -37,23 +37,22 @@ def run_job_process( Subprocess entrypoint. Sends events to `event_queue`. """ - handler = _QueueLogHandler(event_queue) - handler.setLevel(logging.INFO) - - # Attach once to root. data_designer.* loggers should propagate to root. - root = logging.getLogger() - root.addHandler(handler) - root.setLevel(logging.INFO) - dd = logging.getLogger("data_designer") - dd.setLevel(logging.INFO) - dd.propagate = True - event_queue.put({"type": "job.started", "ts": time.time()}) try: - # lazy import so backend can boot even if data-designer isn't installed yet + # Importing data_designer.interface.* triggers DataDesigner logging setup (it clears root handlers), + # so attach our queue handler after that import. from data_designer.config.run_config import RunConfig + import data_designer.interface.data_designer # noqa: F401 + + handler = _QueueLogHandler(event_queue) + handler.setLevel(logging.INFO) + root = logging.getLogger() + root.addHandler(handler) + root.setLevel(logging.INFO) + logging.getLogger("data_designer").setLevel(logging.INFO) + rows = int(run.get("rows") or 1000) dataset_name = str(run.get("dataset_name") or "dataset") run_config_raw = run.get("run_config") or {} @@ -86,4 +85,3 @@ def run_job_process( "stack": traceback.format_exc(limit=20), } ) - diff --git a/studio/backend/routes/data_recipe.py b/studio/backend/routes/data_recipe.py index a7bf901b0d..90c78d16ef 100644 --- a/studio/backend/routes/data_recipe.py +++ b/studio/backend/routes/data_recipe.py @@ -132,7 +132,24 @@ def job_analysis(job_id: str): @router.get("/jobs/{job_id}/events") async def job_events(request: Request, job_id: str): mgr = get_job_manager() - sub = mgr.subscribe(job_id) + last_id = request.headers.get("last-event-id") + after_seq: int | None = None + if last_id: + try: + after_seq = int(str(last_id).strip()) + except Exception: + after_seq = None + + # EventSource can't set custom headers on first connect after a full page refresh, + # so allow resume via query param too: /events?after= + after_q = request.query_params.get("after") + if after_q: + try: + after_seq = int(str(after_q).strip()) + except Exception: + pass + + sub = mgr.subscribe(job_id, after_seq=after_seq) if sub is None: raise HTTPException(status_code=404, detail="job not found") @@ -152,4 +169,3 @@ async def job_events(request: Request, job_id: str): mgr.unsubscribe(sub) return StreamingResponse(gen(), media_type="text/event-stream") -