labtasker-client 2.0.0__py3-none-any.whl

labtasker/journal.py

@@ -0,0 +1,314 @@
+from __future__ import annotations
+
+import json
+import os
+import tempfile
+import threading
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Literal, TypedDict, cast
+
+from labtasker.config import EndpointRecord
+from labtasker.models import ClaimResponse
+from labtasker.types import JSONValue
+from labtasker.validation import (
+    validate_identifier,
+    validate_int64,
+    validate_run_id,
+    validate_task_id,
+)
+
+JournalPhase = Literal["running", "reporting", "acknowledged", "revoked"]
+TerminalAction = Literal["complete", "fail", "unclaim"]
+LOCAL_GITIGNORE = "*\n!.gitignore\n"
+
+
+class RunRecord(TypedDict):
+    schema_version: int
+    endpoint: EndpointRecord
+    queue: str
+    task_id: str
+    run_id: str
+    route: str
+    attempt: int
+    started_at: str
+    finished_at: str | None
+    phase: JournalPhase
+    terminal_action: TerminalAction | None
+    acknowledged_at: str | None
+
+
+class LocalRunJournal:
+    def __init__(self, run_dir: Path, record: RunRecord) -> None:
+        self.run_dir = run_dir
+        self._record = record
+        self._lock = threading.Lock()
+
+    @classmethod
+    def create(
+        cls,
+        *,
+        claim: ClaimResponse,
+        endpoint: EndpointRecord,
+        queue: str,
+        route: str,
+        cwd: Path | None = None,
+    ) -> LocalRunJournal:
+        started_at = claim.task.started_at
+        if started_at is None:
+            raise ValueError("claimed Task is missing started_at")
+        root = (Path.cwd() if cwd is None else cwd).resolve()
+        task_group = f"{_task_slug(claim.task.name)}__{claim.task.id}"
+        run_name = (
+            f"{started_at.astimezone(UTC):%Y%m%dT%H%M%SZ}"
+            f"__attempt-{claim.task.attempt}__{claim.run_id}"
+        )
+        labtasker_dir = root / ".labtasker"
+        labtasker_dir.mkdir(parents=True, exist_ok=True)
+        _ensure_local_gitignore(labtasker_dir)
+        run_dir = labtasker_dir / "runs" / queue / task_group / run_name
+        run_dir.mkdir(parents=True, exist_ok=False)
+        record: RunRecord = {
+            "schema_version": 1,
+            "endpoint": endpoint,
+            "queue": queue,
+            "task_id": claim.task.id,
+            "run_id": claim.run_id,
+            "route": route,
+            "attempt": claim.task.attempt,
+            "started_at": _timestamp(started_at),
+            "finished_at": None,
+            "phase": "running",
+            "terminal_action": None,
+            "acknowledged_at": None,
+        }
+        journal = cls(run_dir.resolve(), record)
+        _atomic_json(journal.task_path, claim.task.model_dump(mode="json"))
+        _atomic_json(journal.run_path, record)
+        journal.log_path.touch(exist_ok=False)
+        return journal
+
+    @classmethod
+    def open(cls, run_dir: Path) -> LocalRunJournal:
+        resolved = run_dir.resolve()
+        parsed = json.loads((resolved / "run.json").read_text(encoding="utf-8"))
+        return cls(resolved, _validate_record(parsed))
+
+    @property
+    def endpoint(self) -> EndpointRecord:
+        return self._record["endpoint"]
+
+    @property
+    def queue(self) -> str:
+        return self._record["queue"]
+
+    @property
+    def task_id(self) -> str:
+        return self._record["task_id"]
+
+    @property
+    def run_id(self) -> str:
+        return self._record["run_id"]
+
+    @property
+    def route(self) -> str:
+        return self._record["route"]
+
+    @property
+    def task_path(self) -> Path:
+        return self.run_dir / "task.json"
+
+    @property
+    def run_path(self) -> Path:
+        return self.run_dir / "run.json"
+
+    @property
+    def result_path(self) -> Path:
+        return self.run_dir / "result.json"
+
+    @property
+    def error_path(self) -> Path:
+        return self.run_dir / "error.json"
+
+    @property
+    def log_path(self) -> Path:
+        return self.run_dir / "run.log"
+
+    @property
+    def phase(self) -> JournalPhase:
+        return self._record["phase"]
+
+    @property
+    def terminal_action(self) -> TerminalAction | None:
+        return self._record["terminal_action"]
+
+    def read_result(self) -> dict[str, JSONValue]:
+        parsed = json.loads(self.result_path.read_text(encoding="utf-8"))
+        if not isinstance(parsed, dict):
+            raise ValueError("result.json must be an object")
+        return cast(dict[str, JSONValue], parsed)
+
+    def reporting(
+        self,
+        action: TerminalAction,
+        payload: dict[str, JSONValue] | None = None,
+    ) -> None:
+        with self._lock:
+            if action == "complete":
+                if payload is None:
+                    raise ValueError("complete journal entry requires a result")
+                _atomic_json_once(self.result_path, payload)
+            elif action == "fail":
+                if payload is None:
+                    raise ValueError("fail journal entry requires an error")
+                _atomic_json_once(self.error_path, payload)
+            elif payload is not None:
+                raise ValueError("unclaim journal entry cannot have a payload")
+            self._record["phase"] = "reporting"
+            self._record["terminal_action"] = action
+            self._record["finished_at"] = None
+            self._record["acknowledged_at"] = None
+            _atomic_json(self.run_path, self._record)
+
+    def acknowledged(self) -> None:
+        now = _timestamp(datetime.now(UTC))
+        with self._lock:
+            self._record["phase"] = "acknowledged"
+            self._record["finished_at"] = now
+            self._record["acknowledged_at"] = now
+            _atomic_json(self.run_path, self._record)
+
+    def revoked(self) -> None:
+        with self._lock:
+            self._record["phase"] = "revoked"
+            self._record["finished_at"] = _timestamp(datetime.now(UTC))
+            self._record["acknowledged_at"] = None
+            _atomic_json(self.run_path, self._record)
+
+
+def _ensure_local_gitignore(labtasker_dir: Path) -> None:
+    try:
+        with (labtasker_dir / ".gitignore").open("x", encoding="utf-8", newline="\n") as stream:
+            stream.write(LOCAL_GITIGNORE)
+    except FileExistsError:
+        pass
+
+
+def _task_slug(name: str | None) -> str:
+    source = name if name else "unnamed"
+    characters: list[str] = []
+    for character in source:
+        if character.isalnum():
+            characters.append(character)
+        elif characters and characters[-1] != "-":
+            characters.append("-")
+    slug = "".join(characters).strip("-") or "unnamed"
+    encoded = slug.encode("utf-8")
+    if len(encoded) <= 80:
+        return slug
+    prefix = encoded[:80]
+    while True:
+        try:
+            slug = prefix.decode("utf-8").rstrip("-")
+            break
+        except UnicodeDecodeError:
+            prefix = prefix[:-1]
+    return slug or "unnamed"
+
+
+def _timestamp(value: datetime) -> str:
+    return value.astimezone(UTC).isoformat().replace("+00:00", "Z")
+
+
+def _atomic_json(path: Path, value: object) -> None:
+    descriptor, temporary_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
+    temporary_path = Path(temporary_name)
+    try:
+        with os.fdopen(descriptor, "w", encoding="utf-8", newline="\n") as stream:
+            json.dump(value, stream, ensure_ascii=False, allow_nan=False, indent=2)
+            stream.write("\n")
+            stream.flush()
+            os.fsync(stream.fileno())
+        temporary_path.replace(path)
+    except BaseException:
+        temporary_path.unlink(missing_ok=True)
+        raise
+
+
+def _atomic_json_once(path: Path, value: object) -> None:
+    descriptor, temporary_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
+    temporary_path = Path(temporary_name)
+    try:
+        with os.fdopen(descriptor, "w", encoding="utf-8", newline="\n") as stream:
+            json.dump(value, stream, ensure_ascii=False, allow_nan=False, indent=2)
+            stream.write("\n")
+            stream.flush()
+            os.fsync(stream.fileno())
+        try:
+            os.link(temporary_path, path)
+        except FileExistsError:
+            existing = json.loads(path.read_text(encoding="utf-8"))
+            if existing != value:
+                raise ValueError(
+                    f"{path.name} already contains a different terminal payload"
+                ) from None
+    finally:
+        temporary_path.unlink(missing_ok=True)
+
+
+def _validate_record(value: object) -> RunRecord:
+    fields = {
+        "schema_version",
+        "endpoint",
+        "queue",
+        "task_id",
+        "run_id",
+        "route",
+        "attempt",
+        "started_at",
+        "finished_at",
+        "phase",
+        "terminal_action",
+        "acknowledged_at",
+    }
+    if not isinstance(value, dict) or set(value) != fields:
+        raise ValueError("run.json does not match journal schema version 1")
+    if value["schema_version"] != 1:
+        raise ValueError("run.json uses an unsupported journal schema")
+    _validate_endpoint(value["endpoint"])
+    validate_identifier(value["queue"], field="queue")
+    validate_task_id(value["task_id"])
+    validate_run_id(value["run_id"])
+    validate_identifier(value["route"], field="route")
+    validate_int64(value["attempt"], field="attempt", positive=True)
+    for field in ("started_at", "finished_at", "acknowledged_at"):
+        if value[field] is not None and not isinstance(value[field], str):
+            raise ValueError(f"run.json {field} must be a timestamp string or null")
+    if value["phase"] not in {"running", "reporting", "acknowledged", "revoked"}:
+        raise ValueError("run.json phase is invalid")
+    if value["terminal_action"] not in {None, "complete", "fail", "unclaim"}:
+        raise ValueError("run.json terminal_action is invalid")
+    return cast(RunRecord, value)
+
+
+def _validate_endpoint(value: object) -> EndpointRecord:
+    fields = {"mode", "url", "socket", "directory", "database"}
+    if not isinstance(value, dict) or set(value) != fields:
+        raise ValueError("run.json endpoint is invalid")
+    mode = value["mode"]
+    if mode == "http":
+        if not isinstance(value["url"], str) or not value["url"]:
+            raise ValueError("run.json HTTP endpoint URL is invalid")
+        if any(value[field] is not None for field in ("socket", "directory", "database")):
+            raise ValueError("run.json HTTP endpoint contains local paths")
+    elif mode == "local":
+        if value["url"] is not None:
+            raise ValueError("run.json local endpoint contains a URL")
+        if any(
+            not isinstance(value[field], str) or not value[field]
+            for field in ("socket", "directory", "database")
+        ):
+            raise ValueError("run.json local endpoint paths are invalid")
+    else:
+        raise ValueError("run.json endpoint mode is invalid")
+    return cast(EndpointRecord, value)

labtasker/local.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+import hashlib
+import importlib.util
+import json
+import os
+import socket
+import subprocess
+import sys
+from collections.abc import Callable
+from dataclasses import dataclass
+from pathlib import Path
+
+import httpx
+
+from labtasker.errors import ConfigError, TransportError
+
+COORDINATOR_TIMEOUT_SECONDS = 35.0
+
+
+@dataclass(frozen=True, slots=True)
+class LocalPaths:
+    directory: Path
+    database: Path
+    log: Path
+    runtime_directory: Path
+    socket: Path
+
+
+@dataclass(frozen=True, slots=True)
+class LocalEnsureResult:
+    started: bool
+    pid: int | None
+    server_version: str | None
+
+
+def require_local_capabilities() -> None:
+    if os.name != "posix" or not hasattr(socket, "AF_UNIX"):
+        raise ConfigError(
+            "invalid_config",
+            "Local mode requires POSIX Unix-domain sockets; configure a URL.",
+            {"source": "default", "field": "url"},
+        )
+
+
+def local_paths(directory: Path | None = None) -> LocalPaths:
+    canonical = (Path.cwd() if directory is None else directory).resolve()
+    digest = hashlib.sha256(os.fsencode(canonical)).hexdigest()
+    runtime_directory = Path("/tmp") / f"labtasker-{os.geteuid()}"
+    local_directory = canonical / ".labtasker"
+    return LocalPaths(
+        directory=canonical,
+        database=local_directory / "server.db",
+        log=local_directory / "server.log",
+        runtime_directory=runtime_directory,
+        socket=runtime_directory / f"{digest}.sock",
+    )
+
+
+def ensure_local_server(paths: LocalPaths, *, emit: Callable[[str], None]) -> LocalEnsureResult:
+    require_local_capabilities()
+    if socket_health(paths):
+        return LocalEnsureResult(started=False, pid=None, server_version=None)
+    if importlib.util.find_spec("labtasker_server") is None:
+        raise ConfigError(
+            "invalid_config",
+            "Local mode requires labtasker-server; install labtasker or configure a URL.",
+            {"source": "default", "field": "url"},
+        )
+
+    emit(f"requesting local daemon ensure directory={paths.directory} socket={paths.socket}")
+    try:
+        result = subprocess.run(
+            [
+                sys.executable,
+                "-m",
+                "labtasker_server",
+                "_ensure-daemon",
+                "--directory",
+                str(paths.directory),
+            ],
+            cwd=paths.directory,
+            stdin=subprocess.DEVNULL,
+            stdout=subprocess.PIPE,
+            stderr=None,
+            text=True,
+            check=False,
+            timeout=COORDINATOR_TIMEOUT_SECONDS,
+        )
+    except subprocess.TimeoutExpired as error:
+        raise _local_transport_error(
+            paths,
+            state="starting",
+            message="The local Server coordinator did not finish within 35 seconds.",
+        ) from error
+    except OSError as error:
+        raise _local_transport_error(
+            paths,
+            state="stopped",
+            message="The local Server coordinator could not be started.",
+        ) from error
+
+    payload = _parse_coordinator_result(result.stdout)
+    if result.returncode != 0 or payload is None or payload.get("ok") is not True:
+        state = payload.get("state") if payload is not None else None
+        message = payload.get("message") if payload is not None else None
+        raise _local_transport_error(
+            paths,
+            state=state if isinstance(state, str) else "unhealthy",
+            message=(
+                message
+                if isinstance(message, str)
+                else "The local Server coordinator failed without a valid result."
+            ),
+            retry_after_seconds=_optional_number(payload, "retry_after_seconds"),
+        )
+    if not socket_health(paths):
+        raise _local_transport_error(
+            paths,
+            state="unhealthy",
+            message="The local Server coordinator returned before its socket was healthy.",
+        )
+
+    pid = payload.get("pid")
+    version = payload.get("version")
+    started = payload.get("started")
+    return LocalEnsureResult(
+        started=started is True,
+        pid=pid if isinstance(pid, int) and not isinstance(pid, bool) and pid > 0 else None,
+        server_version=version if isinstance(version, str) else None,
+    )
+
+
+def socket_transport(paths: LocalPaths) -> httpx.HTTPTransport:
+    return httpx.HTTPTransport(uds=str(paths.socket))
+
+
+def socket_health(paths: LocalPaths, *, timeout: float = 0.2) -> bool:
+    try:
+        with httpx.Client(
+            transport=socket_transport(paths),
+            base_url="http://labtasker",
+            timeout=timeout,
+        ) as client:
+            response = client.get("/health")
+        return response.status_code == 200 and response.json() == {
+            "status": "ok",
+            "api_version": "2",
+            "database": "ok",
+        }
+    except (httpx.HTTPError, ValueError):
+        return False
+
+
+def _parse_coordinator_result(output: str) -> dict[str, object] | None:
+    try:
+        payload = json.loads(output)
+    except (json.JSONDecodeError, TypeError):
+        return None
+    if not isinstance(payload, dict):
+        return None
+    return {str(key): value for key, value in payload.items()}
+
+
+def _optional_number(payload: dict[str, object] | None, field: str) -> float | None:
+    if payload is None:
+        return None
+    value = payload.get(field)
+    if isinstance(value, (int, float)) and not isinstance(value, bool):
+        return float(value)
+    return None
+
+
+def _local_transport_error(
+    paths: LocalPaths,
+    *,
+    state: str,
+    message: str,
+    retry_after_seconds: float | None = None,
+) -> TransportError:
+    details: dict[str, object] = {
+        "state": state,
+        "directory": str(paths.directory),
+        "database": str(paths.database),
+        "socket": str(paths.socket),
+        "log": str(paths.log),
+    }
+    if retry_after_seconds is not None:
+        details["retry_after_seconds"] = retry_after_seconds
+    return TransportError(message, details)