Worker-only imports

Co-authored-by: Copilot <copilot@github.com>

Author: codeSamuraii

docs/QUICK_START.md

@@ -107,6 +107,24 @@ with install_package_as("PyYAML"):
 
 Common mappings (`cv2` → `opencv-python`, `PIL` → `Pillow`, etc.) are built in.
 
+### Worker-only imports
+
+Skip installing packages locally — the worker installs them on demand:
+
+```python
+from pyfuse import worker_only_import
+
+with worker_only_import():
+    import requests
+
+with worker_only_import("opencv-python-headless"):
+    import cv2
+```
+
+The local `requests` and `cv2` resolve to lightweight stubs. They're fine to reference inside a `@trace` function (the worker re-imports them for real), but raise `WorkerOnlyError` if used directly on the client.
+
+Only the names imported literally inside the `with` block are stubbed — real installed packages and their transitive imports are unaffected.
+
 ## Progress, cancellation, and results
 
 ```python

docs/TECHNICAL_OVERVIEW.md

@@ -442,7 +442,8 @@ A single import binding:
 |-------|---------|
 | `statement` | `"import csv"`, `"from os.path import join"` |
 | `bound_name` | `"csv"`, `"join"` |
-| `package` | `"opencv-python"` (from `install_package_as`, or `None`) |
+| `package` | `"opencv-python"` (from `install_package_as` / `worker_only_import`, or `None`) |
+| `worker_only` | `True` if the import was inside a `worker_only_import` block (omitted when `False`) |
 
 Multi-name imports are split into individual objects for per-function tracking.
 
@@ -489,6 +490,24 @@ with install_package_as("opencv-python"):
 
 The worker sees the `package` field on the import and knows to `pip install opencv-python` instead of `pip install cv2`.
 
+### `worker_only_import` context manager
+
+Lets clients reference a package without installing it locally. On the client, imports inside the block resolve to lightweight stub modules (`_WorkerOnlyStub`) that raise `WorkerOnlyError` if used outside a worker context. On the worker, the package is installed via pip and imported normally.
+
+```python
+with worker_only_import():                       # import name == pip name
+    import requests
+
+with worker_only_import("opencv-python-headless"):  # explicit pip name
+    import cv2
+```
+
+Implementation:
+- A meta-path finder is appended to `sys.meta_path` for the duration of the block (real installed packages still win because the finder runs last).
+- The finder's whitelist is populated by parsing the caller's `with` block source via AST — only names the user literally writes inside the block are eligible to be stubbed. This prevents transitive missing imports from real installed packages from being silently swallowed.
+- The AST analyzer marks every `ImportInfo` in the block with `worker_only=True` and records any explicit pip name as `package`.
+- The worker treats worker-only imports identically to regular third-party imports: `extract_third_party_modules` collects them, `_collect_package_hints` honors the pip name, and `pip install` runs before reconstruction.
+
 ## Heartbeat and stall detection
 
 Workers send periodic heartbeats (every 1 second) while executing a task via an `asyncio.Task`. Clients use heartbeat data for stall detection.

examples/package_installation.py

@@ -4,6 +4,11 @@
 pyfuse detects the imports, installs the packages via pip, and executes
 the function -- all automatically.
 
+This example also demonstrates ``worker_only_import``: the client never
+installs ``markdown`` or ``python-dateutil`` -- only the worker does.
+The local imports resolve to lightweight stubs that raise
+``WorkerOnlyError`` if used outside a traced function.
+
 Requires Redis on localhost:6379.  Install: pip install redis
 
 Usage:
@@ -17,12 +22,9 @@
 import asyncio
 from html.parser import HTMLParser
 
-import requests
-import markdown
-
 import pyfuse
 from pyfuse import trace
-from pyfuse import install_package_as
+from pyfuse import install_package_as, worker_only_import
 
 # Some packages have different import and pip names:
 #   import yaml       -> pip install PyYAML
@@ -33,7 +35,18 @@
 with install_package_as("PyYAML"):
     import yaml
 
-with install_package_as("python-dateutil"):
+# `worker_only_import` skips the local install entirely. The package only
+# needs to be available on the worker. The client gets a stub object that
+# raises WorkerOnlyError if used directly, but is fine to reference inside
+# a @trace function (it's serialized and re-imported on the worker).
+
+# No pip name needed when import name == package name:
+with worker_only_import():
+    import requests
+    import markdown
+
+# Pass an explicit pip name when the import name differs from the package:
+with worker_only_import("python-dateutil"):
     from dateutil import parser as date_parser
 
 

pyfuse/__init__.py

@@ -15,11 +15,12 @@
     SignatureError,
     DependencyError,
     ThrottleError,
+    WorkerOnlyError,
 )
 from pyfuse.core.models import ImportInfo, FunctionNode
 from pyfuse.graph.graph import Graph
 from pyfuse.graph.store import Store, MergeResult
-from pyfuse.worker.deps import install_package_as
+from pyfuse.worker.deps import install_package_as, worker_only_import
 from pyfuse.core.pairing import (
     PairingResult,
     generate_pin,
@@ -111,6 +112,7 @@ def pack(func: Callable[..., object], *args: Any, **kwargs: Any) -> Task:
     "disconnect",
     "serve",
     "install_package_as",
+    "worker_only_import",
     "progress",
     # Serialization
     "serialize",
@@ -130,6 +132,7 @@ def pack(func: Callable[..., object], *args: Any, **kwargs: Any) -> Task:
     "ThrottleError",
     "SignatureError",
     "PairingError",
+    "WorkerOnlyError",
     # Graph
     "get_graph",
     "Graph",

pyfuse/__main__.py

@@ -28,7 +28,7 @@
 from pyfuse._venv import temp_venv
 from pyfuse.worker.deps import DEFAULT_IMPORT_TO_PACKAGE
 from pyfuse.worker.remote import serve
-from pyfuse.graph.analyzer import _parse_install_package_as
+from pyfuse.graph.analyzer import _parse_install_package_as, _parse_worker_only_import
 
 
 def _build_worker_cmd(python: str, args: argparse.Namespace) -> list[str]:
@@ -228,6 +228,7 @@ def _detect_script_packages(script: str) -> list[str]:
 
     # module name -> pip package name (None means use default mapping)
     modules: dict[str, str | None] = {}
+    skip: set[str] = set()
     for node in ast.iter_child_nodes(tree):
         if isinstance(node, (ast.Import, ast.ImportFrom)):
             for m in _extract_top_modules(node):
@@ -238,11 +239,18 @@ def _detect_script_packages(script: str) -> list[str]:
                 for child in node.body:
                     for m in _extract_top_modules(child):
                         modules[m] = package
+                continue
+            if _parse_worker_only_import(node) is not False:
+                for child in node.body:
+                    for m in _extract_top_modules(child):
+                        skip.add(m)
 
     packages: dict[str, None] = {}
     for m, explicit_package in sorted(modules.items()):
         if m in sys.stdlib_module_names or m == "pyfuse":
             continue
+        if m in skip:
+            continue
         if _is_local_package(m, script_dir):
             continue
         pkg = explicit_package or DEFAULT_IMPORT_TO_PACKAGE.get(m, m)

pyfuse/core/errors.py

@@ -50,6 +50,10 @@ 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
 # ---------------------------------------------------------------------------

pyfuse/core/models.py

@@ -13,16 +13,19 @@ class ImportInfo:
     statement: str
     bound_name: str
     package: str | None = None
+    worker_only: bool = False
 
-    def to_dict(self) -> dict[str, str]:
+    def to_dict(self) -> dict[str, Any]:
         """Serialize to a plain dict."""
-        d = {"statement": self.statement, "bound_name": self.bound_name}
+        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, str]) -> Self:
+    def from_dict(cls, data: dict[str, Any]) -> Self:
         """Deserialize from a plain dict.
 
         Raises
@@ -34,6 +37,7 @@ def from_dict(cls, data: dict[str, str]) -> Self:
             statement=data["statement"],
             bound_name=data["bound_name"],
             package=data.get("package"),
+            worker_only=bool(data.get("worker_only", False)),
         )
 
 

pyfuse/graph/analyzer.py

@@ -74,6 +74,15 @@ def get_module_imports(func: Callable[..., object]) -> list[ImportInfo]:
                         imports.extend(_extract_import(child, package))
                     elif isinstance(child, ast.ImportFrom):
                         imports.extend(_extract_import_from(child, package))
+                continue
+            wo = _parse_worker_only_import(node)
+            if wo is not False:
+                wo_package = wo if isinstance(wo, str) else None
+                for child in node.body:
+                    if isinstance(child, ast.Import):
+                        imports.extend(_extract_import(child, wo_package, worker_only=True))
+                    elif isinstance(child, ast.ImportFrom):
+                        imports.extend(_extract_import_from(child, wo_package, worker_only=True))
 
     logger.debug(
         "Found %d import bindings in %s", len(imports), source_file
@@ -129,19 +138,57 @@ def _parse_install_package_as(node: ast.With) -> str | None:
     return ctx.args[0].value
 
 
+def _is_worker_only_import_call(expr: ast.expr) -> bool:
+    """Match ``worker_only_import`` or ``pyfuse.worker_only_import``."""
+    if isinstance(expr, ast.Name):
+        return expr.id == "worker_only_import"
+    if isinstance(expr, ast.Attribute):
+        return expr.attr == "worker_only_import"
+    return False
+
+
+def _parse_worker_only_import(node: ast.With) -> str | bool:
+    """Detect ``with worker_only_import([package]):`` blocks.
+
+    Returns ``False`` if not a match, ``True`` if matched without a
+    package argument, or the package string if one was supplied.
+    """
+    if len(node.items) != 1:
+        return False
+    ctx = node.items[0].context_expr
+    if not isinstance(ctx, ast.Call) or not _is_worker_only_import_call(ctx.func):
+        return False
+    if not ctx.args:
+        return True
+    if (
+        len(ctx.args) == 1
+        and isinstance(ctx.args[0], ast.Constant)
+        and isinstance(ctx.args[0].value, str)
+    ):
+        return ctx.args[0].value
+    return True
+
+
 def _extract_import(
-    node: ast.Import, package: str | None = None
+    node: ast.Import,
+    package: str | None = None,
+    worker_only: bool = False,
 ) -> list[ImportInfo]:
     result: list[ImportInfo] = []
     for alias in node.names:
         bound = alias.asname or alias.name.split(".")[0]
         stmt = ast.unparse(ast.Import(names=[alias]))
-        result.append(ImportInfo(statement=stmt, bound_name=bound, package=package))
+        result.append(ImportInfo(
+            statement=stmt, bound_name=bound,
+            package=package, worker_only=worker_only,
+        ))
     return result
 
 
 def _extract_import_from(
-    node: ast.ImportFrom, package: str | None = None
+    node: ast.ImportFrom,
+    package: str | None = None,
+    worker_only: bool = False,
 ) -> list[ImportInfo]:
     result: list[ImportInfo] = []
     for alias in node.names:
@@ -170,7 +217,10 @@ def _extract_import_from(
                 for export_name in exported:
                     stmt = f"from {node.module} import {export_name}"
                     result.append(
-                        ImportInfo(statement=stmt, bound_name=export_name, package=package)
+                        ImportInfo(
+                            statement=stmt, bound_name=export_name,
+                            package=package, worker_only=worker_only,
+                        )
                     )
             except ImportError:
                 warnings.warn(
@@ -185,7 +235,10 @@ def _extract_import_from(
                 module=node.module, names=[alias], level=node.level
             )
         )
-        result.append(ImportInfo(statement=stmt, bound_name=bound, package=package))
+        result.append(ImportInfo(
+            statement=stmt, bound_name=bound,
+            package=package, worker_only=worker_only,
+        ))
     return result
 
 

pyfuse/worker/__init__.py

@@ -1,10 +1,11 @@
-from pyfuse.worker.deps import install_package_as, ensure_dependencies
+from pyfuse.worker.deps import install_package_as, worker_only_import, ensure_dependencies
 from pyfuse.worker.remote import serve, connect, disconnect, submit_remote
 from pyfuse.worker.result import Result, ResultEnvelope
 from pyfuse.worker.worker import Worker, execute
 
 __all__ = [
     "install_package_as",
+    "worker_only_import",
     "ensure_dependencies",
     "serve",
     "connect",