fix(apps): keep transitive torch out of GPU artifacts

runpod/apps/build.py

@@ -1,14 +1,17 @@
 """deploy-time environment build: vendor dependencies into the artifact."""
 
+import csv
 import logging
 import os
+import posixpath
 import re
 import shutil
 import subprocess
 import sys
 import tempfile
 from dataclasses import dataclass, field
 from collections import deque
+from email.parser import HeaderParser
 from pathlib import Path
 from typing import List, Optional, Tuple
 
@@ -59,6 +62,12 @@
 
 _REQ_NAME_RE = re.compile(r"^\s*([A-Za-z0-9][A-Za-z0-9._-]*)")
 
+# a name only when followed by extras, a specifier, a marker, a direct
+# reference, or the end (not a url, path, or archive filename)
+_STRICT_REQ_NAME_RE = re.compile(
+    r"\s*([A-Za-z0-9][A-Za-z0-9._-]*)\s*(?:\[[^\]]*\])?\s*(?:$|[<>=!~;@(])"
+)
+
 
 class BuildError(AppError):
     """environment vendoring failed."""
@@ -79,16 +88,21 @@ def requirement_name(requirement: str) -> str:
     return (match.group(1) if match else requirement).lower().replace("_", "-")
 
 
-def collect_requirements(project_root: Path, app) -> List[str]:
-    """project requirements.txt plus per-resource dependencies, deduped."""
+def _requirements_txt(project_root: Path) -> List[str]:
+    """requirement lines from the project's requirements.txt."""
     requirements: List[str] = []
-
     req_file = project_root / "requirements.txt"
     if req_file.exists():
         for line in req_file.read_text().splitlines():
             line = line.strip()
             if line and not line.startswith("#") and not line.startswith("-"):
                 requirements.append(line)
+    return requirements
+
+
+def collect_requirements(project_root: Path, app) -> List[str]:
+    """project requirements.txt plus per-resource dependencies, deduped."""
+    requirements = _requirements_txt(project_root)
 
     for handle in app.resources.values():
         for dep in handle.spec.dependencies or []:
@@ -135,6 +149,208 @@ def split_exclusions(
     return kept, sorted(set(excluded))
 
 
+def _canonical_name(name: str) -> str:
+    """pep 503 normalized distribution name."""
+    return re.sub(r"[-_.]+", "-", name).lower()
+
+
+def _strict_requirement_name(requirement: str) -> str:
+    """canonical name of a requirement or Requires-Dist line, or "" when
+    the line does not start with a name (url, path, archive filename)."""
+    match = _STRICT_REQ_NAME_RE.match(requirement)
+    if not match:
+        return ""
+    if requirement.strip() == match.group(1) and re.search(
+        r"\.(whl|zip|tar\.gz|tgz)$", match.group(1), re.I
+    ):
+        return ""
+    return _canonical_name(match.group(1))
+
+
+def _dist_info_name(dist_info_dir: str) -> str:
+    """distribution name from a "name-version.dist-info" directory name."""
+    stem = dist_info_dir[: -len(".dist-info")]
+    match = re.search(r"-\d", stem)
+    return stem[: match.start()] if match else stem
+
+
+def _read_record(dist_info: Path) -> Optional[set]:
+    """paths a wheel installed, relative to the install dir, or None.
+
+    pip --target records scripts and data files relative to its staging
+    lib/python dir ("../../bin/torchrun") and then moves them to the top
+    of the target, so one leading "../../" maps onto the target.
+    """
+    record = dist_info / "RECORD"
+    if not record.is_file():
+        return None
+    paths = set()
+    with record.open(newline="", encoding="utf-8", errors="replace") as f:
+        for row in csv.reader(f):
+            if not row:
+                continue
+            path = posixpath.normpath(row[0].replace("\\", "/"))
+            if path.startswith("../../"):
+                path = path[len("../../") :]
+            if not path.startswith("../") and path not in (".", ".."):
+                paths.add(path)
+    return paths
+
+
+def _read_requires(dist_info: Path) -> Optional[set]:
+    """canonical names a wheel depends on, or None when unreadable.
+
+    markers and extras are ignored, which can only keep more wheels.
+    """
+    metadata = dist_info / "METADATA"
+    if not metadata.is_file():
+        return None
+    headers = HeaderParser().parsestr(
+        metadata.read_text(encoding="utf-8", errors="replace")
+    )
+    deps = set()
+    for line in headers.get_all("Requires-Dist") or []:
+        dep = _strict_requirement_name(line)
+        if not dep:
+            return None
+        deps.add(dep)
+    return deps
+
+
+def installed_closure(env_dir: Path, roots: List[str]) -> set:
+    """canonical names of installed distributions the roots pull in,
+    roots included. unreadable metadata counts every installed
+    distribution as pulled in."""
+    requires = {
+        _canonical_name(_dist_info_name(d.name)): _read_requires(d)
+        for d in env_dir.glob("*.dist-info")
+        if d.is_dir()
+    }
+    names = {_strict_requirement_name(r) for r in roots}
+    if "" in names or any(v is None for v in requires.values()):
+        return set(requires) | (names - {""})
+    seen = set()
+    stack = list(names)
+    while stack:
+        node = stack.pop()
+        if node in seen:
+            continue
+        seen.add(node)
+        stack.extend(requires.get(node) or ())
+    return seen
+
+
+def prune_excluded(
+    env_dir: Path,
+    excluded: List[str],
+    roots: List[str],
+    image_provided: List[str],
+) -> List[str]:
+    """remove excluded distributions that pip installed as dependencies.
+
+    split_exclusions only drops excluded packages the user listed. a
+    listed package that depends on one (peft -> torch) still installs it,
+    and on the worker the vendored copy shadows the image's build. this
+    removes every installed excluded distribution, plus the wheels that
+    only image_provided packages depend on (torch's nvidia/cuda, sympy,
+    networkx), since the image installs those alongside its own torch.
+
+    fails closed: without METADATA and RECORD for every wheel, or with a
+    root requirement that has no name (url, path), only the excluded
+    distributions themselves are removed. roots and their dependencies
+    are never removed unless excluded.
+
+    returns the canonical names of removed distributions.
+    """
+    dist_infos = {
+        _canonical_name(_dist_info_name(d.name)): d
+        for d in env_dir.glob("*.dist-info")
+        if d.is_dir()
+    }
+    excluded_set = {_canonical_name(n) for n in excluded}
+    targets = set(dist_infos) & excluded_set
+    if not targets:
+        return []
+
+    records = {name: _read_record(d) for name, d in dist_infos.items()}
+    requires = {name: _read_requires(d) for name, d in dist_infos.items()}
+    root_names = {_strict_requirement_name(r) for r in roots}
+    complete = all(v is not None for v in records.values()) and all(
+        v is not None for v in requires.values()
+    )
+
+    dropped = set(targets)
+    if complete and "" not in root_names:
+        provided = {_canonical_name(n) for n in image_provided}
+        installed = set(dist_infos)
+
+        def closure(start, stop):
+            seen = set()
+            stack = [n for n in start if n in installed]
+            while stack:
+                node = stack.pop()
+                if node in seen:
+                    continue
+                seen.add(node)
+                stack.extend(
+                    d for d in requires[node] if d in installed and d not in stop
+                )
+            return seen
+
+        # a user exclusion removes only the named package, so its
+        # dependencies are protected like the roots
+        pulled = closure(installed & provided, stop=set())
+        seed = (installed - pulled) | (installed & root_names)
+        seed |= installed & (excluded_set - provided)
+        kept = closure(seed - provided, stop=provided)
+        dropped |= installed - kept
+
+    kept_files = set()
+    dropped_files = set()
+    for name, paths in records.items():
+        if paths is None:
+            continue
+        (dropped_files if name in dropped else kept_files).update(paths)
+    dropped_files -= kept_files
+
+    root = env_dir.resolve()
+    touched_tops = set()
+    for rel in dropped_files:
+        path = env_dir / rel
+        try:
+            path.resolve().relative_to(root)
+        except ValueError:
+            continue
+        if path.is_file() or path.is_symlink():
+            path.unlink()
+            touched_tops.add(rel.split("/", 1)[0])
+
+    for name in dropped:
+        shutil.rmtree(dist_infos[name], ignore_errors=True)
+        # wheels without a RECORD: fall back to the import-name directory
+        if records[name] is None:
+            top = env_dir / name.replace("-", "_")
+            if top.is_dir() and not top.is_symlink():
+                shutil.rmtree(top, ignore_errors=True)
+
+    # drop directories the removal emptied
+    for top in touched_tops:
+        top_path = env_dir / top
+        if not top_path.is_dir() or top_path.is_symlink():
+            continue
+        for directory in sorted(
+            (p for p in top_path.rglob("*") if p.is_dir() and not p.is_symlink()),
+            key=lambda p: len(p.parts),
+            reverse=True,
+        ):
+            if not any(directory.iterdir()):
+                directory.rmdir()
+        if not any(top_path.iterdir()):
+            top_path.rmdir()
+
+    return sorted(dropped)
+
+
 def _is_pinnable_name(spec: str) -> bool:
     """true if the spec is a plain requirement (name, optional pin), as
     opposed to a url, git ref, or filesystem path."""
@@ -326,11 +542,16 @@ def build_environment(
 
     user_requirements = collect_requirements(project_root, app)
     # size-prohibitive packages come preinstalled only on the builtin
-    # gpu images. custom images carry no guarantee, and builtin cpu
-    # images ship without them, so every resource must be gpu-on-builtin
-    # for auto-exclusion to be sound (the artifact is shared app-wide)
-    auto_exclude = all(
-        h.spec.gpu and not h.spec.image for h in app.resources.values()
+    # gpu images, and custom images carry no guarantee. builtin cpu
+    # images ship without them, but a vendored pypi torch cannot serve
+    # them either: built off linux it lacks the nvidia libraries and
+    # fails to import, and on linux those libraries exceed the artifact
+    # limit. it would still shadow the gpu image's torch, so an app with
+    # any builtin gpu resource excludes them app-wide (the artifact is
+    # shared), as runpod-flash did
+    resources = list(app.resources.values())
+    auto_exclude = not any(h.spec.image for h in resources) and any(
+        h.spec.gpu for h in resources
     )
     vendored, excluded = split_exclusions(
         user_requirements, exclude, auto_exclude=auto_exclude
@@ -368,6 +589,57 @@ def build_environment(
     )
     vendor(env_dir, all_requirements, python_version, progress=progress)
 
+    # dependencies of vendored packages can pull excluded packages back
+    # in (peft -> torch). a vendored torch shadows the image's build on
+    # the worker and fails its cuda check, so remove them here
+    prune_names = {requirement_name(e) for e in (exclude or [])}
+    if auto_exclude:
+        prune_names |= SIZE_PROHIBITIVE_PACKAGES
+    # workers install manifest-excluded packages their image lacks at
+    # cold start. the gpu images have the torch family, so in a mixed app
+    # only the cpu resources act on it: record it only if they need it
+    cpu_needs = None
+    if auto_exclude and any(not h.spec.gpu for h in resources):
+        cpu_roots = _requirements_txt(project_root)
+        cpu_roots += [
+            dep
+            for h in resources
+            if not h.spec.gpu
+            for dep in h.spec.dependencies or []
+        ]
+        cpu_needs = installed_closure(env_dir, cpu_roots)
+    pruned = prune_excluded(
+        env_dir,
+        sorted(prune_names),
+        all_requirements,
+        sorted(SIZE_PROHIBITIVE_PACKAGES) if auto_exclude else [],
+    )
+    if pruned:
+        log.info(
+            "removed %d dependency packages provided by the worker image "
+            "or excluded: %s",
+            len(pruned),
+            ", ".join(pruned),
+        )
+        excluded = sorted(
+            set(excluded) | ({_canonical_name(n) for n in prune_names} & set(pruned))
+        )
+    if cpu_needs is not None:
+        user_excluded = {_canonical_name(requirement_name(e)) for e in exclude or []}
+        excluded = [
+            name
+            for name in excluded
+            if name in user_excluded
+            or name not in SIZE_PROHIBITIVE_PACKAGES
+            or name in cpu_needs
+        ]
+        for name in sorted(set(excluded) & SIZE_PROHIBITIVE_PACKAGES):
+            log.warning(
+                "cpu resources need %s, which is not bundled; their workers "
+                "install it at cold start, so give them container disk for it",
+                name,
+            )
+
     if not package_spec:
         sync_running_package(env_dir)