From ea1cf17f895a6b5046963c0b8a7524ace3d7bea9 Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Wed, 11 Nov 2026 07:35:43 -0800 Subject: [PATCH] feat(python): export recorded Function job timelines (#4946) Function jobs can now export their recorded execution timeline with `job.export_timeline(path)` or `await job.export_timeline(path)`. The SDK fetches a fresh Chrome Trace Event recording that opens in Perfetto, including worker split spans and timestamped job phases. Batch counts and latency statistics remain details rather than invented batch intervals. ```python job = db.open_job(job_id) job.export_timeline("refresh.trace.json") async_job = await async_db.open_job(job_id) await async_job.export_timeline("refresh.trace.json") ``` The export does not wait for completion or cancel the job. Running jobs return partial recordings. Local operations and SQL query handles remain unsupported, and servers must support the timeline endpoint. The destination is touched only after a successful fetch; network and filesystem errors propagate. Rust freshness wrappers forward the operation without changing the submitting table's read baseline. --- python/python/lancedb/_lancedb.pyi | 1 + python/python/lancedb/job.py | 43 ++++++++++++++++++++++++ python/python/tests/test_job.py | 52 ++++++++++++++++++++++++++++++ python/src/job.rs | 7 ++++ rust/lancedb/src/job.rs | 27 ++++++++++++++++ rust/lancedb/src/remote/db.rs | 43 ++++++++++++++++++++++++ rust/lancedb/src/remote/job.rs | 26 +++++++++++++++ rust/lancedb/src/remote/table.rs | 11 +++++++ 8 files changed, 210 insertions(+) diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 4c5f5758a..4029525e9 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -412,6 +412,7 @@ class Job: ) -> Optional[str]: ... async def cancel(self) -> None: ... async def refresh(self) -> None: ... + async def timeline(self) -> str: ... async def events( self, *, diff --git a/python/python/lancedb/job.py b/python/python/lancedb/job.py index 86a1295cc..094a7b8d5 100644 --- a/python/python/lancedb/job.py +++ b/python/python/lancedb/job.py @@ -7,6 +7,8 @@ import asyncio import reprlib import time from datetime import timedelta +from os import PathLike +from pathlib import Path from typing import Any, Callable, Generic, Optional, TypeVar, cast import pyarrow as pa @@ -422,6 +424,31 @@ class AsyncJob(Generic[T]): description = self._inner._description if self._inner is not None else None return getattr(description, name) if description is not None else None + async def export_timeline(self, path: str | PathLike[str]) -> None: + """Write a fresh recorded Function job timeline as Chrome trace JSON. + + Open the file in Perfetto or Chrome tracing. The recording contains + retained split spans grouped by worker and timestamped job phases. + Batch counts and latency statistics are included as details; individual + batch timestamps are not recorded. Running jobs export a partial snapshot + without waiting or cancelling. Download again to include later work. + + Requires a server-side Function job and a server supporting timeline + export. Unsupported operations raise ``NotImplementedError`` or a server + error. The destination is overwritten only after the trace is fetched. + Filesystem errors propagate to the caller. + + Examples + -------- + >>> await job.export_timeline("refresh.trace.json") # doctest: +SKIP + """ + if self._inner is None: + raise NotImplementedError( + "job timeline export requires a server-side Function job" + ) + trace = await self._inner.timeline() + await asyncio.to_thread(Path(path).write_text, trace, encoding="utf-8") + async def events( self, *, @@ -621,6 +648,22 @@ class Job(Generic[T]): def _field(self, name: str) -> Optional[Any]: return getattr(self._inner, name) if self._inner is not None else None + def export_timeline(self, path: str | PathLike[str]) -> None: + """Write Chrome trace JSON for this Function job. + + See :meth:`AsyncJob.export_timeline`. Open the file in Perfetto or + Chrome tracing. + + Examples + -------- + >>> job.export_timeline("refresh.trace.json") # doctest: +SKIP + """ + if self._inner is None: + raise NotImplementedError( + "job timeline export requires a server-side Function job" + ) + return LOOP.run(self._inner.export_timeline(path)) + def events( self, *, diff --git a/python/python/tests/test_job.py b/python/python/tests/test_job.py index 07b8a8c5e..690f6414b 100644 --- a/python/python/tests/test_job.py +++ b/python/python/tests/test_job.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright The LanceDB Authors import asyncio +import json import time from datetime import timedelta from types import SimpleNamespace @@ -751,3 +752,54 @@ async def test_query_wait_timeout_requires_keeping_original_handle(): assert "continues on the server" in message assert "keep this handle to continue waiting" in message assert "db.open_job" not in message + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync", [False, True]) +async def test_export_timeline_writes_fresh_recording_without_waiting(sync, tmp_path): + trace = { + "traceEvents": [{"ph": "X", "ts": 0, "dur": 1000}], + "metadata": {"partial": True}, + } + inner = SimpleNamespace( + timeline=AsyncMock(return_value=json.dumps(trace)), + wait=AsyncMock(), + cancel=AsyncMock(), + ) + job = AsyncJob(inner) + path = tmp_path / "job.trace.json" + if sync: + await asyncio.to_thread(Job(job).export_timeline, str(path)) + else: + await job.export_timeline(path) + assert json.loads(path.read_text()) == trace + inner.timeline.assert_awaited_once_with() + inner.wait.assert_not_awaited() + inner.cancel.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_export_timeline_failure_preserves_existing_file(tmp_path): + path = tmp_path / "job.trace.json" + path.write_text("previous recording") + job = AsyncJob( + SimpleNamespace( + timeline=AsyncMock(side_effect=RuntimeError("unsupported server")) + ) + ) + with pytest.raises(RuntimeError, match="unsupported server"): + await job.export_timeline(path) + assert path.read_text() == "previous recording" + with pytest.raises(NotImplementedError, match="server-side Function job"): + await AsyncJob(None).export_timeline(path) + with pytest.raises(NotImplementedError, match="server-side Function job"): + Job(None).export_timeline(path) + + +@pytest.mark.asyncio +async def test_export_timeline_reports_filesystem_errors(tmp_path): + job = AsyncJob( + SimpleNamespace(timeline=AsyncMock(return_value='{"traceEvents":[]}')) + ) + with pytest.raises(FileNotFoundError): + await job.export_timeline(tmp_path / "missing" / "trace.json") diff --git a/python/src/job.rs b/python/src/job.rs index f1fd85844..0d411072f 100644 --- a/python/src/job.rs +++ b/python/src/job.rs @@ -152,6 +152,13 @@ impl Job { self.inner.description().map(JobDescription::from) } + pub fn timeline(self_: PyRef<'_, Self>) -> PyResult> { + let inner = self_.inner.clone(); + future_into_py(self_.py(), async move { + Ok(inner.timeline().await.infer_error()?.to_string()) + }) + } + #[pyo3(signature = (*, limit=None, filter=None))] pub fn events( self_: PyRef<'_, Self>, diff --git a/rust/lancedb/src/job.rs b/rust/lancedb/src/job.rs index 8644c6a48..d178542b7 100644 --- a/rust/lancedb/src/job.rs +++ b/rust/lancedb/src/job.rs @@ -85,6 +85,11 @@ pub(crate) trait JobHandle: Send + Sync { async fn describe(&self) -> Result { job_detail_not_supported("describing a job") } + /// Chrome Trace Event JSON for a server-side Function job. + async fn timeline(&self) -> Result { + job_detail_not_supported("job timeline export") + } + /// The job's recorded lifecycle events. async fn events(&self, _request: JobEventsRequest) -> Result> { job_detail_not_supported("job event history") @@ -468,6 +473,28 @@ where } } + /// Fetch a recorded Function job timeline as Chrome Trace Event JSON. + /// + /// The snapshot opens in Perfetto or Chrome tracing. It contains retained + /// split and positioned phase spans; batch statistics remain arguments, + /// because individual batch timestamps are not recorded. Running jobs return + /// a partial recording without waiting for completion. Requires server support; + /// in-process operations and SQL query handles do not support this operation. + /// + /// ```no_run + /// # async fn export(job: &lancedb::job::Job) -> Result<(), Box> { + /// let trace = job.timeline().await?; + /// std::fs::write("job.trace.json", trace.to_string())?; + /// # Ok(()) + /// # } + /// ``` + pub async fn timeline(&self) -> Result { + match &self.inner { + JobInner::Handle { handle, .. } => handle.timeline().await, + JobInner::Completed(_) => job_detail_not_supported("job timeline export"), + } + } + /// Waits until the operation reaches a terminal state. /// /// Returns the endpoint's typed result. Unit-result jobs return `()`. diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index a4952593b..a9d55b511 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -4261,6 +4261,49 @@ mod tests { ); } + #[tokio::test] + async fn test_job_timeline_fetches_recorded_trace_without_waiting() { + let conn = Connection::new_with_handler(|request| { + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); + assert_eq!(body["job_id"], "job-1"); + let body = match request.url().path() { + "/v1/jobs/describe" => { + r#"{"job_id":"job-1","job_type":"function_refresh","job_state":"IN_PROGRESS","creation_ms":1}"# + } + "/v1/jobs/timeline" => { + r#"{"traceEvents":[{"ph":"X","pid":1,"tid":1,"ts":0,"dur":1000,"name":"Split 1"}],"metadata":{"partial":true}}"# + } + path => panic!("Unexpected request: {path}"), + }; + http::Response::builder().status(200).body(body).unwrap() + }); + let job = conn.open_job("job-1").await.unwrap(); + let trace = job.timeline().await.unwrap(); + assert_eq!(trace["traceEvents"][0]["dur"], 1000); + assert_eq!(trace["metadata"]["partial"], true); + } + + #[tokio::test] + async fn test_job_timeline_rejects_an_unrelated_json_response() { + let conn = Connection::new_with_handler(|request| { + let body = if request.url().path() == "/v1/jobs/describe" { + r#"{"job_id":"job-1","job_type":"function_refresh","job_state":"SUCCEEDED","creation_ms":1}"# + } else { + r#"{"message":"not a trace"}"# + }; + http::Response::builder().status(200).body(body).unwrap() + }); + let job = conn.open_job("job-1").await.unwrap(); + assert!( + job.timeline() + .await + .unwrap_err() + .to_string() + .contains("not Chrome Trace Event JSON") + ); + } + #[tokio::test] async fn test_job_events_scope_to_that_job() { let schema = Arc::new(Schema::new(vec![Field::new( diff --git a/rust/lancedb/src/remote/job.rs b/rust/lancedb/src/remote/job.rs index 0c652a3ae..d27c9c666 100644 --- a/rust/lancedb/src/remote/job.rs +++ b/rust/lancedb/src/remote/job.rs @@ -235,6 +235,32 @@ impl JobHandle for RemoteJob { Ok(self.fetch_description().await?.1.into_description()) } + async fn timeline(&self) -> Result { + let request = self + .client + .post("/v1/jobs/timeline") + .json(&serde_json::json!({"job_id": self.job_id})); + let (request_id, response) = self.client.send(request).await?; + let response = self.client.check_response(&request_id, response).await?; + let body = response.text().await.err_to_http(request_id.clone())?; + let trace: serde_json::Value = + serde_json::from_str(&body).map_err(|error| Error::Http { + source: format!("failed to parse job timeline: {error}").into(), + request_id: request_id.clone(), + status_code: None, + })?; + if !trace + .get("traceEvents") + .is_some_and(serde_json::Value::is_array) + { + return Err(Error::Http { + source: "server response is not Chrome Trace Event JSON; job timeline export requires server support".into(), + request_id, status_code: None, + }); + } + Ok(trace) + } + async fn events(&self, request: JobEventsRequest) -> Result> { let mut body = serde_json::json!({ "job_id": self.job_id }); if let Some(limit) = request.limit { diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 37890fd17..3859cb830 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -289,6 +289,10 @@ impl crate::job::JobHandle for FreshnessJob { crate::job::JobHandle::describe(&self.inner).await } + async fn timeline(&self) -> Result { + crate::job::JobHandle::timeline(&self.inner).await + } + async fn events( &self, request: crate::job::JobEventsRequest, @@ -370,6 +374,10 @@ impl crate::job::JobHandle for MaterializedViewRefreshJob { crate::job::JobHandle::describe(&self.inner).await } + async fn timeline(&self) -> Result { + crate::job::JobHandle::timeline(&self.inner).await + } + async fn events( &self, request: crate::job::JobEventsRequest, @@ -12133,6 +12141,8 @@ mod tests { .body(events.clone()) .unwrap() } + "/v1/jobs/timeline" => http::Response::builder().status(200) + .body(br#"{"traceEvents":[],"metadata":{"partial":true}}"#.to_vec()).unwrap(), other => panic!("unexpected path {other}"), } }); @@ -12149,6 +12159,7 @@ mod tests { .await .unwrap(); assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 1); + assert_eq!(job.timeline().await.unwrap()["metadata"]["partial"], true); } /// A newly submitted refresh returns a freshness-wrapped handle. Waiting