labtasker-client 2.0.0__py3-none-any.whl

labtasker/config.py

@@ -0,0 +1,193 @@
+from __future__ import annotations
+
+import os
+import tomllib
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal, TypedDict
+from urllib.parse import urlsplit, urlunsplit
+
+from labtasker.errors import ConfigError
+from labtasker.local import LocalPaths, local_paths, require_local_capabilities
+from labtasker.validation import RequestValidationError, invalid_config, validate_identifier
+
+DEFAULT_QUEUE = "default"
+CONFIG_FIELDS = {"url", "queue", "token"}
+
+
+class EndpointRecord(TypedDict):
+    mode: Literal["http", "local"]
+    url: str | None
+    socket: str | None
+    directory: str | None
+    database: str | None
+
+
+@dataclass(frozen=True, slots=True)
+class ResolvedConfig:
+    url: str | None
+    queue: str
+    token: str | None
+    local: LocalPaths | None
+
+    def public_dict(self) -> dict[str, object]:
+        return {
+            "mode": "local" if self.local is not None else "http",
+            "directory": str(self.local.directory) if self.local is not None else None,
+            "database": str(self.local.database) if self.local is not None else None,
+            "socket": str(self.local.socket) if self.local is not None else None,
+            "url": self.url,
+            "queue": self.queue,
+            "token_configured": self.token is not None,
+        }
+
+    def endpoint_dict(self) -> EndpointRecord:
+        return {
+            "mode": "local" if self.local is not None else "http",
+            "url": self.url,
+            "socket": str(self.local.socket) if self.local is not None else None,
+            "directory": str(self.local.directory) if self.local is not None else None,
+            "database": str(self.local.database) if self.local is not None else None,
+        }
+
+
+def resolve_config(
+    *,
+    url: str | None = None,
+    token: str | None = None,
+    queue: str | None = None,
+    cwd: Path | None = None,
+    environment: dict[str, str] | None = None,
+) -> ResolvedConfig:
+    working_directory = Path.cwd() if cwd is None else cwd
+    environment_values = os.environ if environment is None else environment
+    file_values = _read_config_file(working_directory)
+    constructor_values = {"url": url, "token": token, "queue": queue}
+    environment_fields = {
+        "url": environment_values.get("LABTASKER_URL"),
+        "token": environment_values.get("LABTASKER_TOKEN"),
+        "queue": environment_values.get("LABTASKER_QUEUE"),
+    }
+    defaults: dict[str, str | None] = {
+        "url": None,
+        "token": None,
+        "queue": DEFAULT_QUEUE,
+    }
+
+    effective: dict[str, str | None] = {}
+    sources: dict[str, str] = {}
+    config_path = str(working_directory / ".labtasker" / "config.toml")
+    for field in ("url", "token", "queue"):
+        candidates = [
+            (constructor_values[field], "constructor"),
+            (environment_fields[field], "environment"),
+            (file_values.get(field), config_path),
+            (defaults[field], "default"),
+        ]
+        selected = next(
+            ((candidate, source) for candidate, source in candidates if candidate is not None),
+            (None, "default"),
+        )
+        value, source = selected
+        effective[field] = value
+        sources[field] = source
+
+    effective_url = _validate_url(effective["url"], source=sources["url"])
+    effective_queue = _validate_queue(effective["queue"], source=sources["queue"])
+    effective_token = _validate_token(effective["token"], source=sources["token"])
+    local: LocalPaths | None = None
+    if effective_url is None:
+        if effective_token is not None:
+            raise invalid_config(
+                "A token requires an explicit URL; local mode does not use authentication.",
+                source=sources["token"],
+                field="token",
+            )
+        require_local_capabilities()
+        local = local_paths(working_directory)
+    return ResolvedConfig(
+        url=effective_url,
+        queue=effective_queue,
+        token=effective_token,
+        local=local,
+    )
+
+
+def _read_config_file(cwd: Path) -> dict[str, str]:
+    config_path = cwd / ".labtasker" / "config.toml"
+    legacy_path = cwd / ".labtasker" / "client.toml"
+    if not config_path.exists():
+        if legacy_path.exists():
+            raise ConfigError(
+                "legacy_config_found",
+                "A v1 client.toml was found; create the v2 flat config.toml manually.",
+                {"source": str(legacy_path)},
+            )
+        return {}
+    try:
+        raw = config_path.read_bytes()
+        parsed = tomllib.loads(raw.decode("utf-8"))
+    except (OSError, UnicodeError, tomllib.TOMLDecodeError) as error:
+        raise invalid_config(
+            "The client configuration file could not be read or parsed.",
+            source=str(config_path),
+        ) from error
+    if set(parsed) - CONFIG_FIELDS:
+        raise invalid_config(
+            f"Unknown configuration keys: {sorted(set(parsed) - CONFIG_FIELDS)!r}.",
+            source=str(config_path),
+        )
+    values: dict[str, str] = {}
+    for field, value in parsed.items():
+        if not isinstance(value, str) or not value:
+            raise invalid_config(
+                f"Configuration field '{field}' must be a non-empty string.",
+                source=str(config_path),
+                field=field,
+            )
+        values[field] = value
+    return values
+
+
+def _validate_url(value: str | None, *, source: str) -> str | None:
+    if value is None:
+        return None
+    if not isinstance(value, str) or not value:
+        raise invalid_config("URL must be a non-empty string.", source=source, field="url")
+    try:
+        parsed = urlsplit(value)
+        port = parsed.port
+    except ValueError as error:
+        raise invalid_config("URL is invalid.", source=source, field="url") from error
+    if (
+        parsed.scheme not in {"http", "https"}
+        or not parsed.netloc
+        or parsed.hostname is None
+        or parsed.username is not None
+        or parsed.password is not None
+        or parsed.query
+        or parsed.fragment
+        or (port is not None and not 1 <= port <= 65535)
+    ):
+        raise invalid_config(
+            "URL must be an absolute HTTP(S) base URL without userinfo, query or fragment.",
+            source=source,
+            field="url",
+        )
+    path = parsed.path.rstrip("/")
+    return urlunsplit((parsed.scheme, parsed.netloc, path, "", ""))
+
+
+def _validate_queue(value: str | None, *, source: str) -> str:
+    try:
+        return validate_identifier(value, field="queue")
+    except RequestValidationError as error:
+        raise invalid_config(str(error), source=source, field="queue") from error
+
+
+def _validate_token(value: str | None, *, source: str) -> str | None:
+    if value is None:
+        return None
+    if not isinstance(value, str) or not value:
+        raise invalid_config("Token must be a non-empty string.", source=source, field="token")
+    return value

labtasker/errors.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+class LabtaskerError(Exception):
+    def __init__(self, code: str, message: str, details: dict[str, Any] | None = None) -> None:
+        super().__init__(message)
+        self.code = code
+        self.message = message
+        self.details = {} if details is None else details
+
+    def as_envelope(self) -> dict[str, object]:
+        return {
+            "error": {
+                "code": self.code,
+                "message": self.message,
+                "details": self.details,
+            }
+        }
+
+
+class ConfigError(LabtaskerError):
+    def __init__(self, code: str, message: str, details: dict[str, Any]) -> None:
+        if code not in {"invalid_config", "legacy_config_found"}:
+            raise ValueError(f"Unsupported ConfigError code: {code}")
+        super().__init__(code, message, details)
+
+
+class TransportError(LabtaskerError):
+    def __init__(self, message: str, details: dict[str, Any] | None = None) -> None:
+        super().__init__("transport_error", message, details)
+
+
+class APIError(LabtaskerError):
+    def __init__(
+        self,
+        status_code: int,
+        code: str,
+        message: str,
+        details: dict[str, Any],
+    ) -> None:
+        super().__init__(code, message, details)
+        self.status_code = status_code
+
+
+class TransientError(Exception):
+    """Return the current Task to pending without charging this incident."""
+
+
+class TaskError(Exception):
+    """Fail the current execution and continue the Worker loop."""
+
+
+class FatalWorkerError(Exception):
+    """Fail the current execution and stop the Worker process."""

labtasker/execution.py

@@ -0,0 +1,380 @@
+from __future__ import annotations
+
+import os
+import threading
+import time
+import warnings
+from collections.abc import Callable
+from pathlib import Path
+from typing import Literal
+
+from labtasker.client import Client
+from labtasker.config import EndpointRecord
+from labtasker.errors import ConfigError
+from labtasker.journal import LocalRunJournal
+from labtasker.models import Task, TaskInfo
+from labtasker.types import JSONValue
+from labtasker.validation import RequestValidationError, validate_json_object
+
+CompletionReporter = Callable[[dict[str, JSONValue]], bool]
+ContextKind = Literal["python", "command"]
+
+
+class RunControl:
+    def __init__(
+        self,
+        *,
+        force_stop_timeout: float | None,
+        force_stop: Callable[[], None],
+    ) -> None:
+        self._condition = threading.Condition()
+        self._force_stop_timeout = force_stop_timeout
+        self._force_stop = force_stop
+        self._revoked_at: float | None = None
+        self._revoked_action: str | None = None
+        self._fatal_error: Exception | None = None
+        self._completed = False
+        self._executor_done = False
+        self._watchdog = threading.Thread(
+            target=self._watch_force_stop,
+            name="labtasker-force-stop",
+            daemon=True,
+        )
+        self._watchdog.start()
+
+    @property
+    def revoked(self) -> bool:
+        with self._condition:
+            return self._revoked_at is not None
+
+    @property
+    def revoked_action(self) -> str | None:
+        with self._condition:
+            return self._revoked_action
+
+    @property
+    def completed(self) -> bool:
+        with self._condition:
+            return self._completed
+
+    @property
+    def fatal_error(self) -> Exception | None:
+        with self._condition:
+            return self._fatal_error
+
+    @property
+    def active(self) -> bool:
+        with self._condition:
+            return self._revoked_at is None and not self._completed
+
+    def revoke(self, action: str) -> None:
+        with self._condition:
+            if self._completed or self._revoked_at is not None:
+                return
+            self._revoked_at = time.monotonic()
+            self._revoked_action = action
+            self._condition.notify_all()
+
+    def fail(self, error: Exception) -> None:
+        with self._condition:
+            if self._completed or self._revoked_at is not None:
+                return
+            self._fatal_error = error
+            self._revoked_at = time.monotonic()
+            self._revoked_action = error.__class__.__name__
+            self._condition.notify_all()
+
+    def complete(self) -> None:
+        with self._condition:
+            if self._revoked_at is None:
+                self._completed = True
+            self._condition.notify_all()
+
+    def executor_done(self) -> None:
+        with self._condition:
+            self._executor_done = True
+            self._condition.notify_all()
+
+    def set_force_stop_timeout(self, value: float | None) -> None:
+        with self._condition:
+            if self._completed:
+                raise RuntimeError("The current run has already completed.")
+            self._force_stop_timeout = value
+            self._condition.notify_all()
+
+    def _watch_force_stop(self) -> None:
+        while True:
+            with self._condition:
+                if self._executor_done or self._completed:
+                    return
+                if self._revoked_at is None or self._force_stop_timeout is None:
+                    self._condition.wait()
+                    continue
+                deadline = self._revoked_at + self._force_stop_timeout
+                remaining = deadline - time.monotonic()
+                if remaining > 0:
+                    self._condition.wait(remaining)
+                    continue
+            self._force_stop()
+            return
+
+
+class ExecutionContext:
+    def __init__(
+        self,
+        *,
+        info: TaskInfo,
+        kind: ContextKind,
+        journal: LocalRunJournal,
+        reporter: CompletionReporter,
+        control: RunControl | None,
+    ) -> None:
+        self.info = info
+        self.kind = kind
+        self.journal = journal
+        self.reporter = reporter
+        self.control = control
+        self._lock = threading.Lock()
+        self._finish_started = False
+        self._finished = False
+
+    @property
+    def finished(self) -> bool:
+        with self._lock:
+            return self._finished
+
+    def finish(self, result: dict[str, JSONValue]) -> None:
+        with self._lock:
+            if self._finish_started:
+                raise RuntimeError("finish() has already been called for this execution.")
+            self._finish_started = True
+        _best_effort_journal(lambda: self.journal.reporting("complete", result))
+        accepted = self.reporter(result)
+        if not accepted:
+            _best_effort_journal(self.journal.revoked)
+            raise RuntimeError("The current run was revoked before finish() could complete it.")
+        with self._lock:
+            self._finished = True
+        if self.control is not None:
+            self.control.complete()
+        _best_effort_journal(self.journal.acknowledged)
+
+
+_CONTEXT_LOCK = threading.RLock()
+_ACTIVE_CONTEXT: ExecutionContext | None = None
+_ENV_CONTEXT: ExecutionContext | None = None
+_FORK_HOOK_INSTALLED = False
+
+
+def activate_context(context: ExecutionContext) -> None:
+    global _ACTIVE_CONTEXT, _FORK_HOOK_INSTALLED
+    with _CONTEXT_LOCK:
+        if _ACTIVE_CONTEXT is not None:
+            raise RuntimeError("A Labtasker execution context is already active.")
+        if not _FORK_HOOK_INSTALLED and hasattr(os, "register_at_fork"):
+            os.register_at_fork(after_in_child=_clear_after_fork)
+            _FORK_HOOK_INSTALLED = True
+        _ACTIVE_CONTEXT = context
+
+
+def deactivate_context(context: ExecutionContext) -> None:
+    global _ACTIVE_CONTEXT
+    with _CONTEXT_LOCK:
+        if _ACTIVE_CONTEXT is context:
+            _ACTIVE_CONTEXT = None
+
+
+def active_context_present() -> bool:
+    with _CONTEXT_LOCK:
+        return _ACTIVE_CONTEXT is not None
+
+
+def task_info() -> TaskInfo:
+    context = _get_context()
+    if context is None:
+        raise RuntimeError("No active Labtasker Task execution is available.")
+    return context.info
+
+
+def finish(
+    result: dict[str, JSONValue] | None = None,
+    *,
+    skip_if_no_labtasker: bool = False,
+) -> None:
+    context = _get_context()
+    if context is None:
+        if skip_if_no_labtasker:
+            return
+        raise RuntimeError("No active Labtasker Task execution is available.")
+    normalized = validate_json_object({} if result is None else result, field="result")
+    context.finish(normalized)
+
+
+def cancellation_requested() -> bool:
+    context = _require_python_context()
+    if context.finished or context.control is None:
+        return False
+    return context.control.revoked
+
+
+def set_force_stop_timeout(seconds: float | None) -> None:
+    context = _require_python_context()
+    if context.finished or context.control is None:
+        raise RuntimeError("The current run is no longer cancellable.")
+    context.control.set_force_stop_timeout(_validate_force_stop_timeout(seconds))
+
+
+def _require_python_context() -> ExecutionContext:
+    context = _get_context()
+    if context is None or context.kind != "python":
+        raise RuntimeError("This function requires an active Python Worker execution.")
+    return context
+
+
+def _get_context() -> ExecutionContext | None:
+    with _CONTEXT_LOCK:
+        if _ACTIVE_CONTEXT is not None:
+            return _ACTIVE_CONTEXT
+    return _load_environment_context()
+
+
+def _load_environment_context() -> ExecutionContext | None:
+    global _ENV_CONTEXT
+    with _CONTEXT_LOCK:
+        if _ENV_CONTEXT is not None:
+            return _ENV_CONTEXT
+        names = {
+            "queue": "LABTASKER_QUEUE",
+            "task_id": "LABTASKER_TASK_ID",
+            "run_id": "LABTASKER_RUN_ID",
+            "route": "LABTASKER_ROUTE",
+            "run_dir": "LABTASKER_RUN_DIR",
+        }
+        values = {field: os.environ.get(name) for field, name in names.items()}
+        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):
+            return None
+        if present != set(names):
+            raise ConfigError(
+                "invalid_config",
+                "Inherited Labtasker execution context is incomplete.",
+                {"missing": sorted(set(names) - present)},
+            )
+        url = os.environ.get("LABTASKER_URL")
+        socket = os.environ.get("LABTASKER_SOCKET")
+        local_directory = os.environ.get("LABTASKER_LOCAL_DIRECTORY")
+        token = os.environ.get("LABTASKER_TOKEN")
+        http_endpoint = url is not None and socket is None and local_directory is None
+        local_endpoint = (
+            url is None and socket is not None and local_directory is not None and token is None
+        )
+        if not (http_endpoint or local_endpoint):
+            raise ConfigError(
+                "invalid_config",
+                "Inherited Labtasker execution endpoint is incomplete or ambiguous.",
+                {"expected": ("LABTASKER_URL, or LABTASKER_SOCKET and LABTASKER_LOCAL_DIRECTORY")},
+            )
+        run_dir = Path(values["run_dir"] or "")
+        if not run_dir.is_absolute():
+            raise ConfigError(
+                "invalid_config",
+                "LABTASKER_RUN_DIR must be absolute.",
+                {"field": "LABTASKER_RUN_DIR"},
+            )
+        try:
+            task = Task.model_validate_json((run_dir / "task.json").read_bytes(), strict=True)
+            run_id = values["run_id"] or ""
+            journal = LocalRunJournal.open(run_dir)
+            expected_endpoint: EndpointRecord
+            if http_endpoint:
+                expected_endpoint = {
+                    "mode": "http",
+                    "url": url,
+                    "socket": None,
+                    "directory": None,
+                    "database": None,
+                }
+                client = Client(url=url, token=token, queue=values["queue"])
+            else:
+                directory = Path(local_directory or "")
+                if not directory.is_absolute() or directory.resolve() != directory:
+                    raise ValueError("LABTASKER_LOCAL_DIRECTORY must be a canonical absolute path")
+                client = Client._from_local_directory(
+                    directory,
+                    queue=values["queue"] or "",
+                )
+                local = client.configuration.local
+                assert local is not None
+                if str(local.socket) != socket:
+                    raise ValueError("local execution socket does not match its directory")
+                expected_endpoint = client.configuration.endpoint_dict()
+            if (
+                task.id != values["task_id"]
+                or task.queue != values["queue"]
+                or journal.endpoint != expected_endpoint
+                or journal.queue != values["queue"]
+                or journal.task_id != values["task_id"]
+                or journal.run_id != run_id
+                or journal.route != values["route"]
+            ):
+                raise ValueError("execution environment does not match the local journal")
+            info = TaskInfo(**task.model_dump(), run_id=run_id, run_dir=journal.run_dir)
+        except Exception as error:
+            raise ConfigError(
+                "invalid_config",
+                "Inherited Labtasker execution context could not be loaded.",
+                {"source": str(run_dir)},
+            ) from error
+
+        def report(result: dict[str, JSONValue]) -> bool:
+            from labtasker.worker import report_complete_until_resolved
+
+            return report_complete_until_resolved(
+                client,
+                queue=values["queue"] or "",
+                task_id=values["task_id"] or "",
+                run_id=run_id,
+                result=result,
+            )
+
+        _ENV_CONTEXT = ExecutionContext(
+            info=info,
+            kind="command",
+            journal=journal,
+            reporter=report,
+            control=None,
+        )
+        return _ENV_CONTEXT
+
+
+def _best_effort_journal(operation: Callable[[], None]) -> None:
+    try:
+        operation()
+    except Exception as error:
+        warnings.warn(
+            f"Labtasker could not update the local run journal: {error}",
+            RuntimeWarning,
+            stacklevel=3,
+        )
+
+
+def _validate_force_stop_timeout(value: float | None) -> float | None:
+    if value is None:
+        return None
+    if isinstance(value, bool) or not isinstance(value, (int, float)):
+        raise RequestValidationError(
+            "force_stop_timeout must be a finite non-negative number or None"
+        )
+    normalized = float(value)
+    if not 0 <= normalized < float("inf"):
+        raise RequestValidationError(
+            "force_stop_timeout must be a finite non-negative number or None"
+        )
+    return normalized
+
+
+def _clear_after_fork() -> None:
+    global _ACTIVE_CONTEXT, _ENV_CONTEXT
+    _ACTIVE_CONTEXT = None
+    _ENV_CONTEXT = None