cortexgrid 0.2.85__tar.gz → 0.2.86__tar.gz

{cortexgrid-0.2.85 → cortexgrid-0.2.86}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: cortexgrid
-Version: 0.2.85
+Version: 0.2.86
 Summary: Connect your ML code to the RoboLab compute cluster — Ray, MLflow, and S3
 Project-URL: Homepage, https://github.com/robodatalab/cortexgrid
 Project-URL: Repository, https://github.com/robodatalab/cortexgrid
@@ -13,6 +13,7 @@ Requires-Dist: cloudpickle>=3.0
 Requires-Dist: fabric>=3.2.3
 Requires-Dist: haikunator>=2.1.0
 Requires-Dist: mlflow<4,>=3.11
+Requires-Dist: packaging>=24
 Requires-Dist: pip>=23.0
 Requires-Dist: pydantic-settings>=2.13.1
 Requires-Dist: pydantic>=2.13.3
@@ -115,9 +116,10 @@ print(f"Submitted: {job_id}")
 A separate service — the **jobs control plane** — polls MLflow for pending job requests, matches them against the set of Ray submissions the cluster already has, and submits anything missing. It is also responsible for retrying failed jobs and honouring user-requested stops.
 
 Each submission captures the code and dependencies the entry function needs automatically ([_bundle.py](https://github.com/robodatalab/cortexgrid/blob/main/cortexgrid/_bundle.py)):
-- `bundle(entry)` traces the import graph from the function's source file, resolving each import the way the interpreter does (via `sys.path`), and returns every file needed to run it -- your own modules and third-party packages alike, wherever they live. The standard library is excluded (it ships with the interpreter)
-- Everything ships **as source**: the bundle is staged at each file's import path and tarred into the Ray `working_dir`. Nothing is `pip`-installed on the worker
-- Dependencies the worker image already has are subtracted rather than shipped: `bundle(entry) - worker_provides()`, where `worker_provides()` is the bundle of the packages baked into the ray image (torch and its CUDA stack, ray, mlflow, ...). See [k8s/docker/ray/Dockerfile](https://github.com/robodatalab/cortexgrid/blob/main/k8s/docker/ray/Dockerfile)
+- `bundle(entry)` traces the import graph from the function's source file, resolving each import the way the interpreter does (via `sys.path`). The standard library is excluded (it ships with the interpreter)
+- Your own modules -- anything outside site-packages / dist-packages -- ship **as source**: they are staged at their import paths and tarred into the Ray `working_dir`
+- Third-party packages are recorded as the installed distribution that owns the imported file, pinned to its installed version (`tqdm==4.67.3`), and Ray **pip-installs** them on the worker into a per-node cached virtualenv layered on the image (`runtime_env["pip"]`). An import into site-packages that no installed distribution owns fails `cortexgrid.remote` with `UnownedDependencyError`
+- Distributions the worker image already has are not installed again: `worker_provides()` is the dependency closure of the packages baked into the ray image (torch and its CUDA stack, ray, mlflow, ...), and is subtracted from the pip list. See [k8s/docker/ray/Dockerfile](https://github.com/robodatalab/cortexgrid/blob/main/k8s/docker/ray/Dockerfile) and `_WORKER_BAKED` in `_bundle.py`, which must list the same packages
 - Injects MLflow/S3 credentials so task code running on the DGX can reach all services
 
 ##### Retries

{cortexgrid-0.2.85 → cortexgrid-0.2.86}/cortexgrid/_bundle.py

@@ -3,27 +3,38 @@
 `bundle(seed)` describes what is needed to run the module at `seed`: the local
 files it reaches -- following its import graph and each package's __init__
 chain, resolving imports the way the interpreter does -- and the third-party
-dependencies those files import. A file is local unless it lives in an installed
-package location (site-packages / dist-packages); imports that resolve into one
-are not followed. The standard library is excluded (it ships with the
-interpreter). Bundles of several seeds combine with `BundleDesc.merge`.
+distributions those files import. A file is local unless it lives in an installed
+package location (site-packages / dist-packages). An import that resolves into
+one is recorded as the distribution that installed the file, at its installed
+version, and is not followed: installing that distribution brings its own
+dependencies. The standard library is excluded (it ships with the interpreter).
+Bundles of several seeds combine with `BundleDesc.merge`.
 
 `stage(files, dest)` lays a bundle out under `dest` at each file's import path,
 so `dest` on sys.path (e.g. a Ray working_dir) makes every module importable.
+
+`BundleDesc.pip_requirements(worker_provides())` pins the third-party
+distributions the Ray worker image does not already have, for a Ray `pip`
+runtime_env to install on the worker.
 """
 
 from __future__ import annotations
 
 import ast
-from collections.abc import Iterator
+from collections.abc import Iterable, Iterator
 from dataclasses import dataclass
 import functools
 import importlib.machinery
+import importlib.metadata
 import importlib.util
+import os
 from pathlib import Path
 import shutil
 import sys
 
+from packaging.requirements import Requirement
+from packaging.utils import canonicalize_name
+
 
 ThirdPartyDependencyName = str
 ThirdPartyDependencyVersion = str
@@ -39,17 +50,43 @@ class BundleDesc:
             tp_deps={**self.tp_deps, **other.tp_deps},
         )
 
+    def pip_requirements(
+        self, provided: frozenset[ThirdPartyDependencyName]
+    ) -> list[str]:
+        """The third-party distributions, pinned (`name==version`) and sorted,
+        minus the `provided` ones."""
+        return sorted(
+            f"{name}=={version}"
+            for name, version in self.tp_deps.items()
+            if name not in provided
+        )
+
+
+class UnownedDependencyError(LookupError):
+    """An import resolved into an installed package location, but no installed
+    distribution owns the file, so it can be neither shipped nor installed."""
+
 
 def bundle(seed: Path) -> BundleDesc:
     """What is needed to run the module at `seed`: its local files, each at its
-    real path. Third-party dependency detection is not implemented yet, so
-    `tp_deps` is always empty."""
+    real path, and the third-party distributions they import, keyed by
+    canonical name, each at its installed version.
+
+    Raises UnownedDependencyError for an import that resolves into an installed
+    package location no distribution owns."""
     seed = seed.resolve()
     files: set[Path] = set()
+    tp_deps: dict[ThirdPartyDependencyName, ThirdPartyDependencyVersion] = {}
+    visited: set[Path] = set()
     queue: list[Path] = [seed]
     while queue:
         file = queue.pop()
-        if file in files or not _is_local(file):
+        if file in visited:
+            continue
+        visited.add(file)
+        if not _is_local(file):
+            dist = _owning_distribution(file)
+            tp_deps[canonicalize_name(dist.metadata["Name"])] = dist.version
             continue
         files.add(file)
         queue.extend(_init_chain(file))  # importing a module runs its __init__ chain
@@ -58,7 +95,7 @@ def bundle(seed: Path) -> BundleDesc:
                 dep = _module_file(name)
                 if dep is not None:
                     queue.append(dep)
-    return BundleDesc(local_files=files, tp_deps={})
+    return BundleDesc(local_files=files, tp_deps=tp_deps)
 
 
 def stage(files: set[Path], dest: Path) -> None:
@@ -70,23 +107,120 @@ def stage(files: set[Path], dest: Path) -> None:
         shutil.copy2(file, target)
 
 
-# Distributions already present in the Ray worker image (k8s/docker/ray/Dockerfile).
-# They and their whole dependency trees -- torch's CUDA stack, sympy, numpy, ray,
-# mlflow, ... -- are on the worker already, so a caller subtracts them from a
-# bundle rather than shipping them again.
-_WORKER_BAKED = ("ray", "mlflow", "torch", "smart_open", "dotenv", "psutil")
+# What the Ray worker image pip-installs, as the Dockerfile spells it. They and
+# their dependency trees are on the worker already, so they are never installed
+# there again -- a second copy in the job's virtualenv would shadow the image's.
+#
+# KEEP IN SYNC with k8s/docker/ray/Dockerfile, by hand: every package it
+# pip-installs must be listed here. A package missing here gets installed a
+# second time on the worker; one listed here but no longer in the image is
+# never installed at all.
+_WORKER_BAKED = (
+    "ray[default,serve]",
+    "smart_open[s3]",
+    "mlflow",
+    "python-dotenv",
+    "psutil",
+    "nvidia-cublas-cu12",
+    "nvidia-cudnn-cu12",
+    "nvidia-cuda-nvrtc-cu12",
+    "nvidia-cuda-runtime-cu12",
+    "nvidia-cuda-cupti-cu12",
+    "nvidia-cufft-cu12",
+    "nvidia-curand-cu12",
+    "nvidia-cusolver-cu12",
+    "nvidia-cusparse-cu12",
+    "nvidia-cusparselt-cu12",
+    "nvidia-nccl-cu12",
+    "nvidia-nvshmem-cu12",
+    "nvidia-nvtx-cu12",
+    "nvidia-nvjitlink-cu12",
+    "nvidia-cufile-cu12",
+    "cuda-bindings",
+    "triton",
+    "torch",
+    "filelock",
+    "typing-extensions",
+    "sympy",
+    "networkx",
+    "jinja2",
+    "fsspec",
+)
+
+
+@functools.lru_cache(maxsize=1)
+def worker_provides() -> frozenset[ThirdPartyDependencyName]:
+    """Every distribution the Ray worker image already provides:
+    `distribution_closure(_WORKER_BAKED)`."""
+    return distribution_closure(_WORKER_BAKED)
+
+
+def distribution_closure(
+    requirements: Iterable[str],
+) -> frozenset[ThirdPartyDependencyName]:
+    """Canonical names of the distributions `requirements` name plus everything
+    they depend on, transitively, as this environment's installed metadata
+    declares it -- extras followed where requested, environment markers
+    evaluated here. A requirement not installed here contributes only its own
+    name: its dependencies are unknown."""
+    names: set[ThirdPartyDependencyName] = set()
+    seen: set[tuple[str, frozenset[str]]] = set()
+    queue = [Requirement(spec) for spec in requirements]
+    while queue:
+        requirement = queue.pop()
+        name = canonicalize_name(requirement.name)
+        key = (name, frozenset(requirement.extras))
+        if key in seen:
+            continue
+        seen.add(key)
+        names.add(name)
+        try:
+            dist = importlib.metadata.distribution(name)
+        except importlib.metadata.PackageNotFoundError:
+            continue
+        extras = requirement.extras or {""}
+        for spec in dist.requires or ():
+            dependency = Requirement(spec)
+            if dependency.marker is None or any(
+                dependency.marker.evaluate({"extra": extra}) for extra in extras
+            ):
+                queue.append(dependency)
+    return frozenset(names)
+
+
+def _owning_distribution(file: Path) -> importlib.metadata.Distribution:
+    """The installed distribution whose file list (RECORD) contains `file`.
+    Rebuilds the index once on a miss, in case something was installed since it
+    was built."""
+    path = tuple(sys.path)
+    dist = _distribution_index(path).get(file)
+    if dist is None:
+        _distribution_index.cache_clear()
+        dist = _distribution_index(path).get(file)
+    if dist is None:
+        raise UnownedDependencyError(
+            f"{file} is imported from an installed package location, but no "
+            "installed distribution lists it among its files, so it can be "
+            "neither shipped nor pip-installed on the worker. Install it with "
+            "pip or uv so it carries distribution metadata."
+        )
+    return dist
 
 
 @functools.lru_cache(maxsize=1)
-def worker_provides() -> frozenset[Path]:
-    """Every file the Ray worker image already provides. Subtract from a bundle
-    before shipping: `bundle(entry).local_files - worker_provides()`."""
-    provided: set[Path] = set()
-    for name in _WORKER_BAKED:
-        origin = _module_file(name)
-        if origin is not None:
-            provided |= bundle(origin).local_files
-    return frozenset(provided)
+def _distribution_index(
+    path: tuple[str, ...],
+) -> dict[Path, importlib.metadata.Distribution]:
+    """Every file installed by a distribution found on `path` (sys.path, which
+    the cache is keyed on), mapped to that distribution. Each distribution's
+    root is resolved once; its files are joined onto it lexically, which keeps
+    this fast for environments with tens of thousands of files."""
+    index: dict[Path, importlib.metadata.Distribution] = {}
+    for dist in importlib.metadata.distributions(path=list(path)):
+        root = Path(dist.locate_file("")).resolve()
+        for file in dist.files or ():
+            index[Path(os.path.normpath(root / file))] = dist
+    return index
 
 
 def _is_local(file: Path) -> bool:

{cortexgrid-0.2.85 → cortexgrid-0.2.86}/cortexgrid/jobs.py

@@ -63,6 +63,9 @@ class JobLifecycle:
     retry: bool = False  # static flag set at job creation
     num_gpus: int = 0
     num_cpus: int = 1
+    # static: pinned third-party requirements the worker pip-installs (the
+    # bundle's distributions the Ray image does not already provide)
+    pip_requirements: list[str] = field(default_factory=list)
     history: list[LifecycleEvent] = field(default_factory=list)
 
     def to_json(self) -> str:
@@ -229,11 +232,17 @@ def schedule_remote_job(
     job_id = Haikunator().haikunate(token_length=2, token_chars="0123456789")
     entry_file = Path(inspect.getfile(fn)).resolve()
     driver_file = Path(__file__).with_name("_ray_job_driver.py")
-    files = bundle(entry_file).merge(bundle(driver_file)).local_files - worker_provides()
+    desc = bundle(entry_file).merge(bundle(driver_file))
+    pip_requirements = desc.pip_requirements(worker_provides())
     with tempfile.TemporaryDirectory() as tmp:
         code_root = Path(tmp, "project_code_root")
-        stage(files, code_root)
-        log.info("Submitting job %s (%d files)", job_id, len(files))
+        stage(desc.local_files, code_root)
+        log.info(
+            "Submitting job %s (%d files, pip: %s)",
+            job_id,
+            len(desc.local_files),
+            pip_requirements,
+        )
         Payload(
             experiment_name=experiment_name,
             run_id=run_id,
@@ -252,6 +261,7 @@ def schedule_remote_job(
             retry=retry,
             num_gpus=num_gpus,
             num_cpus=num_cpus,
+            pip_requirements=pip_requirements,
         ).save_to_mlflow()
     return job_id
 

{cortexgrid-0.2.85 → cortexgrid-0.2.86}/cortexgrid/model_serving.py

@@ -19,11 +19,12 @@ the end-to-end design.
 from __future__ import annotations
 
 import inspect
+import json
 import logging
 import shutil
 import tempfile
 import time
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any
 
@@ -90,6 +91,9 @@ class BundleMetadata:
 
     bundle_url: str
     class_import_path: str
+    # pinned third-party requirements the replica pip-installs (the bundle's
+    # distributions the Ray image does not already provide)
+    pip_requirements: list[str] = field(default_factory=list)
 
 
 def bundle_class(
@@ -103,12 +107,18 @@ def bundle_class(
     without holding the class object."""
     entry_file = Path(inspect.getfile(cls)).resolve()
     serve_entry = Path(__file__).with_name("_serve_entry.py")
-    files = bundle(entry_file).merge(bundle(serve_entry)).local_files - worker_provides()
+    desc = bundle(entry_file).merge(bundle(serve_entry))
+    pip_requirements = desc.pip_requirements(worker_provides())
     with tempfile.TemporaryDirectory() as tmp:
         code_root = Path(tmp) / "code"
-        stage(files, code_root)
+        stage(desc.local_files, code_root)
         log.info(
-            "Serve bundle for %s/%s/%s: %d files", family, suffix, run_name, len(files)
+            "Serve bundle for %s/%s/%s: %d files, pip: %s",
+            family,
+            suffix,
+            run_name,
+            len(desc.local_files),
+            pip_requirements,
         )
         zip_base = Path(tmp) / f"{family}__{suffix}"
         shutil.make_archive(str(zip_base), "zip", root_dir=str(code_root))
@@ -119,6 +129,7 @@ def bundle_class(
     return BundleMetadata(
         bundle_url=bundle_url,
         class_import_path=f"{cls.__module__}:{cls.__name__}",
+        pip_requirements=pip_requirements,
     )
 
 
@@ -126,6 +137,13 @@ def _build_application_spec(
     family: str, suffix: str, run_name: str, meta: BundleMetadata
 ) -> dict[str, Any]:
     """Assemble a Ray Serve application schema from pre-bundled metadata."""
+    # working_dir carries the serve-app's own source; Ray pip-installs the
+    # third-party distributions the image lacks into a per-node cached
+    # virtualenv layered on the image. No pip key when there are none, so Ray
+    # builds no virtualenv.
+    runtime_env: dict[str, Any] = {"working_dir": meta.bundle_url}
+    if meta.pip_requirements:
+        runtime_env["pip"] = meta.pip_requirements
     return {
         "name": _app_name(family, suffix, run_name),
         "route_prefix": _route_prefix(family, suffix, run_name),
@@ -140,9 +158,7 @@ def _build_application_spec(
             "suffix": suffix,
             "run_name": run_name,
         },
-        # Every dependency ships as source inside the bundle, so working_dir
-        # alone makes the serve app importable; nothing is pip-installed.
-        "runtime_env": {"working_dir": meta.bundle_url},
+        "runtime_env": runtime_env,
     }
 
 
@@ -150,6 +166,7 @@ def _build_application_spec(
 # `deploy_model` reads back.
 _CLASS_IMPORT_PATH_TAG = "class_import_path"
 _BUNDLE_URL_TAG = "serve_bundle_url"
+_PIP_REQUIREMENTS_TAG = "serve_pip_requirements"
 
 
 def metadata_to_tags(meta: BundleMetadata) -> dict[str, str]:
@@ -159,6 +176,7 @@ def metadata_to_tags(meta: BundleMetadata) -> dict[str, str]:
     return {
         _CLASS_IMPORT_PATH_TAG: meta.class_import_path,
         _BUNDLE_URL_TAG: meta.bundle_url,
+        _PIP_REQUIREMENTS_TAG: json.dumps(meta.pip_requirements),
     }
 
 
@@ -180,6 +198,8 @@ def _load_bundle_metadata(
         return BundleMetadata(
             bundle_url=tags[_BUNDLE_URL_TAG],
             class_import_path=tags[_CLASS_IMPORT_PATH_TAG],
+            # Absent on models saved before dependencies were pip-installed.
+            pip_requirements=json.loads(tags.get(_PIP_REQUIREMENTS_TAG, "[]")),
         )
     except KeyError as exc:
         raise ValueError(

{cortexgrid-0.2.85 → cortexgrid-0.2.86}/docs/cortexgrid/README.md

@@ -88,9 +88,10 @@ print(f"Submitted: {job_id}")
 A separate service — the **jobs control plane** — polls MLflow for pending job requests, matches them against the set of Ray submissions the cluster already has, and submits anything missing. It is also responsible for retrying failed jobs and honouring user-requested stops.
 
 Each submission captures the code and dependencies the entry function needs automatically ([_bundle.py](https://github.com/robodatalab/cortexgrid/blob/main/cortexgrid/_bundle.py)):
-- `bundle(entry)` traces the import graph from the function's source file, resolving each import the way the interpreter does (via `sys.path`), and returns every file needed to run it -- your own modules and third-party packages alike, wherever they live. The standard library is excluded (it ships with the interpreter)
-- Everything ships **as source**: the bundle is staged at each file's import path and tarred into the Ray `working_dir`. Nothing is `pip`-installed on the worker
-- Dependencies the worker image already has are subtracted rather than shipped: `bundle(entry) - worker_provides()`, where `worker_provides()` is the bundle of the packages baked into the ray image (torch and its CUDA stack, ray, mlflow, ...). See [k8s/docker/ray/Dockerfile](https://github.com/robodatalab/cortexgrid/blob/main/k8s/docker/ray/Dockerfile)
+- `bundle(entry)` traces the import graph from the function's source file, resolving each import the way the interpreter does (via `sys.path`). The standard library is excluded (it ships with the interpreter)
+- Your own modules -- anything outside site-packages / dist-packages -- ship **as source**: they are staged at their import paths and tarred into the Ray `working_dir`
+- Third-party packages are recorded as the installed distribution that owns the imported file, pinned to its installed version (`tqdm==4.67.3`), and Ray **pip-installs** them on the worker into a per-node cached virtualenv layered on the image (`runtime_env["pip"]`). An import into site-packages that no installed distribution owns fails `cortexgrid.remote` with `UnownedDependencyError`
+- Distributions the worker image already has are not installed again: `worker_provides()` is the dependency closure of the packages baked into the ray image (torch and its CUDA stack, ray, mlflow, ...), and is subtracted from the pip list. See [k8s/docker/ray/Dockerfile](https://github.com/robodatalab/cortexgrid/blob/main/k8s/docker/ray/Dockerfile) and `_WORKER_BAKED` in `_bundle.py`, which must list the same packages
 - Injects MLflow/S3 credentials so task code running on the DGX can reach all services
 
 ##### Retries

{cortexgrid-0.2.85 → cortexgrid-0.2.86}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "cortexgrid"
-version = "0.2.85"
+version = "0.2.86"
 description = "Connect your ML code to the RoboLab compute cluster — Ray, MLflow, and S3"
 readme = "docs/cortexgrid/README.md"
 license = "Apache-2.0"
@@ -9,6 +9,7 @@ requires-python = ">=3.11"
 dependencies = [
     "ray[default]>=2.9,<3",
     "mlflow>=3.11,<4",
+    "packaging>=24",
     "boto3>=1.34",
     "setuptools>=82.0.1",
     "tqdm>=4.60",