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