From 3f1a94c8d15dde0b3a71f2013b60a73844f815bf Mon Sep 17 00:00:00 2001 From: Wyatt Alt Date: Sat, 18 Jul 2026 17:14:44 -0700 Subject: [PATCH] python: Job/AsyncJob rebound to the platform jobs API (ENT-1956) JobHandle becomes Job (AsyncJobHandle -> AsyncJob), per the reference naming convention (Job the reference vs JobInfo the snapshot), and the implementation moves off the legacy inflight-listing poll onto the platform jobs API: - The reference holds the submission (manifest) id and lazily resolves the platform job id (one /v1/jobs/list call with the manifest-id filter), tolerating async dispatch with a pending grace window. - status/progress/wait read /v1/jobs/describe: terminal states are first-class (DONE / FAILED / CANCELLED) instead of inferred from leaving the inflight listing, progress comes from the owner-written payload, and a failed job raises JobFailedError with the server error. A job that never registers raises instead of hanging. - cancel() drives /v1/jobs/cancel (with a short resolve retry), which the server now propagates to running workers. Connection surfaces gain the three passthroughs (sync + async); every Job-returning API (refresh_column, MV refresh/wait, load_columns) hands out the new type. job_history()/errors() stay on their existing routes (the per-row error store is outside the platform jobs API). Co-Authored-By: Claude Fable 5 --- python/python/lancedb/__init__.py | 8 +- python/python/lancedb/db.py | 51 +++++- python/python/lancedb/remote/table.py | 10 +- python/python/lancedb/table.py | 8 +- python/python/lancedb/udf.py | 226 ++++++++++++++++--------- python/python/tests/test_job_handle.py | 178 ++++++++++++++----- rust/lancedb/src/remote/db.rs | 10 +- 7 files changed, 338 insertions(+), 153 deletions(-) diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 511f91443..b62f598ee 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -23,10 +23,10 @@ from .udf import ( udf, table_udf, Udf, - JobHandle, + Job, JobFailedError, MaterializedView, - AsyncJobHandle, + AsyncJob, AsyncMaterializedView, ) from .lineage import Lineage, Node, Edge, FunctionRef @@ -505,10 +505,10 @@ __all__ = [ "udf", "table_udf", "Udf", - "JobHandle", + "Job", "JobFailedError", "MaterializedView", - "AsyncJobHandle", + "AsyncJob", "AsyncMaterializedView", "Lineage", "Node", diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 72c141dea..1b6883fa1 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -697,13 +697,13 @@ class DBConnection(EnforceOverrides): raise def job(self, job_id: str): - """A `JobHandle` for reconnecting to an inflight job by id -- e.g. an + """A `Job` for reconnecting to an inflight job by id -- e.g. an id you stored, or one returned from the SQL / REST surface. Submit methods (`refresh_column`, `MaterializedView.refresh`) already return a handle directly, so you do not need this to wait on a fresh submission.""" - from .udf import JobHandle + from .udf import Job - return JobHandle(self, job_id) + return Job(self, job_id) def lineage( self, @@ -735,7 +735,7 @@ class DBConnection(EnforceOverrides): ) -> str: """Internal: submit a materialized-view refresh, return the job id. The public surface is ``MaterializedView.refresh()`` (which returns a - `JobHandle`); this stays private so refresh is only reached through the + `Job`); this stays private so refresh is only reached through the handle. ``full=True`` forces a full rebuild (recompute and replace every row) @@ -802,6 +802,24 @@ class DBConnection(EnforceOverrides): """ return LOOP.run(self._conn.cancel_job(job_id)) + def describe_platform_job(self, platform_job_id: str): + """Describe a platform job (POST /v1/jobs/describe): registry-backed + lifecycle state plus the owner-written status payload. None when the + registry has no such job.""" + return LOOP.run(self._conn.describe_platform_job(platform_job_id)) + + def resolve_platform_job_id( + self, manifest_job_id: str, table: "str | None" = None + ): + """Resolve a submission (manifest) job id to its platform job id. + None until the job has registered (dispatch is async).""" + return LOOP.run(self._conn.resolve_platform_job_id(manifest_job_id, table)) + + def cancel_platform_job(self, platform_job_id: str) -> None: + """Cancel a platform job (POST /v1/jobs/cancel). Idempotent on + already-terminal jobs.""" + return LOOP.run(self._conn.cancel_platform_job(platform_job_id)) + def job_history(self, job_id: "str | None" = None): """Durable history of completed server-side jobs (SHOW JOB HISTORY). @@ -2131,13 +2149,13 @@ class AsyncConnection(object): return AsyncMaterializedView(self, name, job_id=job_id) def job(self, job_id: str): - """An `AsyncJobHandle` for reconnecting to an inflight job by id (a + """An `AsyncJob` for reconnecting to an inflight job by id (a stored id, or one from the SQL / REST surface). Submit methods already return a handle, so this is only needed to re-attach to an existing job.""" - from .udf import AsyncJobHandle + from .udf import AsyncJob - return AsyncJobHandle(self, job_id) + return AsyncJob(self, job_id) async def lineage( self, @@ -2164,7 +2182,7 @@ class AsyncConnection(object): max_workers: Optional[int] = None, ) -> str: """Internal: submit a refresh, return the job id. The public surface is - ``AsyncMaterializedView.refresh()`` (returns an `AsyncJobHandle`). + ``AsyncMaterializedView.refresh()`` (returns an `AsyncJob`). ``full=True`` forces a full rebuild (recompute and replace every row) instead of the default incremental refresh. @@ -2219,6 +2237,23 @@ class AsyncConnection(object): """ return await self._inner.cancel_job(job_id) + async def describe_platform_job(self, platform_job_id: str): + """Describe a platform job: registry-backed lifecycle state plus the + owner-written status payload. None when the registry has no such + job.""" + return await self._inner.describe_platform_job(platform_job_id) + + async def resolve_platform_job_id( + self, manifest_job_id: str, table: "str | None" = None + ): + """Resolve a submission (manifest) job id to its platform job id. + None until the job has registered (dispatch is async).""" + return await self._inner.resolve_platform_job_id(manifest_job_id, table) + + async def cancel_platform_job(self, platform_job_id: str) -> None: + """Cancel a platform job. Idempotent on already-terminal jobs.""" + return await self._inner.cancel_platform_job(platform_job_id) + async def job_history(self, job_id: "str | None" = None): """Durable history of completed server-side jobs (SHOW JOB HISTORY). diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index ba084e16a..bfee90497 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -20,7 +20,7 @@ from typing import ( ) if TYPE_CHECKING: - from ..udf import JobHandle + from ..udf import Job import warnings from lancedb import __version__ @@ -940,12 +940,12 @@ class RemoteTable(Table): max_workers: Optional[int] = None, batch_size: Optional[int] = None, priority: Optional[str] = None, - ) -> "JobHandle": + ) -> "Job": """Trigger recompute of computed columns (REFRESH COLUMN). The expression is resolved server-side from each column's stored binding; columns bound to the same struct-returning function - refresh together. Returns a `JobHandle` to wait on, poll, or cancel + refresh together. Returns a `Job` to wait on, poll, or cancel (``tbl.refresh_column("c").wait()``). Server-backed feature (LanceDB Enterprise / Cloud). @@ -954,7 +954,7 @@ class RemoteTable(Table): the function carries. `priority` is a Kueue tier (training | interactive | backfill). """ - from ..udf import JobHandle + from ..udf import Job if isinstance(columns, str): columns = [columns] @@ -968,7 +968,7 @@ class RemoteTable(Table): priority=priority, ) ) - return JobHandle(self._job_conn(), job_id) + return Job(self._job_conn(), job_id) def lineage(self, column=None, *, direction=None, depth=None): """Derived-compute lineage of this table, or one of its columns: diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index a58a0767f..a37f11454 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -3918,12 +3918,12 @@ class LanceTable(Table): max_workers: Optional[int] = None, batch_size: Optional[int] = None, priority: Optional[str] = None, - ) -> "JobHandle": + ) -> "Job": """Trigger recompute of computed columns (REFRESH COLUMN). The expression is resolved server-side from each column's stored binding; columns bound to the same struct-returning function - refresh together. Returns a `JobHandle` to wait on, poll, or cancel + refresh together. Returns a `Job` to wait on, poll, or cancel (``tbl.refresh_column("col").wait()``) -- mirrors `MaterializedView.refresh()`. Server-backed feature (LanceDB Enterprise / Cloud). @@ -3933,7 +3933,7 @@ class LanceTable(Table): the function carries. `priority` is a Kueue tier (training | interactive | backfill). """ - from .udf import JobHandle + from .udf import Job if isinstance(columns, str): columns = [columns] @@ -3947,7 +3947,7 @@ class LanceTable(Table): priority=priority, ) ) - return JobHandle(self._conn, job_id, table=self.name) + return Job(self._conn, job_id, table=self.name) def alter_columns( self, *alterations: Iterable[Dict[str, str]] diff --git a/python/python/lancedb/udf.py b/python/python/lancedb/udf.py index c4f91f893..cb86f8654 100644 --- a/python/python/lancedb/udf.py +++ b/python/python/lancedb/udf.py @@ -19,7 +19,7 @@ Register and use them through the existing connection/table API: db.create_function(embed) # CREATE FUNCTION (once) tbl = db.open_table("docs") tbl.add_columns(computed={"vec": embed("text")}) # bind embed(text) -> vec - tbl.refresh_column("vec").wait() # materialize (returns a JobHandle) + tbl.refresh_column("vec").wait() # materialize (returns a Job) view = db.create_materialized_view("chunks", tbl, ["id", chunk_fn]) `embed("text")` applies the registered function to the `text` column and yields @@ -41,6 +41,7 @@ import inspect import re import sys import textwrap +import json import time import typing @@ -510,12 +511,12 @@ class MaterializedView: A no-op when the view was created with no data.""" if self.job_id is None: return "finished" - return JobHandle(self.conn, self.job_id, table=self.name).wait( + return Job(self.conn, self.job_id, table=self.name).wait( timeout=timeout, poll=poll ) - def refresh(self, full: bool = False) -> "JobHandle": - """Refresh the materialized view; returns a `JobHandle` to wait on, + def refresh(self, full: bool = False) -> "Job": + """Refresh the materialized view; returns a `Job` to wait on, poll, or cancel (``view.refresh().wait()``). ``full=True`` forces a full rebuild (recompute and replace every row) @@ -523,7 +524,7 @@ class MaterializedView: the view's indexes -- they are reindexed by the distributed indexer. """ job_id = self.conn._refresh_materialized_view(self.name, full=full) - return JobHandle(self.conn, job_id, table=self.name) + return Job(self.conn, job_id, table=self.name) def explain_refresh(self, full: bool = False): """Plan a refresh without running it (EXPLAIN REFRESH).""" @@ -570,7 +571,7 @@ _PROGRESS = re.compile(r"(\d+)/(\d+)") class JobFailedError(RuntimeError): - """Raised by ``JobHandle.wait()`` when the server reports the job ``failed``. + """Raised by ``Job.wait()`` when the server reports the job ``failed``. Carries the server-side error so a doomed backfill (e.g. a multi-column ``REFRESH COLUMN`` of a scalar UDF) surfaces its real cause promptly, @@ -583,71 +584,110 @@ class JobFailedError(RuntimeError): super().__init__(f"job {job_id} failed: {error or 'unknown error'}") -class JobHandle: - """A reference to an inflight server-side job, with polling helpers.""" +class Job: + """A reference to a server-side job, backed by the platform jobs API. - #: How long an unseen job is treated as still materializing (submission - #: -> agent cycle -> manifest write is async). + Holds the submission (manifest) id and resolves the platform job id + lazily; ``status``/``progress``/``wait`` read the registry-backed + describe endpoint, so terminal states and errors are first-class. + """ + + #: How long an unresolved job is treated as still materializing + #: (submission -> dispatch -> registry record is async). GRACE_SECONDS = 20.0 + #: Platform lifecycle state -> the user-facing vocabulary. + _STATES = { + "IN_PROGRESS": "running", + "DONE": "finished", + "FAILED": "failed", + "CANCELLED": "cancelled", + } + def __init__(self, conn, job_id: str, table: "str | None" = None): self.conn = conn + #: The submission (manifest) id the launching call handed out. self.id = job_id - #: The job's table, when known (refresh_column / MV refresh). Lets the - #: server resolve this job with an O(1) single-node read; without it the - #: lookup scans the database's active jobs (still correct). + #: The job's table, when known -- narrows platform-id resolution. self.table = table + self._platform_id: "str | None" = None self._created = time.monotonic() - self._seen = False - def _job(self): - # Poll by id (one job), not list_jobs (every active job): the server - # matches the submission/manifest id and reads just this table's node. - return self.conn.get_job(self.id, self.table) + def _resolve(self) -> "str | None": + if self._platform_id is None: + self._platform_id = self.conn.resolve_platform_job_id(self.id, self.table) + return self._platform_id + + def _describe(self): + platform_id = self._resolve() + if platform_id is None: + return None + return self.conn.describe_platform_job(platform_id) + + @staticmethod + def _payload(described) -> dict: + # Older records carry the status-store URI string instead of a + # payload; anything non-dict means "no structured status". + try: + payload = json.loads(described.status_json) + except (TypeError, ValueError): + return {} + return payload if isinstance(payload, dict) else {} def status(self) -> str: - """pending / running / cancelling / stale, or 'finished' once the - job has left the inflight listing.""" - job = self._job() - if job is not None: - self._seen = True - return job.state - if not self._seen and time.monotonic() - self._created < self.GRACE_SECONDS: + """pending / running / finished / failed / cancelled (or unknown + when the job never appeared in the registry).""" + described = self._describe() + if described is not None: + return self._STATES.get(described.job_state, described.job_state) + if time.monotonic() - self._created < self.GRACE_SECONDS: return "pending" - return "finished" + return "unknown" def progress(self) -> "tuple[int, int] | None": - """(units_done, units_total) while running, else None.""" - job = self._job() - if job is not None and job.units_total is not None: - return job.units_done or 0, job.units_total + """(units_done, units_total) once workers have published progress.""" + described = self._describe() + if described is None: + return None + payload = self._payload(described) + if payload.get("units_total") is not None: + return payload.get("units_done") or 0, payload["units_total"] return None def wait(self, timeout: float = 3600.0, poll: float = 2.0) -> str: deadline = time.monotonic() + timeout while time.monotonic() < deadline: - state = self.status() - if state in ("finished", "stale"): - return state - if state == "failed": - # Terminal failure -- surface the server error now, don't block - # until `timeout`. `finalize` wrote it to the job's status node. - job = self._job() - raise JobFailedError(self.id, job.error if job is not None else None) - if state == "pending": + described = self._describe() + if described is None: + if time.monotonic() - self._created > self.GRACE_SECONDS: + raise JobFailedError( + self.id, + "job did not appear in the job registry within the " + "grace period", + ) time.sleep(min(poll, 0.5)) continue - job = self._job() - if job is not None and job.committed: - return "finished" + state = self._STATES.get(described.job_state, described.job_state) + if state == "finished": + return state + if state == "cancelled": + return state + if state == "failed": + raise JobFailedError(self.id, self._payload(described).get("error")) time.sleep(poll) raise TimeoutError(f"job {self.id} still {self.status()} after {timeout}s") def cancel(self) -> None: - # Cancel by the canonical manifest id (what cancel matches), found - # via the submission prefix; fall back to the raw id. - job = self._job() - self.conn.cancel_job(job.job_id if job is not None else self.id) + """Request cancellation. Workers drain cooperatively; poll ``status`` + for the terminal ``cancelled``.""" + deadline = time.monotonic() + 5.0 + while (platform_id := self._resolve()) is None: + if time.monotonic() > deadline: + raise RuntimeError( + f"job {self.id} has not registered yet; retry cancel shortly" + ) + time.sleep(0.5) + self.conn.cancel_platform_job(platform_id) class AsyncMaterializedView: @@ -664,19 +704,19 @@ class AsyncMaterializedView: A no-op when the view was created with no data.""" if self.job_id is None: return "finished" - return await AsyncJobHandle(self.conn, self.job_id, table=self.name).wait( + return await AsyncJob(self.conn, self.job_id, table=self.name).wait( timeout=timeout, poll=poll ) - async def refresh(self, full: bool = False) -> "AsyncJobHandle": - """Refresh the materialized view; returns an `AsyncJobHandle` to wait + async def refresh(self, full: bool = False) -> "AsyncJob": + """Refresh the materialized view; returns an `AsyncJob` to wait on, poll, or cancel. ``full=True`` forces a full rebuild instead of an incremental refresh (indexes are preserved and reindexed by the distributed indexer). """ job_id = await self.conn._refresh_materialized_view(self.name, full=full) - return AsyncJobHandle(self.conn, job_id, table=self.name) + return AsyncJob(self.conn, job_id, table=self.name) async def explain_refresh(self, full: bool = False): return await self.conn.explain_refresh_materialized_view(self.name, full=full) @@ -694,60 +734,80 @@ class AsyncMaterializedView: ) -class AsyncJobHandle: - """Async reference to an inflight server-side job, with polling helpers.""" +class AsyncJob: + """Async reference to a server-side job, backed by the platform jobs API. + + Same contract as `Job` with awaitable methods. + """ GRACE_SECONDS = 20.0 + _STATES = Job._STATES def __init__(self, conn, job_id: str, table: "str | None" = None): self.conn = conn self.id = job_id - #: See JobHandle.table -- enables an O(1) by-id lookup when known. self.table = table + self._platform_id: "str | None" = None self._created = time.monotonic() - self._seen = False - async def _job(self): - # Poll by id, not list_jobs (see JobHandle._job). - return await self.conn.get_job(self.id, self.table) + async def _resolve(self) -> "str | None": + if self._platform_id is None: + self._platform_id = await self.conn.resolve_platform_job_id( + self.id, self.table + ) + return self._platform_id + + async def _describe(self): + platform_id = await self._resolve() + if platform_id is None: + return None + return await self.conn.describe_platform_job(platform_id) async def status(self) -> str: - job = await self._job() - if job is not None: - self._seen = True - return job.state - if not self._seen and time.monotonic() - self._created < self.GRACE_SECONDS: + described = await self._describe() + if described is not None: + return self._STATES.get(described.job_state, described.job_state) + if time.monotonic() - self._created < self.GRACE_SECONDS: return "pending" - return "finished" + return "unknown" async def progress(self) -> "tuple[int, int] | None": - job = await self._job() - if job is not None and job.units_total is not None: - return job.units_done or 0, job.units_total + described = await self._describe() + if described is None: + return None + payload = Job._payload(described) + if payload.get("units_total") is not None: + return payload.get("units_done") or 0, payload["units_total"] return None async def wait(self, timeout: float = 3600.0, poll: float = 2.0) -> str: deadline = time.monotonic() + timeout while time.monotonic() < deadline: - state = await self.status() - if state in ("finished", "stale"): - return state - if state == "failed": - # Terminal failure -- surface the server error now, don't block - # until `timeout`. `finalize` wrote it to the job's status node. - job = await self._job() - raise JobFailedError(self.id, job.error if job is not None else None) - if state == "pending": + described = await self._describe() + if described is None: + if time.monotonic() - self._created > self.GRACE_SECONDS: + raise JobFailedError( + self.id, + "job did not appear in the job registry within the " + "grace period", + ) await asyncio.sleep(min(poll, 0.5)) continue - job = await self._job() - if job is not None and job.committed: - return "finished" + state = self._STATES.get(described.job_state, described.job_state) + if state in ("finished", "cancelled"): + return state + if state == "failed": + raise JobFailedError(self.id, Job._payload(described).get("error")) await asyncio.sleep(poll) - raise TimeoutError( - f"job {self.id} still {await self.status()} after {timeout}s" - ) + raise TimeoutError(f"job {self.id} still {await self.status()} after {timeout}s") async def cancel(self) -> None: - job = await self._job() - await self.conn.cancel_job(job.job_id if job is not None else self.id) + deadline = time.monotonic() + 5.0 + while (platform_id := await self._resolve()) is None: + if time.monotonic() > deadline: + raise RuntimeError( + f"job {self.id} has not registered yet; retry cancel shortly" + ) + await asyncio.sleep(0.5) + await self.conn.cancel_platform_job(platform_id) + diff --git a/python/python/tests/test_job_handle.py b/python/python/tests/test_job_handle.py index 1cb194a1d..10ed12079 100644 --- a/python/python/tests/test_job_handle.py +++ b/python/python/tests/test_job_handle.py @@ -1,92 +1,182 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The LanceDB Authors -"""JobHandle.wait() terminal-state handling. +"""Job / AsyncJob against the platform jobs API. -Regression coverage for the cluster backfill-failure hang: the server reports a -doomed job as ``state="failed"`` within seconds, but ``wait()`` used to ignore -``failed`` and block until its (default 3600s) timeout. These tests pin that a -``failed`` job raises ``JobFailedError`` promptly, carrying the server error. +The reference resolves its submission (manifest) id to a platform job id, +then polls describe for registry-backed state: terminal states are +first-class (DONE / FAILED / CANCELLED), progress comes from the +owner-written status payload, and a failed job raises ``JobFailedError`` +promptly with the server error. """ import asyncio +import json import time import pytest -from lancedb.udf import JobHandle, AsyncJobHandle, JobFailedError +from lancedb.udf import Job, AsyncJob, JobFailedError -class FakeJobInfo: - """Mirror of the pyo3 builtins.JobInfo fields wait()/status() read.""" +class FakeDescription: + """Mirror of the pyo3 PlatformJobDescription fields the Job reads.""" - def __init__(self, state, error=None, committed=False, units_total=None): - self.state = state - self.error = error - self.committed = committed - self.units_total = units_total - self.units_done = None - self.job_id = "job-1" + def __init__(self, job_state, status=None): + self.job_id = "plat-1" + self.job_type = "indexer" + self.job_subtype = "udf" + self.job_state = job_state + self.creation_ms = 0 + self.status_json = json.dumps(status if status is not None else {}) class FakeConn: - """get_job() walks a scripted list of JobInfo (or None) snapshots, holding - the last one once exhausted, so wait() polls a deterministic timeline.""" + """Scripted timeline: resolve returns None until `resolve_after` calls, + then the platform id; describe walks a list of descriptions (holding the + last once exhausted).""" - def __init__(self, snapshots): - self._snaps = list(snapshots) - self.calls = 0 + def __init__(self, descriptions, resolve_after=0): + self._descs = list(descriptions) + self._resolve_after = resolve_after + self.resolve_calls = 0 + self.describe_calls = 0 + self.cancelled = [] - def get_job(self, job_id, table=None): - snap = self._snaps[min(self.calls, len(self._snaps) - 1)] - self.calls += 1 + def resolve_platform_job_id(self, manifest_job_id, table=None): + self.resolve_calls += 1 + if self.resolve_calls <= self._resolve_after: + return None + return "plat-1" + + def describe_platform_job(self, platform_job_id): + assert platform_job_id == "plat-1" + snap = self._descs[min(self.describe_calls, len(self._descs) - 1)] + self.describe_calls += 1 return snap + def cancel_platform_job(self, platform_job_id): + self.cancelled.append(platform_job_id) + class AsyncFakeConn(FakeConn): - async def get_job(self, job_id, table=None): - return FakeConn.get_job(self, job_id, table) + async def resolve_platform_job_id(self, manifest_job_id, table=None): + return FakeConn.resolve_platform_job_id(self, manifest_job_id, table) + + async def describe_platform_job(self, platform_job_id): + return FakeConn.describe_platform_job(self, platform_job_id) + + async def cancel_platform_job(self, platform_job_id): + return FakeConn.cancel_platform_job(self, platform_job_id) + + +def test_status_maps_platform_states(): + for wire, want in [ + ("IN_PROGRESS", "running"), + ("DONE", "finished"), + ("FAILED", "failed"), + ("CANCELLED", "cancelled"), + ]: + job = Job(FakeConn([FakeDescription(wire)]), "job-1", table="t") + assert job.status() == want + + +def test_status_pending_before_resolution(): + job = Job(FakeConn([], resolve_after=10_000), "job-1", table="t") + assert job.status() == "pending" + + +def test_progress_from_status_payload(): + conn = FakeConn( + [ + FakeDescription( + "IN_PROGRESS", + status={"units_done": 3, "units_total": 8, "rows_committed": 100}, + ) + ] + ) + job = Job(conn, "job-1", table="t") + assert job.progress() == (3, 8) + + +def test_progress_none_for_uri_only_status(): + # Older records carry the status-store URI string, not a payload. + desc = FakeDescription("IN_PROGRESS") + desc.status_json = json.dumps("s3://bucket/job/job_status") + job = Job(FakeConn([desc]), "job-1", table="t") + assert job.progress() is None def test_wait_raises_on_failed_promptly(): - # pending -> failed: wait() must raise the server error, not TimeoutError. conn = FakeConn( - [None, FakeJobInfo("failed", error="multi-column backfill needs a STRUCT")] + [ + FakeDescription("IN_PROGRESS"), + FakeDescription( + "FAILED", status={"error": "multi-column backfill needs a STRUCT"} + ), + ] ) - jh = JobHandle(conn, "job-1", table="t") + job = Job(conn, "job-1", table="t") t0 = time.monotonic() with pytest.raises(JobFailedError) as exc: - jh.wait(timeout=30, poll=0.01) + job.wait(timeout=30, poll=0.01) assert time.monotonic() - t0 < 5 # prompt, nowhere near the 30s timeout assert "STRUCT" in str(exc.value) assert exc.value.error == "multi-column backfill needs a STRUCT" assert exc.value.job_id == "job-1" -def test_wait_returns_finished_on_success(): - # running -> finished (job left the inflight listing) returns normally. - conn = FakeConn([FakeJobInfo("running", units_total=2), None]) - jh = JobHandle(conn, "job-1", table="t") - jh._seen = True # already observed, so a None now means "finished" not grace - assert jh.wait(timeout=30, poll=0.01) == "finished" +def test_wait_returns_finished_on_done(): + conn = FakeConn([FakeDescription("IN_PROGRESS"), FakeDescription("DONE")]) + job = Job(conn, "job-1", table="t") + assert job.wait(timeout=30, poll=0.01) == "finished" -def test_wait_returns_finished_on_committed(): - # A committed job that is still listed resolves to finished. - conn = FakeConn([FakeJobInfo("running", committed=True, units_total=2)]) - jh = JobHandle(conn, "job-1", table="t") - jh._seen = True - assert jh.wait(timeout=30, poll=0.01) == "finished" +def test_wait_returns_cancelled(): + conn = FakeConn([FakeDescription("CANCELLED")]) + job = Job(conn, "job-1", table="t") + assert job.wait(timeout=30, poll=0.01) == "cancelled" + + +def test_wait_raises_when_job_never_registers(): + # An unresolved job past the grace window is a lost submission, not an + # eternal "pending" hang. + conn = FakeConn([], resolve_after=10_000) + job = Job(conn, "job-1", table="t") + job.GRACE_SECONDS = 0.05 + job._created = time.monotonic() - 1.0 + with pytest.raises(JobFailedError) as exc: + job.wait(timeout=5, poll=0.01) + assert "registry" in str(exc.value) + + +def test_cancel_resolves_then_cancels(): + conn = FakeConn([FakeDescription("IN_PROGRESS")], resolve_after=1) + job = Job(conn, "job-1", table="t") + job.cancel() + assert conn.cancelled == ["plat-1"] def test_async_wait_raises_on_failed_promptly(): - conn = AsyncFakeConn([None, FakeJobInfo("failed", error="boom")]) - jh = AsyncJobHandle(conn, "job-1", table="t") + conn = AsyncFakeConn( + [FakeDescription("FAILED", status={"error": "boom"})], + ) + job = AsyncJob(conn, "job-1", table="t") async def run(): t0 = time.monotonic() with pytest.raises(JobFailedError) as exc: - await jh.wait(timeout=30, poll=0.01) + await job.wait(timeout=30, poll=0.01) assert time.monotonic() - t0 < 5 assert exc.value.error == "boom" asyncio.run(run()) + + +def test_async_wait_returns_finished(): + conn = AsyncFakeConn([FakeDescription("IN_PROGRESS"), FakeDescription("DONE")]) + job = AsyncJob(conn, "job-1", table="t") + + async def run(): + assert await job.wait(timeout=30, poll=0.01) == "finished" + + asyncio.run(run()) diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index f548fb4ab..e92d34ce3 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -21,8 +21,8 @@ use crate::Error; use crate::database::{ CloneTableRequest, CreateFunctionRequest, CreateMaterializedViewRequest, CreateTableMode, CreateTableRequest, Database, DatabaseOptions, FunctionInfo, JobErrorInfo, JobHistoryInfo, - JobInfo, MaterializedViewInfo, MvRefreshPlan, OpenTableRequest, ReadConsistency, - RefreshMaterializedViewRequest, TableLineageRequest, TableNamesRequest, + JobInfo, MaterializedViewInfo, MvRefreshPlan, OpenTableRequest, PlatformJobDescription, + ReadConsistency, RefreshMaterializedViewRequest, TableLineageRequest, TableNamesRequest, }; use crate::error::Result; use crate::remote::util::stream_as_body; @@ -158,7 +158,7 @@ struct RemoteJobEntry { error: Option, } -#[derive(Deserialize)] +#[derive(serde::Deserialize)] struct RemoteDescribePlatformJobResponse { job_id: String, job_type: String, @@ -171,13 +171,13 @@ struct RemoteDescribePlatformJobResponse { status: serde_json::Value, } -#[derive(Deserialize)] +#[derive(serde::Deserialize)] struct RemoteListPlatformJobsResponse { #[serde(default)] jobs: Vec, } -#[derive(Deserialize)] +#[derive(serde::Deserialize)] struct RemotePlatformJobRow { job_id: String, }