labtasker-client 2.4.0__tar.gz → 2.5.0__tar.gz

{labtasker_client-2.4.0 → labtasker_client-2.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: labtasker-client
-Version: 2.4.0
+Version: 2.5.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
@@ -8,10 +8,12 @@ Author-email: lcf <luocfprime@gmail.com>
 License-Expression: Apache-2.0
 License-File: LICENSE
 Requires-Python: >=3.10
-Requires-Dist: httpx[socks]<1,>=0.25.2
+Requires-Dist: httpx<1,>=0.25.2
 Requires-Dist: packaging>=21.3
 Requires-Dist: pydantic<3,>=2.12; python_version >= '3.14'
 Requires-Dist: pydantic<3,>=2.5.3; python_version < '3.13'
 Requires-Dist: pydantic<3,>=2.9; python_version >= '3.13' and python_version < '3.14'
 Requires-Dist: tomli>=2.0.1; python_version < '3.11'
 Requires-Dist: typer<1,>=0.19
+Provides-Extra: socks
+Requires-Dist: socksio==1.*; extra == 'socks'

{labtasker_client-2.4.0 → labtasker_client-2.5.0}/pyproject.toml

@@ -4,14 +4,14 @@ build-backend = "hatchling.build"
 
 [project]
 name = "labtasker-client"
-version = "2.4.0"
+version = "2.5.0"
 description = "A small task queue for parallel model inference and evaluation"
 requires-python = ">=3.10"
 license = "Apache-2.0"
 license-files = ["LICENSE"]
 authors = [{ name = "lcf", email = "luocfprime@gmail.com" }]
 dependencies = [
-  "httpx[socks]>=0.25.2,<1",
+  "httpx>=0.25.2,<1",
   "packaging>=21.3",
   "pydantic>=2.5.3,<3; python_version < '3.13'",
   "pydantic>=2.9,<3; python_version >= '3.13' and python_version < '3.14'",
@@ -20,6 +20,9 @@ dependencies = [
   "typer>=0.19,<1",
 ]
 
+[project.optional-dependencies]
+socks = ["socksio==1.*"]
+
 [project.urls]
 Homepage = "https://github.com/luocfprime/labtasker"
 Repository = "https://github.com/luocfprime/labtasker.git"

{labtasker_client-2.4.0 → labtasker_client-2.5.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.0"
+__version__ = "2.5.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.0 → labtasker_client-2.5.0}/src/labtasker/cli.py

@@ -16,6 +16,7 @@ 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.execution import report_worker_telemetry as report_current_worker_telemetry
 from labtasker.types import TaskOrderField, TaskStatus, TaskUpdate
 from labtasker.validation import RequestValidationError, validate_grouping, validate_json_object
 
@@ -126,6 +127,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 +151,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 +160,7 @@ def worker_loop(
             idle_timeout=idle_timeout,
             max_consecutive_failures=max_consecutive_failures,
             force_stop_timeout=force_stop_timeout,
+            metadata=worker_metadata,
         )
     except (TemplateSyntaxError, RequestValidationError) as error:
         raise typer.BadParameter(str(error)) from error
@@ -465,6 +472,28 @@ def worker_count(
     _write_json({"count": result} if isinstance(result, int) else result)
 
 
+@worker_app.command("telemetry")
+def worker_telemetry_report(
+    data: Annotated[
+        str,
+        typer.Option(help="Latest Worker telemetry as one strict JSON object."),
+    ],
+) -> None:
+    """Replace telemetry for the current Worker invocation.
+
+    This command is available inside a command launched by ``labtasker loop``.
+    It performs one best-effort synchronous report and prints whether the Server
+    accepted it.
+    """
+    try:
+        telemetry = _json_object(data, option="--data")
+        reported = _invoke(lambda: report_current_worker_telemetry(telemetry))
+    except RuntimeError as error:
+        typer.echo(str(error), err=True)
+        raise typer.Exit(1) from error
+    _write_json({"reported": reported})
+
+
 def _count_options(
     group_by: list[str] | None, allowed: set[str], limit: int | None, cursor: str | None
 ) -> dict[str, Any]:

{labtasker_client-2.4.0 → labtasker_client-2.5.0}/src/labtasker/client.py

@@ -607,6 +607,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,
         *,

{labtasker_client-2.4.0 → labtasker_client-2.5.0}/src/labtasker/command_worker.py

@@ -27,7 +27,7 @@ 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,
@@ -57,19 +57,25 @@ 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,
 ) -> 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
         _preflight(client, queue_name)
-        with ObservationReporter(client.configuration, normalized_route) as observer:
+        with ObservationReporter(
+            client.configuration, normalized_route, normalized_metadata
+        ) as observer:
             idle_deadline: float | None = None
             while True:
                 claim = client._claim(
@@ -102,6 +108,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,6 +131,7 @@ def _run_command_claim(
     queue: str,
     route: str,
     force_stop_timeout: float | None,
+    worker_id: str,
 ) -> _ExecutionResult:
     try:
         journal = LocalRunJournal.create(
@@ -153,7 +161,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 +497,7 @@ def _command_environment(
     journal: LocalRunJournal,
     queue: str,
     route: str,
+    worker_id: str,
 ) -> dict[str, str]:
     environment = dict(os.environ)
     environment.update(
@@ -498,6 +507,7 @@ 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

{labtasker_client-2.4.0 → labtasker_client-2.5.0}/src/labtasker/execution.py

@@ -18,6 +18,7 @@ from labtasker.validation import RequestValidationError, validate_json_object
 
 CompletionReporter = Callable[[dict[str, JSONValue]], bool]
 ProgressReporter = Callable[[dict[str, JSONValue]], bool]
+WorkerTelemetryReporter = Callable[[dict[str, JSONValue]], bool]
 ContextKind = Literal["python", "command"]
 
 
@@ -130,6 +131,7 @@ class ExecutionContext:
         reporter: CompletionReporter,
         control: RunControl | None,
         progress_reporter: ProgressReporter | None = None,
+        worker_telemetry_reporter: WorkerTelemetryReporter | None = None,
     ) -> None:
         self.info = info
         self.kind = kind
@@ -137,6 +139,7 @@ class ExecutionContext:
         self.reporter = reporter
         self.control = control
         self.progress_reporter = progress_reporter
+        self.worker_telemetry_reporter = worker_telemetry_reporter
         self._lock = threading.Lock()
         self._finish_started = False
         self._finished = False
@@ -175,6 +178,12 @@ class ExecutionContext:
             raise RuntimeError("Progress reporting is unavailable for this execution.")
         return reporter(progress)
 
+    def report_worker_telemetry(self, telemetry: dict[str, JSONValue]) -> bool:
+        reporter = self.worker_telemetry_reporter
+        if reporter is None:
+            raise RuntimeError("Worker telemetry reporting is unavailable for this execution.")
+        return reporter(telemetry)
+
 
 _CONTEXT_LOCK = threading.RLock()
 _ACTIVE_CONTEXT: ExecutionContext | None = None
@@ -240,6 +249,20 @@ def report_progress(
     return context.report_progress(normalized)
 
 
+def report_worker_telemetry(
+    telemetry: dict[str, JSONValue],
+    *,
+    skip_if_no_labtasker: bool = False,
+) -> bool:
+    context = _get_context()
+    if context is None:
+        if skip_if_no_labtasker:
+            return False
+        raise RuntimeError("No active Labtasker Task execution is available.")
+    normalized = validate_json_object(telemetry, field="telemetry")
+    return context.report_worker_telemetry(normalized)
+
+
 def cancellation_requested() -> bool:
     context = _require_python_context()
     if context.finished or context.control is None:
@@ -281,6 +304,7 @@ def _load_environment_context() -> ExecutionContext | None:
             "run_dir": "LABTASKER_RUN_DIR",
         }
         values = {field: os.environ.get(name) for field, name in names.items()}
+        worker_id = os.environ.get("LABTASKER_WORKER_ID")
         execution_fields = {"task_id", "run_id", "route", "run_dir"}
         present = {field for field, value in values.items() if value is not None}
         if not (present & execution_fields):
@@ -379,6 +403,31 @@ def _load_environment_context() -> ExecutionContext | None:
                 progress=progress,
             )
 
+        worker_telemetry_reporter: WorkerTelemetryReporter | None = None
+        if worker_id is not None:
+            from labtasker.validation import validate_worker_id
+
+            try:
+                normalized_worker_id = validate_worker_id(worker_id)
+            except RequestValidationError as error:
+                raise ConfigError(
+                    "invalid_config",
+                    "Inherited Labtasker Worker context is invalid.",
+                    {"field": "LABTASKER_WORKER_ID"},
+                ) from error
+
+            def report_telemetry(telemetry: dict[str, JSONValue]) -> bool:
+                from labtasker.worker import report_worker_telemetry_once
+
+                return report_worker_telemetry_once(
+                    client,
+                    queue=values["queue"] or "",
+                    worker_id=normalized_worker_id,
+                    telemetry=telemetry,
+                )
+
+            worker_telemetry_reporter = report_telemetry
+
         _ENV_CONTEXT = ExecutionContext(
             info=info,
             kind="command",
@@ -386,6 +435,7 @@ def _load_environment_context() -> ExecutionContext | None:
             reporter=report,
             control=None,
             progress_reporter=progress_reporter,
+            worker_telemetry_reporter=worker_telemetry_reporter,
         )
         return _ENV_CONTEXT
 

{labtasker_client-2.4.0 → labtasker_client-2.5.0}/src/labtasker/models.py

@@ -4,7 +4,7 @@ from datetime import datetime, timezone
 from pathlib import Path
 from typing import Literal
 
-from pydantic import BaseModel, ConfigDict, field_validator, model_validator
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
 from labtasker.types import JSONValue, TaskStatus
 from labtasker.validation import (
@@ -290,6 +290,9 @@ class WorkerObservation(ResponseModel):
     route: str
     status: Literal["idle", "busy"]
     task_id: str | None
+    metadata: dict[str, JSONValue] = Field(default_factory=dict)
+    telemetry: dict[str, JSONValue] | None = None
+    telemetry_updated_at: datetime | None = None
     last_seen_at: datetime
     expires_at: datetime
 
@@ -310,10 +313,25 @@ class WorkerObservation(ResponseModel):
     def validate_task(cls, value: str | None) -> str | None:
         return None if value is None else validate_task_id(value)
 
-    @field_validator("last_seen_at", "expires_at")
+    @field_validator("metadata", "telemetry")
     @classmethod
-    def validate_time(cls, value: datetime) -> datetime:
-        return _utc_datetime(value)
+    def validate_worker_objects(
+        cls, value: dict[str, JSONValue] | None, info: object
+    ) -> dict[str, JSONValue] | None:
+        if value is None:
+            return None
+        return validate_json_object(value, field=getattr(info, "field_name", "worker"))
+
+    @field_validator("telemetry_updated_at", "last_seen_at", "expires_at")
+    @classmethod
+    def validate_time(cls, value: datetime | None) -> datetime | None:
+        return None if value is None else _utc_datetime(value)
+
+    @model_validator(mode="after")
+    def validate_telemetry_fields(self) -> WorkerObservation:
+        if (self.telemetry is None) != (self.telemetry_updated_at is None):
+            raise ValueError("telemetry and telemetry_updated_at must be both null or both present")
+        return self
 
 
 class WorkerPage(ResponseModel):

{labtasker_client-2.4.0 → labtasker_client-2.5.0}/src/labtasker/observations.py

@@ -7,6 +7,7 @@ import os
 import secrets
 import threading
 import time
+from copy import deepcopy
 from types import TracebackType
 
 import httpx
@@ -14,6 +15,8 @@ import httpx
 from labtasker.client import REQUEST_TIMEOUT_SECONDS
 from labtasker.config import ResolvedConfig
 from labtasker.local import socket_transport
+from labtasker.types import JSONValue
+from labtasker.validation import validate_json_object
 
 REPORT_INTERVAL_SECONDS = 60.0
 SHUTDOWN_WAIT_SECONDS = 1.0
@@ -39,10 +42,18 @@ def _make_http_client(configuration: ResolvedConfig) -> httpx.Client:
 
 
 class ObservationReporter:
-    def __init__(self, configuration: ResolvedConfig, route: str) -> None:
+    def __init__(
+        self,
+        configuration: ResolvedConfig,
+        route: str,
+        metadata: dict[str, JSONValue] | None = None,
+    ) -> None:
         self.id = f"w_{secrets.token_urlsafe(9)}"
         self._configuration = configuration
         self._route = route
+        self._metadata = deepcopy(
+            validate_json_object({} if metadata is None else metadata, field="metadata")
+        )
         self._pid = os.getpid()
         self._lock = threading.Lock()
         self._wake = threading.Event()
@@ -139,6 +150,7 @@ class ObservationReporter:
                             "route": self._route,
                             "status": "idle" if task_id is None else "busy",
                             "task_id": task_id,
+                            "metadata": self._metadata,
                         },
                     )
                     if response.status_code != 204:

{labtasker_client-2.4.0 → labtasker_client-2.5.0}/src/labtasker/worker.py

@@ -36,7 +36,7 @@ from labtasker.models import ClaimResponse, TaskInfo
 from labtasker.observations import ObservationReporter
 from labtasker.tee import WorkerTee, configure_worker_logger
 from labtasker.types import JSONValue
-from labtasker.validation import RequestValidationError, validate_identifier
+from labtasker.validation import RequestValidationError, validate_identifier, validate_json_object
 
 P = ParamSpec("P")
 R = TypeVar("R")
@@ -133,11 +133,15 @@ def loop(
     idle_timeout: float = 300.0,
     force_stop_timeout: float | None = None,
     max_consecutive_failures: int = 5,
+    metadata: dict[str, JSONValue] | None = None,
 ) -> Callable[[Callable[P, R]], Callable[P, None]]:
     _FailureGuard(max_consecutive_failures)
     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"
+    )
 
     def decorate(function: Callable[P, R]) -> Callable[P, None]:
         binding = compile_binding(function)
@@ -157,6 +161,7 @@ def loop(
                 idle_timeout=normalized_idle_timeout,
                 force_stop_timeout=normalized_force_stop_timeout,
                 max_consecutive_failures=max_consecutive_failures,
+                metadata=normalized_metadata,
             )
 
         return run
@@ -174,6 +179,7 @@ def _run_python_worker(
     idle_timeout: float,
     force_stop_timeout: float | None,
     max_consecutive_failures: int = 5,
+    metadata: dict[str, JSONValue] | None = None,
 ) -> None:
     guard = _FailureGuard(max_consecutive_failures)
     _guard_worker_topology()
@@ -181,7 +187,7 @@ def _run_python_worker(
         configure_worker_logger()
         queue_name = client.configuration.queue
         _preflight(client, queue_name)
-        with ObservationReporter(client.configuration, route) as observer:
+        with ObservationReporter(client.configuration, route, metadata) as observer:
             idle_deadline: float | None = None
             while True:
                 claim = client._claim(route=route, run_id=_generate_run_id(), queue=queue_name)
@@ -213,6 +219,7 @@ def _run_python_worker(
                     queue=queue_name,
                     route=route,
                     force_stop_timeout=force_stop_timeout,
+                    worker_id=observer.id,
                 )
 
                 guard.observe(result, claim.task.id)
@@ -230,6 +237,7 @@ def _run_python_claim(
     queue: str,
     route: str,
     force_stop_timeout: float | None,
+    worker_id: str,
 ) -> _ExecutionResult:
     try:
         journal = LocalRunJournal.create(
@@ -271,6 +279,14 @@ def _run_python_claim(
             control=control,
         )
 
+    def worker_telemetry_reporter(telemetry: dict[str, JSONValue]) -> bool:
+        return report_worker_telemetry_once(
+            client,
+            queue=queue,
+            worker_id=worker_id,
+            telemetry=telemetry,
+        )
+
     info = TaskInfo(
         **claim.task.model_dump(),
         run_id=claim.run_id,
@@ -283,6 +299,7 @@ def _run_python_claim(
         reporter=report_complete,
         control=control,
         progress_reporter=progress_reporter,
+        worker_telemetry_reporter=worker_telemetry_reporter,
     )
     heartbeat = Heartbeat(
         client,
@@ -403,6 +420,28 @@ def report_progress_once(
     return True
 
 
+def report_worker_telemetry_once(
+    client: Client,
+    *,
+    queue: str,
+    worker_id: str,
+    telemetry: dict[str, JSONValue],
+) -> bool:
+    try:
+        client._report_worker_telemetry(
+            worker_id=worker_id,
+            telemetry=telemetry,
+            queue=queue,
+        )
+    except TransportError as error:
+        logger.warning("Worker telemetry transport error; continuing Task: %s", error.message)
+        return False
+    except LabtaskerError as error:
+        logger.warning("Worker telemetry report failed; continuing Task: %s", error.message)
+        return False
+    return True
+
+
 def _report_complete(
     client: Client,
     journal: LocalRunJournal,