labtasker-client 2.4.1__tar.gz → 2.6.0__tar.gz

{labtasker_client-2.4.1 → labtasker_client-2.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: labtasker-client
-Version: 2.4.1
+Version: 2.6.0
 Summary: A small task queue for parallel model inference and evaluation
 Project-URL: Homepage, https://github.com/luocfprime/labtasker
 Project-URL: Repository, https://github.com/luocfprime/labtasker.git

{labtasker_client-2.4.1 → labtasker_client-2.6.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "labtasker-client"
-version = "2.4.1"
+version = "2.6.0"
 description = "A small task queue for parallel model inference and evaluation"
 requires-python = ">=3.10"
 license = "Apache-2.0"

{labtasker_client-2.4.1 → labtasker_client-2.6.0}/src/labtasker/__init__.py

@@ -31,6 +31,7 @@ from labtasker.execution import (
     cancellation_requested,
     finish,
     report_progress,
+    report_worker_telemetry,
     set_force_stop_timeout,
     task_info,
 )
@@ -49,7 +50,7 @@ from labtasker.models import (
 from labtasker.types import JSONValue, TaskOrderField, TaskStatus, TaskUpdate
 from labtasker.worker import loop
 
-__version__ = "2.4.1"
+__version__ = "2.6.0"
 
 __all__ = [
     "APIError",
@@ -89,6 +90,7 @@ __all__ = [
     "list_workers",
     "loop",
     "report_progress",
+    "report_worker_telemetry",
     "requeue_task",
     "set_force_stop_timeout",
     "submit_task",

{labtasker_client-2.4.1 → labtasker_client-2.6.0}/src/labtasker/cli.py

@@ -3,6 +3,9 @@ from __future__ import annotations
 import json
 import logging
 from collections.abc import Callable
+from contextvars import ContextVar
+from dataclasses import dataclass
+from pathlib import Path
 from typing import Annotated, Any, TypeVar, cast
 
 import typer
@@ -15,7 +18,6 @@ from labtasker.command_template import TemplateSyntaxError
 from labtasker.command_worker import run_command_worker
 from labtasker.config import resolve_config
 from labtasker.errors import LabtaskerError
-from labtasker.execution import report_progress as report_current_progress
 from labtasker.types import TaskOrderField, TaskStatus, TaskUpdate
 from labtasker.validation import RequestValidationError, validate_grouping, validate_json_object
 
@@ -58,6 +60,15 @@ app.add_typer(config_app, name="config")
 logger = logging.getLogger("labtasker.cli")
 
 
+@dataclass(frozen=True, slots=True)
+class CLIState:
+    labtasker_root: Path | None
+    auto_start_local_server: bool
+
+
+_CLI_STATE: ContextVar[CLIState | None] = ContextVar("labtasker_cli_state", default=None)
+
+
 def _version_callback(value: bool) -> None:
     if value:
         typer.echo(f"labtasker-client {__version__}")
@@ -75,8 +86,20 @@ def main(
             help="Show the Client package version and exit.",
         ),
     ] = False,
+    labtasker_root: Annotated[
+        Path | None,
+        typer.Option(help="Exact configuration, journal, and managed-local root."),
+    ] = None,
+    auto_start_local_server: Annotated[
+        bool,
+        typer.Option(
+            "--auto-start-local-server",
+            help="Allow this invocation to start a managed-local daemon.",
+        ),
+    ] = False,
 ) -> None:
     """Submit, inspect, and execute Labtasker v2 Tasks."""
+    _CLI_STATE.set(CLIState(labtasker_root, auto_start_local_server))
 
 
 class _SeparatedCommand(TyperCommand):
@@ -126,6 +149,10 @@ def worker_loop(
             )
         ),
     ] = None,
+    metadata: Annotated[
+        str,
+        typer.Option(help="Static Worker metadata as one strict JSON object."),
+    ] = "{}",
 ) -> None:
     """Claim matching Tasks and execute one child command for each claim.
 
@@ -146,6 +173,7 @@ def worker_loop(
         argv.pop(0)
     if not argv:
         raise typer.BadParameter("COMMAND is required after --")
+    worker_metadata = _json_object(metadata, option="--metadata")
     try:
         run_command_worker(
             argv,
@@ -154,6 +182,9 @@ def worker_loop(
             idle_timeout=idle_timeout,
             max_consecutive_failures=max_consecutive_failures,
             force_stop_timeout=force_stop_timeout,
+            metadata=worker_metadata,
+            labtasker_root=_cli_state().labtasker_root,
+            auto_start_local_server=_cli_state().auto_start_local_server,
         )
     except (TemplateSyntaxError, RequestValidationError) as error:
         raise typer.BadParameter(str(error)) from error
@@ -170,29 +201,6 @@ def worker_loop(
         raise typer.Exit(1) from error
 
 
-@app.command("progress")
-def progress_report(
-    data: Annotated[
-        str,
-        typer.Option(help="Latest progress as one strict JSON object."),
-    ],
-) -> None:
-    """Replace the current Task run's progress snapshot.
-
-    This command is available inside a command launched by ``labtasker loop``.
-    It prints whether the best-effort report was accepted; transport failures
-    and confirmed revocation return ``reported: false`` without failing the
-    command workload.
-    """
-    try:
-        progress = _json_object(data, option="--data")
-        reported = _invoke(lambda: report_current_progress(progress))
-    except RuntimeError as error:
-        typer.echo(str(error), err=True)
-        raise typer.Exit(1) from error
-    _write_json({"reported": reported})
-
-
 @task_app.command("submit")
 def task_submit(
     args: Annotated[
@@ -599,7 +607,10 @@ def queue_delete(
     name: Annotated[str, typer.Argument(help="Queue name to delete.")],
     cascade: Annotated[
         bool,
-        typer.Option(help="Also permanently delete every non-running Task in the Queue."),
+        typer.Option(
+            "--cascade",
+            help="Also permanently delete every non-running Task in the Queue.",
+        ),
     ] = False,
 ) -> None:
     """Permanently delete one Queue.
@@ -618,14 +629,30 @@ def config_show() -> None:
     .labtasker/config.toml, then built-in defaults. The token value is never
     printed.
     """
-    _write_json(_invoke(lambda: resolve_config().public_dict()))
+    state = _cli_state()
+    _write_json(
+        _invoke(
+            lambda: resolve_config(
+                labtasker_root=state.labtasker_root,
+                auto_start_local_server=state.auto_start_local_server,
+            ).public_dict()
+        )
+    )
 
 
 def _with_client(operation: Callable[[Client], T]) -> T:
-    with Client() as client:
+    state = _cli_state()
+    with Client(
+        labtasker_root=state.labtasker_root,
+        auto_start_local_server=state.auto_start_local_server,
+    ) as client:
         return operation(client)
 
 
+def _cli_state() -> CLIState:
+    return _CLI_STATE.get() or CLIState(None, False)
+
+
 def _invoke(operation: Callable[[], T]) -> T:
     try:
         return operation()

{labtasker_client-2.4.1 → labtasker_client-2.6.0}/src/labtasker/client.py

@@ -13,12 +13,7 @@ from pydantic import TypeAdapter, ValidationError
 
 from labtasker.config import ResolvedConfig, resolve_config
 from labtasker.errors import APIError, TransportError
-from labtasker.local import (
-    ensure_local_server,
-    local_paths,
-    require_local_capabilities,
-    socket_transport,
-)
+from labtasker.local import ensure_local_server, socket_transport
 from labtasker.models import (
     BulkUpdateResult,
     ClaimResponse,
@@ -53,7 +48,7 @@ from labtasker.validation import (
 
 T = TypeVar("T")
 ModelT = TypeVar("ModelT", bound=ResponseModel)
-REQUEST_TIMEOUT_SECONDS = 10.0
+REQUEST_TIMEOUT_SECONDS = 15.0
 MAX_RETRY_ATTEMPTS = 3
 RETRY_BACKOFF_SECONDS = (0.05, 0.1)
 QUEUE_LIST_ADAPTER = TypeAdapter(list[Queue])
@@ -63,21 +58,36 @@ class Client:
     def __init__(
         self,
         url: str | None = None,
+        socket: str | Path | None = None,
+        labtasker_root: str | Path | None = None,
+        auto_start_local_server: bool = False,
         token: str | None = None,
         queue: str | None = None,
     ) -> None:
-        self._initialize(resolve_config(url=url, token=token, queue=queue))
+        self._initialize(
+            resolve_config(
+                url=url,
+                socket=socket,
+                labtasker_root=labtasker_root,
+                auto_start_local_server=auto_start_local_server,
+                token=token,
+                queue=queue,
+            )
+        )
 
     @classmethod
-    def _from_local_directory(cls, directory: Path, *, queue: str) -> Client:
-        require_local_capabilities()
+    def _from_socket(cls, socket: Path, *, queue: str) -> Client:
         client = cls.__new__(cls)
         client._initialize(
             ResolvedConfig(
                 url=None,
+                socket=socket,
+                managed_local=False,
+                labtasker_root=Path("/"),
                 queue=validate_identifier(queue, field="queue"),
                 token=None,
-                local=local_paths(directory),
+                auto_start_local_server=False,
+                local=None,
             )
         )
         return client
@@ -87,22 +97,22 @@ class Client:
         headers = {}
         if self._config.token is not None:
             headers["Authorization"] = f"Bearer {self._config.token}"
-        if self._config.local is None:
-            assert self._config.url is not None
+        if self._config.url is not None:
             self._http = httpx.Client(
                 base_url=f"{self._config.url}/api/v2/",
                 headers=headers,
                 timeout=REQUEST_TIMEOUT_SECONDS,
             )
         else:
+            assert self._config.socket is not None
             self._http = httpx.Client(
                 base_url="http://labtasker/api/v2/",
-                transport=socket_transport(self._config.local),
+                transport=socket_transport(self._config.socket),
                 timeout=REQUEST_TIMEOUT_SECONDS,
             )
         self._closed = False
         self._endpoint_announced = False
-        self._local_ready = False
+        self._local_ensure_result: tuple[int | None, str | None] | None = None
         self._server_version: str | None = None
         self._warned_server_versions: set[Version] = set()
 
@@ -148,7 +158,7 @@ class Client:
             )
 
     @property
-    def configuration(self) -> ResolvedConfig:
+    def _configuration(self) -> ResolvedConfig:
         return self._config
 
     def submit_task(
@@ -538,7 +548,7 @@ class Client:
             path=f"queues/{queue_name}/tasks/claim",
             json={"route": normalized_route, "run_id": normalized_run_id},
             parser=_parse_claim,
-            retry=True,
+            recover_local_connect=False,
         )
 
     def _health(self) -> HealthResponse:
@@ -561,6 +571,7 @@ class Client:
         task_id: str,
         run_id: str,
         queue: str | None = None,
+        recover_local_connect: bool = True,
     ) -> HeartbeatResponse:
         return self._run_action(
             "heartbeat",
@@ -569,6 +580,7 @@ class Client:
             queue=queue,
             body={},
             parser=lambda response: _parse_model(response, HeartbeatResponse, {200}),
+            recover_local_connect=recover_local_connect,
         )
 
     def _complete(
@@ -607,6 +619,27 @@ class Client:
             parser=lambda response: _parse_none(response, {204}),
         )
 
+    def _report_worker_telemetry(
+        self,
+        *,
+        worker_id: str,
+        telemetry: dict[str, JSONValue],
+        queue: str | None = None,
+    ) -> None:
+        self._ensure_open()
+        queue_name = self._queue(queue)
+        normalized_telemetry = validate_json_object(telemetry, field="telemetry")
+        from labtasker.validation import validate_worker_id
+
+        normalized_worker_id = validate_worker_id(worker_id)
+        self._call(
+            operation="report_worker_telemetry",
+            method="POST",
+            path=f"queues/{queue_name}/workers/{normalized_worker_id}/telemetry",
+            json={"telemetry": normalized_telemetry},
+            parser=lambda response: _parse_none(response, {204}),
+        )
+
     def _fail(
         self,
         *,
@@ -686,6 +719,7 @@ class Client:
         queue: str | None,
         body: dict[str, object],
         parser: Callable[[httpx.Response], T],
+        recover_local_connect: bool = True,
     ) -> T:
         self._ensure_open()
         queue_name = self._queue(queue)
@@ -697,6 +731,7 @@ class Client:
             path=f"queues/{queue_name}/tasks/{normalized_task_id}/{action}",
             json={"run_id": normalized_run_id, **body},
             parser=parser,
+            recover_local_connect=recover_local_connect,
         )
 
     def _queue(self, queue: str | None) -> str:
@@ -712,9 +747,9 @@ class Client:
         json: object | None = None,
         params: dict[str, str | int] | None = None,
         retry: bool = False,
+        recover_local_connect: bool = True,
     ) -> T:
         self._ensure_open()
-        self._prepare_endpoint()
         attempts = MAX_RETRY_ATTEMPTS if retry else 1
         last_transport_error: TransportError | None = None
         local_connect_recovery_used = False
@@ -725,13 +760,14 @@ class Client:
             except httpx.RequestError as error:
                 last_transport_error = self._connection_error(operation)
                 can_recover_local_connect = (
-                    self._config.local is not None
+                    recover_local_connect
+                    and self._config.managed_local
+                    and self._config.auto_start_local_server
                     and isinstance(error, (httpx.ConnectError, httpx.ConnectTimeout))
                     and not local_connect_recovery_used
                 )
                 if can_recover_local_connect:
                     local_connect_recovery_used = True
-                    self._local_ready = False
                     self._ensure_local_available()
                     if attempt + 1 == attempts:
                         attempts += 1
@@ -739,7 +775,7 @@ class Client:
                     raise last_transport_error from error
             else:
                 self._observe_server_version(response)
-                self._announce_http_endpoint()
+                self._announce_endpoint()
                 if response.is_error:
                     try:
                         api_error = _parse_api_error(response)
@@ -773,32 +809,55 @@ class Client:
             raise AssertionError("Request loop ended without a result or error.")
         raise last_transport_error
 
+    def _repair_local_connection(self, error: TransportError) -> None:
+        """Repair an opted-in managed-local connection without replaying a request."""
+        cause = error.__cause__
+        if (
+            self._config.managed_local
+            and self._config.auto_start_local_server
+            and isinstance(cause, (httpx.ConnectError, httpx.ConnectTimeout))
+        ):
+            self._ensure_local_available()
+
     def _ensure_open(self) -> None:
         if self._closed:
             raise RuntimeError("Client is closed.")
 
     @property
     def _operation_endpoint_details(self) -> dict[str, object]:
-        if self._config.local is None:
+        if self._config.url is not None:
             return {"url": self._config.url}
         return {
-            "directory": str(self._config.local.directory),
-            "socket": str(self._config.local.socket),
+            "labtasker_root": str(self._config.labtasker_root),
+            "socket": str(self._config.socket),
         }
 
-    def _prepare_endpoint(self) -> None:
-        if self._config.local is not None and not self._local_ready:
-            self._ensure_local_available()
-
-    def _announce_http_endpoint(self) -> None:
-        if self._endpoint_announced or self._config.local is not None:
+    def _announce_endpoint(self) -> None:
+        if self._endpoint_announced:
             return
-        assert self._config.url is not None
-        transport = self._config.url.partition(":")[0]
-        print(
-            f"[labtasker] connected server=remote transport={transport} url={self._config.url}",
-            file=sys.stderr,
-        )
+        if self._config.url is not None:
+            transport = self._config.url.partition(":")[0]
+            message = (
+                f"[labtasker] connected server=remote transport={transport} url={self._config.url}"
+            )
+        elif self._config.managed_local:
+            assert self._config.local is not None and self._config.socket is not None
+            pid, version = self._local_ensure_result or (None, None)
+            message = (
+                "[labtasker] connected server=local transport=unix "
+                f"labtasker_root={self._config.labtasker_root} "
+                f"database={self._config.local.database} socket={self._config.socket}"
+            )
+            if pid is not None:
+                message += f" pid={pid}"
+            if version is not None:
+                message += f" version={version}"
+        else:
+            assert self._config.socket is not None
+            message = (
+                f"[labtasker] connected server=external transport=unix socket={self._config.socket}"
+            )
+        print(message, file=sys.stderr)
         self._endpoint_announced = True
 
     def _ensure_local_available(self) -> None:
@@ -806,16 +865,7 @@ class Client:
         if paths is None:
             return
         result = ensure_local_server(paths, emit=self._emit_local_transition)
-        pid = result.pid if result.pid is not None else "unknown"
-        version = result.server_version if result.server_version is not None else "unknown"
-        print(
-            "[labtasker] connected server=local transport=unix "
-            f"directory={paths.directory} database={paths.database} socket={paths.socket} "
-            f"pid={pid} version={version}",
-            file=sys.stderr,
-        )
-        self._endpoint_announced = True
-        self._local_ready = True
+        self._local_ensure_result = (result.pid, result.server_version)
 
     @staticmethod
     def _emit_local_transition(message: str) -> None:
@@ -823,16 +873,25 @@ class Client:
 
     def _connection_error(self, operation: str) -> TransportError:
         details: dict[str, object] = {"operation": operation}
-        if self._config.local is None:
+        if self._config.url is not None:
             details["url"] = self._config.url
+        elif not self._config.managed_local:
+            details["socket"] = str(self._config.socket)
         else:
+            assert self._config.local is not None
             details.update(
                 {
                     "state": "unhealthy",
-                    "directory": str(self._config.local.directory),
+                    "labtasker_root": str(self._config.labtasker_root),
                     "database": str(self._config.local.database),
-                    "socket": str(self._config.local.socket),
+                    "socket": str(self._config.socket),
                     "log": str(self._config.local.log),
+                    "remedies": [
+                        "rerun with --auto-start-local-server",
+                        "launch labtasker-server serve --connection socket --daemon "
+                        f"--labtasker-root {self._config.labtasker_root}",
+                        "configure LABTASKER_URL or LABTASKER_SOCKET",
+                    ],
                 }
             )
         return TransportError("The Labtasker Server could not be reached.", details)

{labtasker_client-2.4.1 → labtasker_client-2.6.0}/src/labtasker/command_worker.py

@@ -27,17 +27,16 @@ from labtasker.models import ClaimResponse
 from labtasker.observations import ObservationReporter
 from labtasker.tee import configure_worker_logger
 from labtasker.types import JSONValue
-from labtasker.validation import validate_identifier
+from labtasker.validation import validate_identifier, validate_json_object
 from labtasker.worker import (
-    POLL_INTERVAL_SECONDS,
     Heartbeat,
     _best_effort_unclaim,
     _ExecutionResult,
     _FailureGuard,
     _finish_journal,
-    _generate_run_id,
     _guard_worker_topology,
     _journal_best_effort,
+    _next_claim,
     _preflight,
     _report_until_resolved,
     _safe_diagnostic_text,
@@ -57,36 +56,41 @@ def run_command_worker(
     idle_timeout: float = 300.0,
     force_stop_timeout: float | None = None,
     max_consecutive_failures: int = 5,
+    metadata: dict[str, JSONValue] | None = None,
+    labtasker_root: Path | None = None,
+    auto_start_local_server: bool = False,
 ) -> None:
     guard = _FailureGuard(max_consecutive_failures)
     templates = compile_argv(argv)
     normalized_route = validate_identifier(route, field="route")
     normalized_idle_timeout = _validate_idle_timeout(idle_timeout)
     normalized_force_stop_timeout = _validate_force_stop_timeout(force_stop_timeout)
+    normalized_metadata = validate_json_object(
+        {} if metadata is None else metadata, field="metadata"
+    )
     _guard_command_worker_platform()
     _guard_worker_topology()
     configure_worker_logger()
-    with Client(queue=queue) as client:
-        queue_name = client.configuration.queue
+    with Client(
+        queue=queue,
+        labtasker_root=labtasker_root,
+        auto_start_local_server=auto_start_local_server,
+    ) as client:
+        queue_name = client._configuration.queue
         _preflight(client, queue_name)
-        with ObservationReporter(client.configuration, normalized_route) as observer:
-            idle_deadline: float | None = None
+        with ObservationReporter(
+            client._configuration, normalized_route, normalized_metadata
+        ) as observer:
             while True:
-                claim = client._claim(
+                claim = _next_claim(
+                    client,
                     route=normalized_route,
-                    run_id=_generate_run_id(),
                     queue=queue_name,
+                    idle_timeout=normalized_idle_timeout,
                 )
                 if claim is None:
-                    now = time.monotonic()
-                    if idle_deadline is None:
-                        idle_deadline = now + normalized_idle_timeout
-                    if now >= idle_deadline:
-                        logger.info("Worker idle timeout reached; stopping normally.")
-                        return
-                    time.sleep(min(POLL_INTERVAL_SECONDS, idle_deadline - now))
-                    continue
-                idle_deadline = None
+                    logger.info("Worker idle timeout reached; stopping normally.")
+                    return
                 observer.activity(claim.task.id)
                 logger.info(
                     "Claimed Task %s as run %s (attempt %d, route %s).",
@@ -102,6 +106,7 @@ def run_command_worker(
                     queue=queue_name,
                     route=normalized_route,
                     force_stop_timeout=normalized_force_stop_timeout,
+                    worker_id=observer.id,
                 )
 
                 guard.observe(result, claim.task.id)
@@ -124,13 +129,15 @@ def _run_command_claim(
     queue: str,
     route: str,
     force_stop_timeout: float | None,
+    worker_id: str,
 ) -> _ExecutionResult:
     try:
         journal = LocalRunJournal.create(
             claim=claim,
-            endpoint=client.configuration.endpoint_dict(),
+            endpoint=client._configuration.endpoint_dict(),
             queue=queue,
             route=route,
+            labtasker_root=client._configuration.labtasker_root,
         )
     except Exception:
         _best_effort_unclaim(client, claim, queue)
@@ -153,7 +160,7 @@ def _run_command_claim(
             return _report_command_failure(
                 client, journal, claim, queue, "TaskBindingError", str(error), control=control
             )
-        environment = _command_environment(client, claim, journal, queue, route)
+        environment = _command_environment(client, claim, journal, queue, route, worker_id)
         try:
             if _interactive_terminal():
                 process = _run_pty(
@@ -489,6 +496,7 @@ def _command_environment(
     journal: LocalRunJournal,
     queue: str,
     route: str,
+    worker_id: str,
 ) -> dict[str, str]:
     environment = dict(os.environ)
     environment.update(
@@ -498,20 +506,22 @@ def _command_environment(
             "LABTASKER_RUN_ID": claim.run_id,
             "LABTASKER_ROUTE": route,
             "LABTASKER_RUN_DIR": str(journal.run_dir),
+            "LABTASKER_WORKER_ID": worker_id,
         }
     )
-    configuration = client.configuration
-    if configuration.local is None:
+    configuration = client._configuration
+    if configuration.url is not None:
         assert configuration.url is not None
         environment["LABTASKER_URL"] = configuration.url
         environment.pop("LABTASKER_SOCKET", None)
-        environment.pop("LABTASKER_LOCAL_DIRECTORY", None)
     else:
-        environment["LABTASKER_SOCKET"] = str(configuration.local.socket)
-        environment["LABTASKER_LOCAL_DIRECTORY"] = str(configuration.local.directory)
+        assert configuration.socket is not None
+        environment["LABTASKER_SOCKET"] = str(configuration.socket)
         environment.pop("LABTASKER_URL", None)
+    environment.pop("LABTASKER_ROOT", None)
+    environment.pop("LABTASKER_LOCAL_DIRECTORY", None)
     token = configuration.token
-    if token is None or configuration.local is not None:
+    if token is None or configuration.url is None:
         environment.pop("LABTASKER_TOKEN", None)
     else:
         environment["LABTASKER_TOKEN"] = token