offwork 0.4.0__py3-none-any.whl

offwork/_venv.py

@@ -0,0 +1,174 @@
+"""Temporary virtual environment management for offwork."""
+
+import os
+import sys
+import time
+import venv
+import atexit
+import shutil
+import signal
+import asyncio
+import logging
+import tempfile
+from types import FrameType
+from pathlib import Path
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from collections.abc import Sequence, AsyncIterator
+
+logger = logging.getLogger(__name__)
+
+_STALE_THRESHOLD_SECS = 60 * 60 * 24  # 24 hours
+_DEFAULT_PREFIX = "offwork-tmp-"
+
+
+def _find_project_root() -> Path | None:
+    """Walk up from this file to find the directory containing pyproject.toml."""
+    current = Path(__file__).resolve().parent
+    while current != current.parent:
+        if (current / "pyproject.toml").exists():
+            return current
+        current = current.parent
+    return None
+
+
+def _venv_python(venv_dir: Path) -> Path:
+    """Return the path to the Python executable inside a venv."""
+    if sys.platform == "win32":
+        return venv_dir / "Scripts" / "python.exe"
+    return venv_dir / "bin" / "python"
+
+
+def cleanup_stale_venvs(
+    prefix: str = _DEFAULT_PREFIX,
+    max_age_secs: float = _STALE_THRESHOLD_SECS,
+) -> list[str]:
+    """Remove stale offwork temp directories left behind by crashed processes.
+
+    Returns the list of directories that were removed.
+    """
+    tmproot = Path(tempfile.gettempdir())
+    cutoff = time.time() - max_age_secs
+    removed: list[str] = []
+
+    for entry in tmproot.iterdir():
+        if not entry.name.startswith(prefix) or not entry.is_dir():
+            continue
+        try:
+            mtime = entry.stat().st_mtime
+        except OSError:
+            continue
+        if mtime < cutoff:
+            logger.info("Removing stale temp venv: %s", entry)
+            shutil.rmtree(entry, ignore_errors=True)
+            removed.append(str(entry))
+
+    return removed
+
+
+@dataclass
+class TempVenv:
+    """A temporary virtual environment."""
+
+    venv_dir: Path
+    python: Path
+
+    async def pip_install(
+        self, *packages: str, extra_args: Sequence[str] = ()
+    ) -> None:
+        """Install packages via pip in this venv."""
+        if not packages:
+            return
+        cmd = [str(self.python), "-m", "pip", "install", *packages, *extra_args]
+        logger.info("Installing in temp venv: %s", " ".join(packages))
+        proc = await asyncio.create_subprocess_exec(
+            *cmd,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+        _stdout, stderr = await proc.communicate()
+        if proc.returncode != 0:
+            raise RuntimeError(
+                f"pip install failed in temp venv:\n{stderr.decode() if stderr else ''}"
+            )
+
+
+@asynccontextmanager
+async def temp_venv(
+    *,
+    install_offwork: bool = True,
+    extras: Sequence[str] = (),
+    prefix: str = _DEFAULT_PREFIX,
+) -> AsyncIterator[TempVenv]:
+    """Create a temporary venv, optionally install offwork, yield, then clean up.
+
+    Cleanup is guaranteed on normal exit, exceptions, SIGTERM, SIGINT, and
+    ``atexit``.  Stale directories from previously crashed processes are
+    cleaned up on entry.
+
+    Parameters
+    ----------
+    install_offwork:
+        Install offwork into the venv (editable from source tree, or from PyPI).
+    extras:
+        Optional extras to install, e.g. ``("redis",)``.
+    prefix:
+        Prefix for the temporary directory name.
+    """
+    # Opportunistically clean up leftovers from previous crashes.
+    cleanup_stale_venvs(prefix)
+
+    tmpdir = tempfile.mkdtemp(prefix=prefix)
+
+    def _cleanup() -> None:
+        if os.path.isdir(tmpdir):
+            logger.info("Cleaning up temporary venv at %s", tmpdir)
+            print(f"Cleaning up temporary venv at {tmpdir}...", file=sys.stderr)
+            shutil.rmtree(tmpdir, ignore_errors=True)
+
+    # Safety net: atexit ensures cleanup even if the context manager is
+    # bypassed (e.g. an unhandled exception in calling code outside the
+    # ``with`` block, or ``sys.exit()`` called elsewhere).
+    atexit.register(_cleanup)
+
+    # SIGTERM normally kills the process before context-manager __exit__
+    # runs.  Convert it into a clean SystemExit so the finally block and
+    # atexit handlers execute.
+    prev_sigterm = signal.getsignal(signal.SIGTERM)
+
+    def _sigterm_handler(signum: int, frame: FrameType | None) -> None:
+        raise SystemExit(128 + signum)
+
+    signal.signal(signal.SIGTERM, _sigterm_handler)
+
+    try:
+        venv_dir = Path(tmpdir) / "venv"
+        logger.info("Creating temporary venv at %s", venv_dir)
+        print("Creating temporary virtual environment...", file=sys.stderr)
+        loop = asyncio.get_running_loop()
+        await loop.run_in_executor(
+            None, lambda: venv.create(str(venv_dir), with_pip=True)
+        )
+
+        python = _venv_python(venv_dir)
+        if not python.exists():
+            raise RuntimeError(f"venv Python not found at {python}")
+
+        tv = TempVenv(venv_dir=venv_dir, python=python)
+
+        if install_offwork:
+            root = _find_project_root()
+            if root is not None:
+                spec = str(root)
+            else:
+                spec = "offwork"
+            if extras:
+                spec += f"[{','.join(extras)}]"
+            print("Installing offwork into temporary venv...", file=sys.stderr)
+            await tv.pip_install(spec, extra_args=["--quiet"])
+
+        yield tv
+    finally:
+        _cleanup()
+        atexit.unregister(_cleanup)
+        signal.signal(signal.SIGTERM, prev_sigterm)

offwork/core/__init__.py

@@ -0,0 +1,15 @@
+from offwork.core.task import Task, resolve_args
+from offwork.core.errors import Error, RemoteError, WorkerError, DependencyError
+from offwork.core.models import ImportInfo, FunctionNode
+from offwork.core.version import _VERSION
+
+__all__ = [
+    "Task",
+    "resolve_args",
+    "Error",
+    "RemoteError",
+    "WorkerError",
+    "DependencyError",
+    "ImportInfo",
+    "FunctionNode",
+]

offwork/core/errors.py

@@ -0,0 +1,83 @@
+"""Exception hierarchy for offwork."""
+
+import sys
+import types
+
+
+class Error(Exception):
+    """Raised when offwork cannot trace or analyze a function."""
+
+
+class WorkerError(Error):
+    """Raised when worker execution fails."""
+
+
+class DependencyError(Error):
+    """Raised when dependency installation fails."""
+
+
+class RemoteError(WorkerError):
+    """Raised on the client side when a remote execution fails.
+
+    Carries the worker-side traceback and formats it cleanly when the
+    exception is printed, suppressing the noisy client-side frames.
+    """
+
+    remote_traceback: str | None
+
+    def __init__(self, message: str, remote_traceback: str | None = None) -> None:
+        super().__init__(message)
+        self.remote_traceback = remote_traceback
+
+
+class TaskStalled(Error):
+    """Raised when a worker stops sending heartbeats for a task."""
+
+
+class TaskCancelled(Error):
+    """Raised when a task is cancelled before or during execution."""
+
+
+class SignatureError(Error):
+    """Raised when HMAC signature verification of a task fails."""
+
+
+class PairingError(Error):
+    """Raised when the PIN-based pairing protocol fails."""
+
+
+class ThrottleError(Error):
+    """Raised when a task is rejected due to rate limiting."""
+
+
+class WorkerOnlyError(Error):
+    """Raised when a worker-only import stub is used on the client."""
+
+
+# ---------------------------------------------------------------------------
+# Custom excepthook: suppress client traceback for RemoteError
+# ---------------------------------------------------------------------------
+
+_original_excepthook = sys.excepthook
+
+
+def _offwork_excepthook(
+    exc_type: type[BaseException],
+    exc_value: BaseException,
+    exc_tb: types.TracebackType | None,
+) -> None:
+    if isinstance(exc_value, RemoteError) and exc_value.remote_traceback:
+        tb = exc_value.remote_traceback
+        # Replace "Traceback (most recent call last):" with our header
+        tb = tb.replace(
+            "Traceback (most recent call last):",
+            "Worker traceback (most recent call last):",
+            1,
+        )
+        sys.stderr.write(f"\n{tb}")
+        return
+    _original_excepthook(exc_type, exc_value, exc_tb)
+
+
+if sys.excepthook is not _offwork_excepthook:
+    sys.excepthook = _offwork_excepthook

offwork/core/models.py

@@ -0,0 +1,174 @@
+"""Data models for function nodes and import bindings."""
+
+import json
+import hashlib
+from typing import Any, Self
+from dataclasses import field, dataclass
+
+
+@dataclass(frozen=True)
+class ImportInfo:
+    """A single import binding brought into scope by an import statement."""
+
+    statement: str
+    bound_name: str
+    package: str | None = None
+    worker_only: bool = False
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize to a plain dict."""
+        d: dict[str, Any] = {"statement": self.statement, "bound_name": self.bound_name}
+        if self.package is not None:
+            d["package"] = self.package
+        if self.worker_only:
+            d["worker_only"] = True
+        return d
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> Self:
+        """Deserialize from a plain dict.
+
+        Raises
+        ------
+        KeyError
+            If required fields (``statement``, ``bound_name``) are missing.
+        """
+        return cls(
+            statement=data["statement"],
+            bound_name=data["bound_name"],
+            package=data.get("package"),
+            worker_only=bool(data.get("worker_only", False)),
+        )
+
+
+@dataclass
+class FunctionNode:
+    """A node in the dependency graph representing one traced function."""
+
+    qualified_name: str
+    name: str
+    module: str
+    source: str
+    imports: list[ImportInfo] = field(default_factory=list)
+    dependencies: list[str] = field(default_factory=list)
+    owner_class: str | None = None
+    closure_vars: dict[str, str] = field(default_factory=dict)
+    closure_func_refs: dict[str, str] = field(default_factory=dict)
+    module_vars: dict[str, str] = field(default_factory=dict)
+    class_bases: list[str] = field(default_factory=list)
+    class_keywords: dict[str, str] = field(default_factory=dict)
+    class_attrs: list[str] = field(default_factory=list)
+    class_decorators: list[str] = field(default_factory=list)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize to a plain dict including all fields."""
+        d: dict[str, Any] = {
+            "qualified_name": self.qualified_name,
+            "name": self.name,
+            "module": self.module,
+            "source": self.source,
+            "imports": [imp.to_dict() for imp in self.imports],
+            "dependencies": list(self.dependencies),
+            "owner_class": self.owner_class,
+        }
+        if self.closure_vars:
+            d["closure_vars"] = dict(self.closure_vars)
+        if self.closure_func_refs:
+            d["closure_func_refs"] = dict(self.closure_func_refs)
+        if self.module_vars:
+            d["module_vars"] = dict(self.module_vars)
+        if self.class_bases:
+            d["class_bases"] = list(self.class_bases)
+        if self.class_keywords:
+            d["class_keywords"] = dict(self.class_keywords)
+        if self.class_attrs:
+            d["class_attrs"] = list(self.class_attrs)
+        if self.class_decorators:
+            d["class_decorators"] = list(self.class_decorators)
+        return d
+
+    def to_content_blob(self) -> dict[str, Any]:
+        """Return only the fields stored in the content-addressable store.
+
+        Excludes ``qualified_name`` and ``dependencies`` which are graph
+        topology, not intrinsic content.
+        """
+        blob: dict[str, Any] = {
+            "name": self.name,
+            "module": self.module,
+            "source": self.source,
+            "imports": [imp.to_dict() for imp in self.imports],
+            "owner_class": self.owner_class,
+        }
+        if self.closure_vars:
+            blob["closure_vars"] = dict(self.closure_vars)
+        if self.closure_func_refs:
+            blob["closure_func_refs"] = dict(self.closure_func_refs)
+        if self.module_vars:
+            blob["module_vars"] = dict(self.module_vars)
+        if self.class_bases:
+            blob["class_bases"] = list(self.class_bases)
+        if self.class_keywords:
+            blob["class_keywords"] = dict(self.class_keywords)
+        if self.class_attrs:
+            blob["class_attrs"] = list(self.class_attrs)
+        if self.class_decorators:
+            blob["class_decorators"] = list(self.class_decorators)
+        return blob
+
+    def content_hash(self) -> str:
+        """Compute a deterministic content hash for this node.
+
+        The hash covers the node's own content but NOT its dependencies,
+        so adding/removing edges doesn't invalidate existing hashes.
+        """
+        canonical = {
+            "name": self.name,
+            "module": self.module,
+            "source": self.source,
+            "imports": [
+                imp.to_dict()
+                for imp in sorted(self.imports, key=lambda i: i.statement)
+            ],
+            "owner_class": self.owner_class,
+            "closure_vars": dict(sorted(self.closure_vars.items())),
+            "closure_func_refs": dict(sorted(self.closure_func_refs.items())),
+            "module_vars": dict(sorted(self.module_vars.items())),
+            "class_bases": list(self.class_bases),
+            "class_keywords": dict(sorted(self.class_keywords.items())),
+            "class_attrs": list(self.class_attrs),
+            "class_decorators": list(self.class_decorators),
+        }
+        raw = json.dumps(canonical, sort_keys=True, separators=(",", ":"))
+        return hashlib.sha256(raw.encode("utf-8")).hexdigest()[:16]
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> Self:
+        """Deserialize from a plain dict.
+
+        Raises
+        ------
+        KeyError
+            If required fields (``qualified_name``, ``name``, ``module``,
+            ``source``, ``imports``, ``dependencies``) are missing.
+        """
+        required = ("qualified_name", "name", "module", "source", "imports", "dependencies")
+        missing = [k for k in required if k not in data]
+        if missing:
+            raise KeyError(f"FunctionNode missing required fields: {', '.join(missing)}")
+        return cls(
+            qualified_name=data["qualified_name"],
+            name=data["name"],
+            module=data["module"],
+            source=data["source"],
+            imports=[ImportInfo.from_dict(imp) for imp in data["imports"]],
+            dependencies=data["dependencies"],
+            owner_class=data.get("owner_class"),
+            closure_vars=data.get("closure_vars", {}),
+            closure_func_refs=data.get("closure_func_refs", {}),
+            module_vars=data.get("module_vars", {}),
+            class_bases=data.get("class_bases", []),
+            class_keywords=data.get("class_keywords", {}),
+            class_attrs=data.get("class_attrs", []),
+            class_decorators=data.get("class_decorators", []),
+        )