From d82fe8520fac608bee8efb71cde77d8aebf5dad8 Mon Sep 17 00:00:00 2001 From: Wyatt Alt Date: Tue, 1 Sep 2026 13:20:59 +0000 Subject: [PATCH] feat: pause and resume jobs from the Python SDK Cancellation is the only job control the SDK exposes, and it is terminal, so a long-running server-side job cannot be parked and picked up again. This adds pause_job and resume_job to the Rust core connection and the Python bindings (sync and async), posting to the server's /v1/jobs/pause and /v1/jobs/resume endpoints. A pause parks the job until it is resumed: its workers drain and stop. The outcome strings mirror the server's answers -- a job finalizing its results reports "committing" and cannot be parked, and a resume before the drain is confirmed reports "still_pausing"; both are retried rather than failed. Resuming re-queues the job and its workers pick their work back up from checkpoints. Local connections report the operations as unsupported, like the rest of the jobs API. --- python/python/lancedb/_lancedb.pyi | 2 + python/python/lancedb/db.py | 53 ++++++++++++++++++ python/python/lancedb/remote/db.py | 16 ++++++ python/python/tests/test_remote_db.py | 23 ++++++++ python/src/connection.rs | 24 +++++++++ rust/lancedb/src/connection.rs | 14 ++++- rust/lancedb/src/database.rs | 33 ++++++++++++ rust/lancedb/src/remote/db.rs | 77 ++++++++++++++++++++++++++- rust/lancedb/src/remote/job.rs | 22 ++++++++ 9 files changed, 262 insertions(+), 2 deletions(-) diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 2b2691139..3c8584c48 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -154,6 +154,8 @@ class Connection(object): async def list_jobs(self) -> List[JobInfo]: ... async def get_job(self, job_id: str) -> Optional[JobDescription]: ... async def cancel_job(self, job_id: str) -> bool: ... + async def pause_job(self, job_id: str) -> str: ... + async def resume_job(self, job_id: str) -> str: ... async def job_history( self, job_id: Optional[str] = None ) -> List[pa.RecordBatch]: ... diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index ecaae42f8..6f499e129 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -753,6 +753,26 @@ class DBConnection(EnforceOverrides): "cancel_job is not supported for this connection type" ) + def pause_job(self, job_id: str) -> str: + """Pause a server-side job by id. + + The job's workers drain and it stays parked until resumed. Returns + "pausing", "already_paused", or "committing" -- a job finalizing its + results cannot be parked; retry shortly. + """ + raise NotImplementedError("pause_job is not supported for this connection type") + + def resume_job(self, job_id: str) -> str: + """Resume a paused server-side job by id. + + Its workers pick their work back up from checkpoints. Returns + "resumed", "still_pausing" -- the pause's worker drain is not + confirmed yet; retry shortly -- or "not_paused". + """ + raise NotImplementedError( + "resume_job is not supported for this connection type" + ) + def job_history(self, job_id: Optional[str] = None) -> List[pa.RecordBatch]: """The lifecycle event history of a server-side job, as Arrow batches. @@ -1450,6 +1470,22 @@ class LanceDBConnection(DBConnection): """ return LOOP.run(self._conn.cancel_job(job_id)) + @override + def pause_job(self, job_id: str) -> str: + """Pause a server-side job by id. + + Returns "pausing", "already_paused", or "committing". + """ + return LOOP.run(self._conn.pause_job(job_id)) + + @override + def resume_job(self, job_id: str) -> str: + """Resume a paused server-side job by id. + + Returns "resumed", "still_pausing", or "not_paused". + """ + return LOOP.run(self._conn.resume_job(job_id)) + @override def job_history(self, job_id: Optional[str] = None) -> List[pa.RecordBatch]: """The lifecycle event history of a server-side job, as Arrow batches. @@ -2281,6 +2317,23 @@ class AsyncConnection(object): """ return await self._inner.cancel_job(job_id) + async def pause_job(self, job_id: str) -> str: + """Pause a server-side job by id. + + The job's workers drain and it stays parked until resumed. Returns + "pausing", "already_paused", or "committing" -- a job finalizing its + results cannot be parked; retry shortly. + """ + return await self._inner.pause_job(job_id) + + async def resume_job(self, job_id: str) -> str: + """Resume a paused server-side job by id. + + Its workers pick their work back up from checkpoints. Returns + "resumed", "still_pausing" -- retry shortly -- or "not_paused". + """ + return await self._inner.resume_job(job_id) + async def job_history(self, job_id: Optional[str] = None) -> List[pa.RecordBatch]: """The lifecycle event history of a server-side job, as Arrow batches. diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 27e21d200..272060865 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -776,6 +776,22 @@ class RemoteDBConnection(DBConnection): """ return LOOP.run(self._conn.cancel_job(job_id)) + @override + def pause_job(self, job_id: str) -> str: + """Pause a server-side job by id. + + Returns "pausing", "already_paused", or "committing". + """ + return LOOP.run(self._conn.pause_job(job_id)) + + @override + def resume_job(self, job_id: str) -> str: + """Resume a paused server-side job by id. + + Returns "resumed", "still_pausing", or "not_paused". + """ + return LOOP.run(self._conn.resume_job(job_id)) + @override def job_history(self, job_id: Optional[str] = None) -> List[pa.RecordBatch]: """The lifecycle event history of a server-side job, as Arrow batches. diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 01e2cc4c5..ffd1de641 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -2534,6 +2534,26 @@ def test_remote_connection_jobs_surface(): request.send_header("Content-Type", "application/json") request.end_headers() request.wfile.write(b'{"job_id": "job-1"}') + elif request.path == "/v1/jobs/pause": + if payload["job_id"] != "job-1": + request.send_response(404) + request.end_headers() + return + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(b'{"job_id": "job-1", "paused": true}') + elif request.path == "/v1/jobs/resume": + if payload["job_id"] != "job-1": + request.send_response(404) + request.end_headers() + return + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write( + b'{"job_id": "job-1", "resumed": false, "still_pausing": true}' + ) elif request.path == "/v1/jobs/query_events": assert payload["job_id"] == "job-1" request.send_response(200) @@ -2562,6 +2582,9 @@ def test_remote_connection_jobs_surface(): assert db.cancel_job("job-1") is True assert db.cancel_job("missing") is False + assert db.pause_job("job-1") == "pausing" + assert db.resume_job("job-1") == "still_pausing" + batches = db.job_history("job-1") assert len(batches) == 1 assert batches[0].num_rows == 2 diff --git a/python/src/connection.rs b/python/src/connection.rs index fc835f805..dc64fae89 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -666,6 +666,30 @@ impl Connection { }) } + pub fn pause_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult> { + let inner = self_.get_inner()?.clone(); + future_into_py(self_.py(), async move { + let status = inner.pause_job(&job_id).await.infer_error()?; + Ok(match status { + lancedb::database::PauseJobStatus::Pausing => "pausing", + lancedb::database::PauseJobStatus::AlreadyPaused => "already_paused", + lancedb::database::PauseJobStatus::Committing => "committing", + }) + }) + } + + pub fn resume_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult> { + let inner = self_.get_inner()?.clone(); + future_into_py(self_.py(), async move { + let status = inner.resume_job(&job_id).await.infer_error()?; + Ok(match status { + lancedb::database::ResumeJobStatus::Resumed => "resumed", + lancedb::database::ResumeJobStatus::StillPausing => "still_pausing", + lancedb::database::ResumeJobStatus::NotPaused => "not_paused", + }) + }) + } + #[pyo3(signature = (job_id=None))] pub fn job_history( self_: PyRef<'_, Self>, diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index 943ad51b7..9297341f2 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -24,7 +24,7 @@ use crate::data::scannable::Scannable; use crate::database::listing::ListingDatabase; use crate::database::{ CloneTableRequest, Database, DatabaseOptions, JobDescription, JobInfo, OpenTableRequest, - ReadConsistency, TableNamesRequest, + PauseJobStatus, ReadConsistency, ResumeJobStatus, TableNamesRequest, }; use crate::embeddings::{EmbeddingRegistry, MemoryRegistry}; use crate::error::{Error, Result}; @@ -590,6 +590,18 @@ impl Connection { self.internal.cancel_job(job_id.as_ref()).await } + /// Pause a server-side job by id. Its workers drain and it stays parked + /// until resumed; see [`PauseJobStatus`] for the outcomes. + pub async fn pause_job(&self, job_id: impl AsRef) -> Result { + self.internal.pause_job(job_id.as_ref()).await + } + + /// Resume a paused server-side job by id. Its workers pick their work + /// back up from checkpoints; see [`ResumeJobStatus`] for the outcomes. + pub async fn resume_job(&self, job_id: impl AsRef) -> Result { + self.internal.resume_job(job_id.as_ref()).await + } + /// The lifecycle event history of a server-side job (all jobs when /// `job_id` is `None`), as recorded Arrow batches. pub async fn job_history(&self, job_id: Option<&str>) -> Result> { diff --git a/rust/lancedb/src/database.rs b/rust/lancedb/src/database.rs index 775b0b579..5d6fda4c4 100644 --- a/rust/lancedb/src/database.rs +++ b/rust/lancedb/src/database.rs @@ -235,6 +235,29 @@ pub struct JobDescription { pub failure: Option, } +/// The server's answer to a pause request. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PauseJobStatus { + /// The pause was accepted; workers drain and the job stays parked. + Pausing, + /// The job was already paused, so a repeated pause changed nothing. + AlreadyPaused, + /// The job is finalizing its results and cannot be parked right now. + /// The commit is the short tail of a long job; retry shortly. + Committing, +} + +/// The server's answer to a resume request. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ResumeJobStatus { + /// The job re-entered the queue and will run again. + Resumed, + /// The pause's worker drain is not confirmed yet; retry shortly. + StillPausing, + /// The job was not paused, so there was nothing to resume. + NotPaused, +} + fn job_op_not_supported(what: &str) -> Result { Err(crate::error::Error::NotSupported { message: format!("{} is not supported by this database", what), @@ -331,6 +354,16 @@ pub trait Database: async fn cancel_job(&self, _job_id: &str) -> Result { job_op_not_supported("cancel_job") } + /// Pause a job by id. The job's workers drain and it stays parked until + /// resumed; see [`PauseJobStatus`] for the outcomes. + async fn pause_job(&self, _job_id: &str) -> Result { + job_op_not_supported("pause_job") + } + /// Resume a paused job by id. It re-enters the queue and its workers pick + /// their work back up from checkpoints; see [`ResumeJobStatus`]. + async fn resume_job(&self, _job_id: &str) -> Result { + job_op_not_supported("resume_job") + } /// The lifecycle event history of a job (all jobs when `job_id` is /// `None`), as recorded Arrow batches. async fn job_history(&self, _job_id: Option<&str>) -> Result> { diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 39b258a63..7a23327e1 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -26,7 +26,9 @@ use crate::database::{ use crate::error::Result; use crate::function::{FunctionRegistrationRequest, FunctionVersion}; use crate::job::Job; -use crate::remote::job::{DescribeJobResponse, RemoteJob, job_state_to_client}; +use crate::remote::job::{ + DescribeJobResponse, PauseJobResponse, RemoteJob, ResumeJobResponse, job_state_to_client, +}; use crate::remote::util::stream_as_body; use crate::table::BaseTable; @@ -684,6 +686,40 @@ impl Database for RemoteDatabase { } } + async fn pause_job(&self, job_id: &str) -> Result { + let req = self + .client + .post("/v1/jobs/pause") + .json(&serde_json::json!({ "job_id": job_id })); + let (request_id, rsp) = self.client.send(req).await?; + let rsp = self.client.check_response(&request_id, rsp).await?; + let body: PauseJobResponse = rsp.json().await.err_to_http(request_id)?; + Ok(if body.paused { + crate::database::PauseJobStatus::Pausing + } else if body.committing { + crate::database::PauseJobStatus::Committing + } else { + crate::database::PauseJobStatus::AlreadyPaused + }) + } + + async fn resume_job(&self, job_id: &str) -> Result { + let req = self + .client + .post("/v1/jobs/resume") + .json(&serde_json::json!({ "job_id": job_id })); + let (request_id, rsp) = self.client.send(req).await?; + let rsp = self.client.check_response(&request_id, rsp).await?; + let body: ResumeJobResponse = rsp.json().await.err_to_http(request_id)?; + Ok(if body.resumed { + crate::database::ResumeJobStatus::Resumed + } else if body.still_pausing { + crate::database::ResumeJobStatus::StillPausing + } else { + crate::database::ResumeJobStatus::NotPaused + }) + } + async fn job_history(&self, job_id: Option<&str>) -> Result> { let mut body = serde_json::json!({}); if let Some(job_id) = job_id { @@ -2619,6 +2655,45 @@ mod tests { assert!(!conn.cancel_job("nope").await.unwrap()); } + #[tokio::test] + async fn test_pause_and_resume_job() { + use crate::database::{PauseJobStatus, ResumeJobStatus}; + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.url().path(), "/v1/jobs/pause"); + http::Response::builder() + .status(200) + .body(r#"{"job_id": "job-1", "paused": true}"#) + .unwrap() + }); + assert_eq!( + conn.pause_job("job-1").await.unwrap(), + PauseJobStatus::Pausing + ); + + let conn = Connection::new_with_handler(|_| { + http::Response::builder() + .status(200) + .body(r#"{"job_id": "job-1", "paused": false, "committing": true}"#) + .unwrap() + }); + assert_eq!( + conn.pause_job("job-1").await.unwrap(), + PauseJobStatus::Committing + ); + + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.url().path(), "/v1/jobs/resume"); + http::Response::builder() + .status(200) + .body(r#"{"job_id": "job-1", "resumed": false, "still_pausing": true}"#) + .unwrap() + }); + assert_eq!( + conn.resume_job("job-1").await.unwrap(), + ResumeJobStatus::StillPausing + ); + } + #[tokio::test] async fn test_job_history_parses_arrow_stream() { 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 0d41dbb35..5d36d033d 100644 --- a/rust/lancedb/src/remote/job.rs +++ b/rust/lancedb/src/remote/job.rs @@ -73,6 +73,28 @@ pub(super) struct ReportedFailure { retryable: Option, } +/// Forward-compatible `/v1/jobs/pause` wire envelope. +#[derive(Deserialize)] +pub(super) struct PauseJobResponse { + /// False when the job was already paused, so a repeated pause changed + /// nothing. + #[serde(default)] + pub(super) paused: bool, + /// The job is finalizing its results and cannot be parked right now. + #[serde(default)] + pub(super) committing: bool, +} + +/// Forward-compatible `/v1/jobs/resume` wire envelope. +#[derive(Deserialize)] +pub(super) struct ResumeJobResponse { + #[serde(default)] + pub(super) resumed: bool, + /// The pause's worker drain is not confirmed yet. + #[serde(default)] + pub(super) still_pausing: bool, +} + /// Forward-compatible `/v1/jobs/describe` wire envelope. #[derive(Deserialize)] pub(super) struct DescribeJobResponse {