From 71b72bd1726cae16df604a025d8b9b962b707c06 Mon Sep 17 00:00:00 2001 From: Justin Date: Fri, 16 Oct 2026 18:33:53 -0400 Subject: [PATCH 1/7] fix(apps): keep transitive torch out of GPU artifacts Packages that depend on torch (peft, accelerate, sentence-transformers) installed it into the artifact even on builtin GPU images, because only top-level requirements were excluded. The vendored PyPI torch shadowed the image's build and failed to load its CUDA libraries, so workers crash-looped. runpod-flash stripped these at archive time. After vendoring, remove every installed excluded distribution, plus the wheels only the torch family depends on (nvidia/cuda, sympy, networkx), using each wheel's RECORD and Requires-Dist. Fails closed to name-only removal when metadata is missing or a requirement has no name. Script paths pip --target records as ../../bin are mapped onto the target. Auto-exclusion now applies when any resource is GPU on a builtin image and none uses a custom image, matching runpod-flash for mixed CPU+GPU apps. A vendored PyPI torch cannot import on the builtin CPU images either, so CPU resources lose nothing. --- runpod/apps/build.py | 226 +++++++++++++++++++++++++++++++++- tests/test_apps/test_build.py | 197 +++++++++++++++++++++++++++++ 2 files changed, 418 insertions(+), 5 deletions(-) diff --git a/runpod/apps/build.py b/runpod/apps/build.py index b438f0ac6..a8e88381e 100644 --- a/runpod/apps/build.py +++ b/runpod/apps/build.py @@ -1,7 +1,9 @@ """deploy-time environment build: vendor dependencies into the artifact.""" +import csv import logging import os +import posixpath import re import shutil import subprocess @@ -9,6 +11,7 @@ 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.""" @@ -135,6 +144,185 @@ 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 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 +514,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 +561,29 @@ 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 + 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 not package_spec: sync_running_package(env_dir) diff --git a/tests/test_apps/test_build.py b/tests/test_apps/test_build.py index b757867e1..22500a7d2 100644 --- a/tests/test_apps/test_build.py +++ b/tests/test_apps/test_build.py @@ -13,6 +13,7 @@ SIZE_PROHIBITIVE_PACKAGES, BuildError, collect_requirements, + prune_excluded, requirement_name, runpod_requirement, split_exclusions, @@ -253,3 +254,199 @@ def stop_blocked_process(): process.wait() assert not expired.is_set() assert seen == [(1, "numpy")] + + +def _wheel(env, name, version="1.0", requires=(), files=(), record=True, metadata=True): + """fake a pip --target install of one wheel into env.""" + dist_info = env / f"{name.replace('-', '_')}-{version}.dist-info" + dist_info.mkdir(parents=True) + if metadata: + lines = [f"Name: {name}", f"Version: {version}"] + lines += [f"Requires-Dist: {r}" for r in requires] + (dist_info / "METADATA").write_text("\n".join(lines) + "\n") + rel_files = list(files) or [f"{name.replace('-', '_')}/__init__.py"] + for rel in rel_files: + target = env / (rel[len("../../"):] if rel.startswith("../../") else rel) + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text("") + if record: + rows = rel_files + [f"{dist_info.name}/METADATA", f"{dist_info.name}/RECORD"] + (dist_info / "RECORD").write_text("".join(f"{r},,\n" for r in rows)) + + +def _installed(env): + return sorted(d.name.split("-")[0] for d in env.glob("*.dist-info")) + + +TORCH = sorted(SIZE_PROHIBITIVE_PACKAGES) + + +class TestPruneExcluded: + def _peft_env(self, env): + _wheel(env, "peft", requires=["torch>=1.13", "numpy"]) + _wheel(env, "numpy") + _wheel( + env, + "torch", + requires=[ + "sympy", + "numpy", + 'nvidia-cublas-cu12==12.8; platform_system == "Linux"', + ], + files=[ + "torch/__init__.py", + "torch/__pycache__/__init__.cpython-312.pyc", + "functorch/__init__.py", + "torchgen/__init__.py", + "../../bin/torchrun", + ], + ) + _wheel(env, "sympy", requires=["mpmath"]) + _wheel(env, "mpmath") + _wheel( + env, + "nvidia-cublas-cu12", + files=["nvidia/cublas/lib/libcublas.so", "nvidia/__init__.py"], + ) + _wheel( + env, + "nvidia-cudnn-cu12", + files=["nvidia/cudnn/lib/libcudnn.so", "nvidia/__init__.py"], + ) + + def test_transitive_torch_and_its_tree_removed(self, tmp_path): + self._peft_env(tmp_path) + _wheel(tmp_path, "faster-whisper", requires=["nvidia-cudnn-cu12"]) + + dropped = prune_excluded( + tmp_path, TORCH, ["peft", "faster-whisper"], TORCH + ) + + assert dropped == ["mpmath", "nvidia-cublas-cu12", "sympy", "torch"] + assert _installed(tmp_path) == [ + "faster_whisper", + "numpy", + "nvidia_cudnn_cu12", + "peft", + ] + for gone in ("torch", "functorch", "torchgen", "sympy", "bin/torchrun"): + assert not (tmp_path / gone).exists(), gone + # shared namespace dir keeps the files a kept wheel installed + assert (tmp_path / "nvidia/cudnn/lib/libcudnn.so").is_file() + assert (tmp_path / "nvidia/__init__.py").is_file() + assert not (tmp_path / "nvidia/cublas").exists() + + def test_root_requirement_is_kept(self, tmp_path): + self._peft_env(tmp_path) + prune_excluded(tmp_path, TORCH, ["peft", "sympy==1.14"], TORCH) + assert "sympy" in _installed(tmp_path) + assert "mpmath" in _installed(tmp_path) + + def test_nothing_excluded_installed_is_noop(self, tmp_path): + _wheel(tmp_path, "numpy") + assert prune_excluded(tmp_path, TORCH, ["numpy"], TORCH) == [] + assert _installed(tmp_path) == ["numpy"] + + def test_missing_record_fails_closed(self, tmp_path): + self._peft_env(tmp_path) + _wheel(tmp_path, "oddball", record=False) + dropped = prune_excluded(tmp_path, TORCH, ["peft", "oddball"], TORCH) + # only the excluded package itself goes + assert dropped == ["torch"] + assert not (tmp_path / "torch").exists() + assert "sympy" in _installed(tmp_path) + + def test_missing_record_on_excluded_falls_back_to_dir(self, tmp_path): + _wheel(tmp_path, "peft", requires=["torch"]) + _wheel(tmp_path, "torch", record=False) + assert prune_excluded(tmp_path, TORCH, ["peft"], TORCH) == ["torch"] + assert not (tmp_path / "torch").exists() + + def test_unnamed_root_fails_closed(self, tmp_path): + self._peft_env(tmp_path) + dropped = prune_excluded( + tmp_path, TORCH, ["peft", "https://example.com/pkg.tar.gz"], TORCH + ) + assert dropped == ["torch"] + + def test_unreadable_requires_dist_fails_closed(self, tmp_path): + self._peft_env(tmp_path) + _wheel(tmp_path, "weird", requires=["https://example.com/x.whl"]) + dropped = prune_excluded(tmp_path, TORCH, ["peft", "weird"], TORCH) + assert dropped == ["torch"] + + def test_user_exclude_keeps_its_dependencies(self, tmp_path): + self._peft_env(tmp_path) + _wheel(tmp_path, "bigpkg", requires=["sympy"]) + dropped = prune_excluded( + tmp_path, TORCH + ["bigpkg"], ["peft", "bigpkg"], TORCH + ) + assert "bigpkg" in dropped + # sympy is torch-only except that bigpkg needs it + assert "sympy" in _installed(tmp_path) + + def test_no_image_provided_removes_only_named(self, tmp_path): + # custom image: a user --exclude of torch drops torch, not its tree + self._peft_env(tmp_path) + assert prune_excluded(tmp_path, ["torch"], ["peft"], []) == ["torch"] + assert "sympy" in _installed(tmp_path) + + +class TestBuildEnvironmentPrunes: + def _app(self, **kwargs): + app = App("a") + + @app.queue(name="q", dependencies=["peft"], **kwargs) + def q(): + pass + + return app + + def _fake_vendor(self, env_dir, requirements, python_version, **kwargs): + if "peft" in requirements: + _wheel(env_dir, "peft", requires=["torch"]) + _wheel(env_dir, "torch", requires=["sympy"]) + _wheel(env_dir, "sympy") + + def test_gpu_builtin_image_prunes_transitive_torch(self, tmp_path): + from runpod.apps.build import build_environment + + with patch("runpod.apps.build.vendor", side_effect=self._fake_vendor): + result = build_environment( + self._app(gpu="AMPERE_16"), tmp_path, scratch_dir=tmp_path / "s" + ) + assert result.excluded == ["torch"] + assert _installed(result.env_dir) == ["peft"] + + def test_mixed_cpu_gpu_app_prunes_transitive_torch(self, tmp_path): + from runpod.apps.build import build_environment + + app = self._app(gpu="AMPERE_16") + + @app.queue(name="prep", cpu="cpu3c-1-2", dependencies=["pandas"]) + def prep(): + pass + + with patch("runpod.apps.build.vendor", side_effect=self._fake_vendor): + result = build_environment(app, tmp_path, scratch_dir=tmp_path / "s") + assert result.excluded == ["torch"] + assert _installed(result.env_dir) == ["peft"] + + def test_custom_image_keeps_transitive_torch(self, tmp_path): + from runpod.apps.build import build_environment + + app = self._app(gpu="AMPERE_16", image="me/custom:1") + with patch("runpod.apps.build.vendor", side_effect=self._fake_vendor): + result = build_environment(app, tmp_path, scratch_dir=tmp_path / "s") + assert result.excluded == [] + assert _installed(result.env_dir) == ["peft", "sympy", "torch"] + + def test_cpu_keeps_transitive_torch(self, tmp_path): + from runpod.apps.build import build_environment + + with patch("runpod.apps.build.vendor", side_effect=self._fake_vendor): + result = build_environment( + self._app(cpu="cpu3c-1-2"), tmp_path, scratch_dir=tmp_path / "s" + ) + assert result.excluded == [] + assert _installed(result.env_dir) == ["peft", "sympy", "torch"] From 9c9166f8052a83cdb8406cb6fe8ed5db69ba4c79 Mon Sep 17 00:00:00 2001 From: Justin Date: Fri, 16 Oct 2026 18:40:00 -0400 Subject: [PATCH 2/7] fix(apps): record torch for cpu workers only when they need it Workers pip-install manifest excludedPackages their image lacks at cold start. In a mixed CPU+GPU app the GPU images already have the torch family, so only CPU workers act on the list, and a CPU stage that never uses torch tried to install CUDA torch onto its default 10 GB disk and failed to boot. In mixed apps, keep a torch-family name in the manifest only when the CPU resources' dependency closure (their decorator dependencies plus requirements.txt) reaches it, and warn that those workers install it at cold start. Unreadable metadata keeps every name, as before. --- runpod/apps/build.py | 62 +++++++++++++++++++++++++++++++++-- tests/test_apps/test_build.py | 35 +++++++++++++++++++- 2 files changed, 93 insertions(+), 4 deletions(-) diff --git a/runpod/apps/build.py b/runpod/apps/build.py index a8e88381e..cf2500320 100644 --- a/runpod/apps/build.py +++ b/runpod/apps/build.py @@ -88,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 []: @@ -212,6 +217,29 @@ def _read_requires(dist_info: Path) -> Optional[set]: 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], @@ -567,6 +595,19 @@ def build_environment( 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), @@ -583,6 +624,21 @@ def build_environment( 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) diff --git a/tests/test_apps/test_build.py b/tests/test_apps/test_build.py index 22500a7d2..a5d6d6ac4 100644 --- a/tests/test_apps/test_build.py +++ b/tests/test_apps/test_build.py @@ -429,9 +429,42 @@ def prep(): with patch("runpod.apps.build.vendor", side_effect=self._fake_vendor): result = build_environment(app, tmp_path, scratch_dir=tmp_path / "s") - assert result.excluded == ["torch"] + # torch comes from the gpu image; the cpu stage never needs it, + # so its workers must not install it at cold start + assert result.excluded == [] assert _installed(result.env_dir) == ["peft"] + def test_mixed_app_records_torch_only_when_cpu_needs_it(self, tmp_path): + from runpod.apps.build import build_environment + + app = self._app(gpu="AMPERE_16") + + @app.queue(name="prep", cpu="cpu3c-1-2", dependencies=["peft"]) + def prep(): + pass + + with patch("runpod.apps.build.vendor", side_effect=self._fake_vendor): + result = build_environment(app, tmp_path, scratch_dir=tmp_path / "s") + # the cpu workers install torch at cold start, so it stays listed + assert result.excluded == ["torch"] + + def test_mixed_app_direct_gpu_torch_not_recorded(self, tmp_path): + from runpod.apps.build import build_environment + + app = App("a") + + @app.queue(name="g", gpu="AMPERE_16", dependencies=["torch"]) + def g(): + pass + + @app.queue(name="c", cpu="cpu3c-1-2", dependencies=["pandas"]) + def c(): + pass + + with patch("runpod.apps.build.vendor"): + result = build_environment(app, tmp_path, scratch_dir=tmp_path / "s") + assert result.excluded == [] + def test_custom_image_keeps_transitive_torch(self, tmp_path): from runpod.apps.build import build_environment From 1889c3f6344f7a97c8a0b90d99af3839fe358520 Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 19 Oct 2026 23:09:08 -0400 Subject: [PATCH 3/7] fix(apps): count torch behind an extra only when the cpu stage asks for it transformers lists torch under its [torch] extra. the cpu-needs closure followed every Requires-Dist line, so a mixed app whose cpu stage used plain transformers recorded torch and its workers installed it at cold start. --- runpod/apps/build.py | 66 +++++++++++++++++++++++++++++------ tests/test_apps/test_build.py | 62 ++++++++++++++++++++++++++++++++ 2 files changed, 117 insertions(+), 11 deletions(-) diff --git a/runpod/apps/build.py b/runpod/apps/build.py index cf2500320..b2da53c98 100644 --- a/runpod/apps/build.py +++ b/runpod/apps/build.py @@ -13,7 +13,7 @@ from collections import deque from email.parser import HeaderParser from pathlib import Path -from typing import List, Optional, Tuple +from typing import Dict, List, Optional, Tuple from .errors import AppError from .images import DEFAULT_PYTHON_VERSION, SUPPORTED_PYTHON_VERSIONS @@ -217,27 +217,71 @@ def _read_requires(dist_info: Path) -> Optional[set]: return deps +_EXTRA_MARKER_RE = re.compile(r"""\bextra\s*==\s*['"]([^'"]+)['"]""") +_REQ_EXTRAS_RE = re.compile(r"\s*[A-Za-z0-9][A-Za-z0-9._-]*\s*\[([^\]]*)\]") + + +def _requirement_extras(requirement: str) -> set: + """canonical extras a requirement asks for ("transformers[torch]").""" + match = _REQ_EXTRAS_RE.match(requirement) + if not match: + return set() + return {_canonical_name(e.strip()) for e in match.group(1).split(",") if e.strip()} + + +def _read_requires_by_extra(dist_info: Path) -> Optional[list]: + """(dependency, its extras, extras that gate it) per Requires-Dist + line, or None when unreadable. an ungated line has no gating extras. + other markers 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") + ) + entries = [] + for line in headers.get_all("Requires-Dist") or []: + dep = _strict_requirement_name(line) + if not dep: + return None + marker = line.split(";", 1)[1] if ";" in line else "" + gates = set() + # gated only when every "or" branch names an extra; any other + # branch (extra == "torch" or python_version >= "3.12") can + # apply without the extra, so the line counts as ungated + branches = re.split(r"\bor\b", marker) + if marker.strip() and all(_EXTRA_MARKER_RE.search(b) for b in branches): + gates = {_canonical_name(e) for e in _EXTRA_MARKER_RE.findall(marker)} + entries.append((dep, _requirement_extras(line.split(";", 1)[0]), gates)) + return entries + + 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.""" + roots included. a dependency behind an extra counts only when that + extra was asked for (transformers does not pull torch, while + transformers[torch] does). unreadable metadata counts every + installed distribution as pulled in.""" requires = { - _canonical_name(_dist_info_name(d.name)): _read_requires(d) + _canonical_name(_dist_info_name(d.name)): _read_requires_by_extra(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) + activated: Dict[str, set] = {} + stack = [(_strict_requirement_name(r), _requirement_extras(r)) for r in roots] while stack: - node = stack.pop() - if node in seen: + node, extras = stack.pop() + known = activated.get(node) + if known is not None and extras <= known: continue - seen.add(node) - stack.extend(requires.get(node) or ()) - return seen + active = activated[node] = (known or set()) | extras + for dep, dep_extras, gates in requires.get(node) or (): + if not gates or gates & active: + stack.append((dep, dep_extras)) + return set(activated) def prune_excluded( diff --git a/tests/test_apps/test_build.py b/tests/test_apps/test_build.py index a5d6d6ac4..42f74e272 100644 --- a/tests/test_apps/test_build.py +++ b/tests/test_apps/test_build.py @@ -448,6 +448,68 @@ def prep(): # the cpu workers install torch at cold start, so it stays listed assert result.excluded == ["torch"] + def test_mixed_app_ignores_torch_behind_an_unrequested_extra(self, tmp_path): + from runpod.apps.build import build_environment + + app = self._app(gpu="AMPERE_16") + + @app.queue(name="prep", cpu="cpu3c-1-2", dependencies=["transformers"]) + def prep(): + pass + + def vendor(env_dir, requirements, python_version, **kwargs): + self._fake_vendor(env_dir, requirements, python_version) + _wheel(env_dir, "transformers", requires=['torch>=2; extra == "torch"']) + + with patch("runpod.apps.build.vendor", side_effect=vendor): + result = build_environment(app, tmp_path, scratch_dir=tmp_path / "s") + # transformers without [torch] never imports it on the cpu stage + assert result.excluded == [] + + def test_mixed_app_records_torch_for_a_requested_extra(self, tmp_path): + from runpod.apps.build import build_environment + + app = self._app(gpu="AMPERE_16") + + @app.queue(name="prep", cpu="cpu3c-1-2", dependencies=["transformers[torch]"]) + def prep(): + pass + + def vendor(env_dir, requirements, python_version, **kwargs): + self._fake_vendor(env_dir, requirements, python_version) + _wheel( + env_dir, + "transformers", + requires=['torch>=2; extra == "torch" or extra == "all"'], + ) + + with patch("runpod.apps.build.vendor", side_effect=vendor): + result = build_environment(app, tmp_path, scratch_dir=tmp_path / "s") + assert result.excluded == ["torch"] + + def test_mixed_app_records_torch_when_a_marker_branch_skips_the_extra( + self, tmp_path + ): + from runpod.apps.build import build_environment + + app = self._app(gpu="AMPERE_16") + + @app.queue(name="prep", cpu="cpu3c-1-2", dependencies=["transformers"]) + def prep(): + pass + + def vendor(env_dir, requirements, python_version, **kwargs): + self._fake_vendor(env_dir, requirements, python_version) + _wheel( + env_dir, + "transformers", + requires=['torch; extra == "torch" or python_version >= "3.12"'], + ) + + with patch("runpod.apps.build.vendor", side_effect=vendor): + result = build_environment(app, tmp_path, scratch_dir=tmp_path / "s") + assert result.excluded == ["torch"] + def test_mixed_app_direct_gpu_torch_not_recorded(self, tmp_path): from runpod.apps.build import build_environment From ce9fc88f00e1f9676b5ca999cdc18e11aaec9ff2 Mon Sep 17 00:00:00 2001 From: Justin Date: Tue, 20 Oct 2026 12:28:31 -0400 Subject: [PATCH 4/7] refactor(apps): split prune and metadata helpers out of the build --- runpod/apps/build.py | 280 ++++++++++++++++++---------------- tests/test_apps/test_build.py | 50 +++--- 2 files changed, 171 insertions(+), 159 deletions(-) diff --git a/runpod/apps/build.py b/runpod/apps/build.py index b2da53c98..23eac2aec 100644 --- a/runpod/apps/build.py +++ b/runpod/apps/build.py @@ -67,6 +67,8 @@ _STRICT_REQ_NAME_RE = re.compile( r"\s*([A-Za-z0-9][A-Za-z0-9._-]*)\s*(?:\[[^\]]*\])?\s*(?:$|[<>=!~;@(])" ) +_REQ_EXTRAS_RE = re.compile(r"\s*[A-Za-z0-9][A-Za-z0-9._-]*\s*\[([^\]]*)\]") +_EXTRA_MARKER_RE = re.compile(r"""\bextra\s*==\s*['"]([^'"]+)['"]""") class BuildError(AppError): @@ -197,28 +199,35 @@ def _read_record(dist_info: Path) -> Optional[set]: return paths -def _read_requires(dist_info: Path) -> Optional[set]: - """canonical names a wheel depends on, or None when unreadable. +def _installed_dist_infos(env_dir: Path) -> Dict[str, Path]: + """canonical distribution name -> its .dist-info dir in env_dir.""" + return { + _canonical_name(_dist_info_name(d.name)): d + for d in env_dir.glob("*.dist-info") + if d.is_dir() + } - markers and extras are ignored, which can only keep more wheels. - """ + +def _requires_dist(dist_info: Path) -> Optional[List[str]]: + """a wheel's Requires-Dist lines, or None when METADATA is missing or + a line does not start with a name.""" 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 + lines = headers.get_all("Requires-Dist") or [] + return None if any(not _strict_requirement_name(line) for line in lines) else lines -_EXTRA_MARKER_RE = re.compile(r"""\bextra\s*==\s*['"]([^'"]+)['"]""") -_REQ_EXTRAS_RE = re.compile(r"\s*[A-Za-z0-9][A-Za-z0-9._-]*\s*\[([^\]]*)\]") +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. + """ + lines = _requires_dist(dist_info) + return None if lines is None else {_strict_requirement_name(line) for line in lines} def _requirement_extras(requirement: str) -> set: @@ -229,30 +238,34 @@ def _requirement_extras(requirement: str) -> set: return {_canonical_name(e.strip()) for e in match.group(1).split(",") if e.strip()} +def _gating_extras(marker: str) -> set: + """extras a marker requires, or an empty set when it can apply + without one. every "or" branch must name an extra: in + 'extra == "torch" or python_version >= "3.12"' the second branch + applies on its own, so nothing gates the line.""" + branches = re.split(r"\bor\b", marker) + if not marker.strip() or not all(_EXTRA_MARKER_RE.search(b) for b in branches): + return set() + return {_canonical_name(e) for e in _EXTRA_MARKER_RE.findall(marker)} + + def _read_requires_by_extra(dist_info: Path) -> Optional[list]: """(dependency, its extras, extras that gate it) per Requires-Dist - line, or None when unreadable. an ungated line has no gating extras. - other markers are ignored, which can only keep more wheels.""" - metadata = dist_info / "METADATA" - if not metadata.is_file(): + line, or None when unreadable. other markers are ignored, which can + only keep more wheels.""" + lines = _requires_dist(dist_info) + if lines is None: return None - headers = HeaderParser().parsestr( - metadata.read_text(encoding="utf-8", errors="replace") - ) entries = [] - for line in headers.get_all("Requires-Dist") or []: - dep = _strict_requirement_name(line) - if not dep: - return None - marker = line.split(";", 1)[1] if ";" in line else "" - gates = set() - # gated only when every "or" branch names an extra; any other - # branch (extra == "torch" or python_version >= "3.12") can - # apply without the extra, so the line counts as ungated - branches = re.split(r"\bor\b", marker) - if marker.strip() and all(_EXTRA_MARKER_RE.search(b) for b in branches): - gates = {_canonical_name(e) for e in _EXTRA_MARKER_RE.findall(marker)} - entries.append((dep, _requirement_extras(line.split(";", 1)[0]), gates)) + for line in lines: + requirement, _, marker = line.partition(";") + entries.append( + ( + _strict_requirement_name(line), + _requirement_extras(requirement), + _gating_extras(marker), + ) + ) return entries @@ -263,15 +276,14 @@ def installed_closure(env_dir: Path, roots: List[str]) -> set: transformers[torch] does). unreadable metadata counts every installed distribution as pulled in.""" requires = { - _canonical_name(_dist_info_name(d.name)): _read_requires_by_extra(d) - for d in env_dir.glob("*.dist-info") - if d.is_dir() + name: _read_requires_by_extra(d) + for name, d in _installed_dist_infos(env_dir).items() } - names = {_strict_requirement_name(r) for r in roots} + stack = [(_strict_requirement_name(r), _requirement_extras(r)) for r in roots] + names = {name for name, _ in stack} if "" in names or any(v is None for v in requires.values()): return set(requires) | (names - {""}) activated: Dict[str, set] = {} - stack = [(_strict_requirement_name(r), _requirement_extras(r)) for r in roots] while stack: node, extras = stack.pop() known = activated.get(node) @@ -284,6 +296,49 @@ def installed_closure(env_dir: Path, roots: List[str]) -> set: return set(activated) +def _dependency_closure(requires: Dict[str, set], start: set, stop: set) -> set: + """installed distributions reachable from start, not walking into stop.""" + seen = set() + stack = [n for n in start if n in requires] + while stack: + node = stack.pop() + if node in seen: + continue + seen.add(node) + stack.extend(d for d in requires[node] if d in requires and d not in stop) + return seen + + +def _remove_files(env_dir: Path, rel_paths: set) -> set: + """delete files under env_dir, returning the top-level entries touched. + paths that resolve outside env_dir are skipped.""" + root = env_dir.resolve() + touched = set() + for rel in rel_paths: + path = env_dir / rel + if not path.resolve().is_relative_to(root): + continue + if path.is_file() or path.is_symlink(): + path.unlink() + touched.add(rel.split("/", 1)[0]) + return touched + + +def _remove_empty_dirs(env_dir: Path, tops: set) -> None: + """remove directories under each top that the deletions emptied.""" + for top in tops: + top_path = env_dir / top + if not top_path.is_dir() or top_path.is_symlink(): + continue + subdirs = [p for p in top_path.rglob("*") if p.is_dir() and not p.is_symlink()] + # deepest first, so a parent is empty by the time it is checked + for directory in sorted(subdirs, key=lambda p: len(p.parts), reverse=True): + if not any(directory.iterdir()): + directory.rmdir() + if not any(top_path.iterdir()): + top_path.rmdir() + + def prune_excluded( env_dir: Path, excluded: List[str], @@ -306,92 +361,43 @@ def prune_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() - } + dist_infos = _installed_dist_infos(env_dir) excluded_set = {_canonical_name(n) for n in excluded} - targets = set(dist_infos) & excluded_set - if not targets: + dropped = set(dist_infos) & excluded_set + if not dropped: 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() - ) + complete = None not in records.values() and None not 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) + provided = {_canonical_name(n) for n in image_provided} + pulled_by_provided = _dependency_closure(requires, installed & provided, set()) + # keep everything the roots, unrelated wheels, and user exclusions + # (which remove only the named package) need, without walking + # into image-provided packages + seed = (installed - pulled_by_provided) | (installed & root_names) seed |= installed & (excluded_set - provided) - kept = closure(seed - provided, stop=provided) + kept = _dependency_closure(requires, seed - provided, stop=provided) dropped |= installed - kept - kept_files = set() dropped_files = set() + kept_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]) + (dropped_files if name in dropped else kept_files).update(paths or ()) + touched = _remove_files(env_dir, dropped_files - kept_files) 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() + # a wheel without a RECORD: fall back to its import-name directory + fallback = env_dir / name.replace("-", "_") + if records[name] is None and fallback.is_dir() and not fallback.is_symlink(): + shutil.rmtree(fallback, ignore_errors=True) + _remove_empty_dirs(env_dir, touched) return sorted(dropped) @@ -565,6 +571,41 @@ def vendor( ) +def _cpu_requirements_closure(project_root: Path, resources, env_dir: Path) -> set: + """installed distributions the cpu resources' requirements pull in.""" + roots = _requirements_txt(project_root) + roots += [ + dep for h in resources if not h.spec.gpu for dep in h.spec.dependencies or [] + ] + return installed_closure(env_dir, roots) + + +def _exclusions_for_cpu_workers( + excluded: List[str], user_exclude: List[str], cpu_needs: set +) -> List[str]: + """the manifest exclusions for a mixed cpu and gpu app. + + workers pip-install manifest-excluded packages their image lacks at + cold start. the gpu images ship the torch family, so only cpu workers + act on those entries: keep one only when the cpu resources need it. + """ + user_excluded = {_canonical_name(requirement_name(e)) for e in user_exclude} + kept = [ + 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(kept) & 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, + ) + return kept + + def build_environment( app, project_root: Path, @@ -639,19 +680,9 @@ def build_environment( 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) + cpu_needs = _cpu_requirements_closure(project_root, resources, env_dir) pruned = prune_excluded( env_dir, sorted(prune_names), @@ -669,20 +700,7 @@ def build_environment( 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, - ) + excluded = _exclusions_for_cpu_workers(excluded, exclude or [], cpu_needs) if not package_spec: sync_running_package(env_dir) diff --git a/tests/test_apps/test_build.py b/tests/test_apps/test_build.py index 42f74e272..56a80618c 100644 --- a/tests/test_apps/test_build.py +++ b/tests/test_apps/test_build.py @@ -12,6 +12,7 @@ from runpod.apps.build import ( SIZE_PROHIBITIVE_PACKAGES, BuildError, + build_environment, collect_requirements, prune_excluded, requirement_name, @@ -266,7 +267,7 @@ def _wheel(env, name, version="1.0", requires=(), files=(), record=True, metadat (dist_info / "METADATA").write_text("\n".join(lines) + "\n") rel_files = list(files) or [f"{name.replace('-', '_')}/__init__.py"] for rel in rel_files: - target = env / (rel[len("../../"):] if rel.startswith("../../") else rel) + target = env / (rel[len("../../") :] if rel.startswith("../../") else rel) target.parent.mkdir(parents=True, exist_ok=True) target.write_text("") if record: @@ -278,7 +279,7 @@ def _installed(env): return sorted(d.name.split("-")[0] for d in env.glob("*.dist-info")) -TORCH = sorted(SIZE_PROHIBITIVE_PACKAGES) +IMAGE_PROVIDED = sorted(SIZE_PROHIBITIVE_PACKAGES) class TestPruneExcluded: @@ -319,7 +320,7 @@ def test_transitive_torch_and_its_tree_removed(self, tmp_path): _wheel(tmp_path, "faster-whisper", requires=["nvidia-cudnn-cu12"]) dropped = prune_excluded( - tmp_path, TORCH, ["peft", "faster-whisper"], TORCH + tmp_path, IMAGE_PROVIDED, ["peft", "faster-whisper"], IMAGE_PROVIDED ) assert dropped == ["mpmath", "nvidia-cublas-cu12", "sympy", "torch"] @@ -338,19 +339,23 @@ def test_transitive_torch_and_its_tree_removed(self, tmp_path): def test_root_requirement_is_kept(self, tmp_path): self._peft_env(tmp_path) - prune_excluded(tmp_path, TORCH, ["peft", "sympy==1.14"], TORCH) + prune_excluded( + tmp_path, IMAGE_PROVIDED, ["peft", "sympy==1.14"], IMAGE_PROVIDED + ) assert "sympy" in _installed(tmp_path) assert "mpmath" in _installed(tmp_path) def test_nothing_excluded_installed_is_noop(self, tmp_path): _wheel(tmp_path, "numpy") - assert prune_excluded(tmp_path, TORCH, ["numpy"], TORCH) == [] + assert prune_excluded(tmp_path, IMAGE_PROVIDED, ["numpy"], IMAGE_PROVIDED) == [] assert _installed(tmp_path) == ["numpy"] def test_missing_record_fails_closed(self, tmp_path): self._peft_env(tmp_path) _wheel(tmp_path, "oddball", record=False) - dropped = prune_excluded(tmp_path, TORCH, ["peft", "oddball"], TORCH) + dropped = prune_excluded( + tmp_path, IMAGE_PROVIDED, ["peft", "oddball"], IMAGE_PROVIDED + ) # only the excluded package itself goes assert dropped == ["torch"] assert not (tmp_path / "torch").exists() @@ -359,27 +364,34 @@ def test_missing_record_fails_closed(self, tmp_path): def test_missing_record_on_excluded_falls_back_to_dir(self, tmp_path): _wheel(tmp_path, "peft", requires=["torch"]) _wheel(tmp_path, "torch", record=False) - assert prune_excluded(tmp_path, TORCH, ["peft"], TORCH) == ["torch"] + assert prune_excluded(tmp_path, IMAGE_PROVIDED, ["peft"], IMAGE_PROVIDED) == [ + "torch" + ] assert not (tmp_path / "torch").exists() def test_unnamed_root_fails_closed(self, tmp_path): self._peft_env(tmp_path) dropped = prune_excluded( - tmp_path, TORCH, ["peft", "https://example.com/pkg.tar.gz"], TORCH + tmp_path, + IMAGE_PROVIDED, + ["peft", "https://example.com/pkg.tar.gz"], + IMAGE_PROVIDED, ) assert dropped == ["torch"] def test_unreadable_requires_dist_fails_closed(self, tmp_path): self._peft_env(tmp_path) _wheel(tmp_path, "weird", requires=["https://example.com/x.whl"]) - dropped = prune_excluded(tmp_path, TORCH, ["peft", "weird"], TORCH) + dropped = prune_excluded( + tmp_path, IMAGE_PROVIDED, ["peft", "weird"], IMAGE_PROVIDED + ) assert dropped == ["torch"] def test_user_exclude_keeps_its_dependencies(self, tmp_path): self._peft_env(tmp_path) _wheel(tmp_path, "bigpkg", requires=["sympy"]) dropped = prune_excluded( - tmp_path, TORCH + ["bigpkg"], ["peft", "bigpkg"], TORCH + tmp_path, IMAGE_PROVIDED + ["bigpkg"], ["peft", "bigpkg"], IMAGE_PROVIDED ) assert "bigpkg" in dropped # sympy is torch-only except that bigpkg needs it @@ -409,8 +421,6 @@ def _fake_vendor(self, env_dir, requirements, python_version, **kwargs): _wheel(env_dir, "sympy") def test_gpu_builtin_image_prunes_transitive_torch(self, tmp_path): - from runpod.apps.build import build_environment - with patch("runpod.apps.build.vendor", side_effect=self._fake_vendor): result = build_environment( self._app(gpu="AMPERE_16"), tmp_path, scratch_dir=tmp_path / "s" @@ -419,8 +429,6 @@ def test_gpu_builtin_image_prunes_transitive_torch(self, tmp_path): assert _installed(result.env_dir) == ["peft"] def test_mixed_cpu_gpu_app_prunes_transitive_torch(self, tmp_path): - from runpod.apps.build import build_environment - app = self._app(gpu="AMPERE_16") @app.queue(name="prep", cpu="cpu3c-1-2", dependencies=["pandas"]) @@ -435,8 +443,6 @@ def prep(): assert _installed(result.env_dir) == ["peft"] def test_mixed_app_records_torch_only_when_cpu_needs_it(self, tmp_path): - from runpod.apps.build import build_environment - app = self._app(gpu="AMPERE_16") @app.queue(name="prep", cpu="cpu3c-1-2", dependencies=["peft"]) @@ -449,8 +455,6 @@ def prep(): assert result.excluded == ["torch"] def test_mixed_app_ignores_torch_behind_an_unrequested_extra(self, tmp_path): - from runpod.apps.build import build_environment - app = self._app(gpu="AMPERE_16") @app.queue(name="prep", cpu="cpu3c-1-2", dependencies=["transformers"]) @@ -467,8 +471,6 @@ def vendor(env_dir, requirements, python_version, **kwargs): assert result.excluded == [] def test_mixed_app_records_torch_for_a_requested_extra(self, tmp_path): - from runpod.apps.build import build_environment - app = self._app(gpu="AMPERE_16") @app.queue(name="prep", cpu="cpu3c-1-2", dependencies=["transformers[torch]"]) @@ -490,8 +492,6 @@ def vendor(env_dir, requirements, python_version, **kwargs): def test_mixed_app_records_torch_when_a_marker_branch_skips_the_extra( self, tmp_path ): - from runpod.apps.build import build_environment - app = self._app(gpu="AMPERE_16") @app.queue(name="prep", cpu="cpu3c-1-2", dependencies=["transformers"]) @@ -511,8 +511,6 @@ def vendor(env_dir, requirements, python_version, **kwargs): assert result.excluded == ["torch"] def test_mixed_app_direct_gpu_torch_not_recorded(self, tmp_path): - from runpod.apps.build import build_environment - app = App("a") @app.queue(name="g", gpu="AMPERE_16", dependencies=["torch"]) @@ -528,8 +526,6 @@ def c(): assert result.excluded == [] def test_custom_image_keeps_transitive_torch(self, tmp_path): - from runpod.apps.build import build_environment - app = self._app(gpu="AMPERE_16", image="me/custom:1") with patch("runpod.apps.build.vendor", side_effect=self._fake_vendor): result = build_environment(app, tmp_path, scratch_dir=tmp_path / "s") @@ -537,8 +533,6 @@ def test_custom_image_keeps_transitive_torch(self, tmp_path): assert _installed(result.env_dir) == ["peft", "sympy", "torch"] def test_cpu_keeps_transitive_torch(self, tmp_path): - from runpod.apps.build import build_environment - with patch("runpod.apps.build.vendor", side_effect=self._fake_vendor): result = build_environment( self._app(cpu="cpu3c-1-2"), tmp_path, scratch_dir=tmp_path / "s" From 936d6efe992ce391f97ed7d928c459c3817bf431 Mon Sep 17 00:00:00 2001 From: Justin Date: Tue, 20 Oct 2026 13:00:22 -0400 Subject: [PATCH 5/7] refactor(apps): split pruning into choosing and deleting distributions --- runpod/apps/build.py | 72 ++++++++++++++++++++++++++++++++------------ 1 file changed, 52 insertions(+), 20 deletions(-) diff --git a/runpod/apps/build.py b/runpod/apps/build.py index 23eac2aec..e8449f62d 100644 --- a/runpod/apps/build.py +++ b/runpod/apps/build.py @@ -269,6 +269,16 @@ def _read_requires_by_extra(dist_info: Path) -> Optional[list]: return entries +def _active_requirements(entries: list, extras: set) -> list: + """(dependency, its extras) for each Requires-Dist entry that an install + with these extras pulls in. non-extra markers are ignored.""" + return [ + (dep, dep_extras) + for dep, dep_extras, gates in entries + if not gates or gates & extras + ] + + def installed_closure(env_dir: Path, roots: List[str]) -> set: """canonical names of installed distributions the roots pull in, roots included. a dependency behind an extra counts only when that @@ -290,9 +300,7 @@ def installed_closure(env_dir: Path, roots: List[str]) -> set: if known is not None and extras <= known: continue active = activated[node] = (known or set()) | extras - for dep, dep_extras, gates in requires.get(node) or (): - if not gates or gates & active: - stack.append((dep, dep_extras)) + stack.extend(_active_requirements(requires.get(node) or (), active)) return set(activated) @@ -363,27 +371,52 @@ def prune_excluded( """ dist_infos = _installed_dist_infos(env_dir) excluded_set = {_canonical_name(n) for n in excluded} - dropped = set(dist_infos) & excluded_set - if not dropped: + if not set(dist_infos) & excluded_set: return [] - records = {name: _read_record(d) for name, d in dist_infos.items()} + dropped = _distributions_to_drop( + dist_infos, records, excluded_set, roots, image_provided + ) + _delete_distributions(env_dir, dist_infos, records, dropped) + return sorted(dropped) + + +def _distributions_to_drop( + dist_infos: Dict[str, Path], + records: Dict[str, Optional[set]], + excluded: set, + roots: List[str], + image_provided: List[str], +) -> set: + """the excluded distributions, plus the wheels only image-provided + packages need. just the excluded ones when metadata is incomplete + or a root has no name.""" + installed = set(dist_infos) + targets = installed & excluded requires = {name: _read_requires(d) for name, d in dist_infos.items()} root_names = {_strict_requirement_name(r) for r in roots} - complete = None not in records.values() and None not in requires.values() - - if complete and "" not in root_names: - installed = set(dist_infos) - provided = {_canonical_name(n) for n in image_provided} - pulled_by_provided = _dependency_closure(requires, installed & provided, set()) - # keep everything the roots, unrelated wheels, and user exclusions - # (which remove only the named package) need, without walking - # into image-provided packages - seed = (installed - pulled_by_provided) | (installed & root_names) - seed |= installed & (excluded_set - provided) - kept = _dependency_closure(requires, seed - provided, stop=provided) - dropped |= installed - kept + if None in records.values() or None in requires.values() or "" in root_names: + return targets + + provided = {_canonical_name(n) for n in image_provided} + pulled_by_provided = _dependency_closure(requires, installed & provided, set()) + # keep what the roots, unrelated wheels, and user exclusions (which + # remove only the named package) need, without walking into + # image-provided packages + seed = (installed - pulled_by_provided) | (installed & root_names) + seed |= installed & (excluded - provided) + kept = _dependency_closure(requires, seed - provided, stop=provided) + return targets | (installed - kept) + +def _delete_distributions( + env_dir: Path, + dist_infos: Dict[str, Path], + records: Dict[str, Optional[set]], + dropped: set, +) -> None: + """delete the dropped distributions' files, keeping any file a kept + distribution also installed.""" dropped_files = set() kept_files = set() for name, paths in records.items(): @@ -398,7 +431,6 @@ def prune_excluded( shutil.rmtree(fallback, ignore_errors=True) _remove_empty_dirs(env_dir, touched) - return sorted(dropped) def _is_pinnable_name(spec: str) -> bool: From a50ca33e8f6fc9cbb55aad9b5b9342baba75e87e Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Tue, 20 Oct 2026 13:04:34 -0700 Subject: [PATCH 6/7] fix(apps): respect grouped dependency extra markers --- requirements.txt | 1 + runpod/apps/build.py | 43 +++++++++++++------- tests/test_apps/test_build.py | 74 +++++++++++++++++++++++++++++++++++ 3 files changed, 104 insertions(+), 14 deletions(-) diff --git a/requirements.txt b/requirements.txt index 22eef2d7a..e8fd5fa68 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,6 +8,7 @@ cloudpickle >= 3.1.2 colorama >= 0.4.6, < 0.4.7 cryptography >= 50.0.2 fastapi[all] >= 0.143.0 +packaging >= 24.2 paramiko >= 5.0.0 pathspec >= 1.1.1 prettytable >= 3.18.0 diff --git a/runpod/apps/build.py b/runpod/apps/build.py index e8449f62d..6e63241fb 100644 --- a/runpod/apps/build.py +++ b/runpod/apps/build.py @@ -15,6 +15,8 @@ from pathlib import Path from typing import Dict, List, Optional, Tuple +from packaging.markers import InvalidMarker, Marker + from .errors import AppError from .images import DEFAULT_PYTHON_VERSION, SUPPORTED_PYTHON_VERSIONS @@ -68,7 +70,6 @@ r"\s*([A-Za-z0-9][A-Za-z0-9._-]*)\s*(?:\[[^\]]*\])?\s*(?:$|[<>=!~;@(])" ) _REQ_EXTRAS_RE = re.compile(r"\s*[A-Za-z0-9][A-Za-z0-9._-]*\s*\[([^\]]*)\]") -_EXTRA_MARKER_RE = re.compile(r"""\bextra\s*==\s*['"]([^'"]+)['"]""") class BuildError(AppError): @@ -238,19 +239,28 @@ def _requirement_extras(requirement: str) -> set: return {_canonical_name(e.strip()) for e in match.group(1).split(",") if e.strip()} -def _gating_extras(marker: str) -> set: - """extras a marker requires, or an empty set when it can apply - without one. every "or" branch must name an extra: in - 'extra == "torch" or python_version >= "3.12"' the second branch - applies on its own, so nothing gates the line.""" - branches = re.split(r"\bor\b", marker) - if not marker.strip() or not all(_EXTRA_MARKER_RE.search(b) for b in branches): - return set() - return {_canonical_name(e) for e in _EXTRA_MARKER_RE.findall(marker)} +def _extra_marker_applies(markers: list, extra: str) -> bool: + """evaluate a parsed marker's extra conditions, preserving grouping. + + non-extra comparisons are conservatively true: the builder's platform + and python patch version do not describe the worker environment. + """ + groups = [True] + for marker in markers: + if isinstance(marker, list): + groups[-1] &= _extra_marker_applies(marker, extra) + elif isinstance(marker, tuple): + lhs, _, rhs = marker + if lhs.serialize() == "extra" or rhs.serialize() == "extra": + comparison = Marker(" ".join(part.serialize() for part in marker)) + groups[-1] &= comparison.evaluate({"extra": extra}) + elif marker == "or": + groups.append(True) + return any(groups) def _read_requires_by_extra(dist_info: Path) -> Optional[list]: - """(dependency, its extras, extras that gate it) per Requires-Dist + """(dependency, its extras, parsed marker) per Requires-Dist line, or None when unreadable. other markers are ignored, which can only keep more wheels.""" lines = _requires_dist(dist_info) @@ -259,11 +269,15 @@ def _read_requires_by_extra(dist_info: Path) -> Optional[list]: entries = [] for line in lines: requirement, _, marker = line.partition(";") + try: + parsed = Marker(marker) if marker.strip() else None + except InvalidMarker: + return None entries.append( ( _strict_requirement_name(line), _requirement_extras(requirement), - _gating_extras(marker), + parsed, ) ) return entries @@ -274,8 +288,9 @@ def _active_requirements(entries: list, extras: set) -> list: with these extras pulls in. non-extra markers are ignored.""" return [ (dep, dep_extras) - for dep, dep_extras, gates in entries - if not gates or gates & extras + for dep, dep_extras, marker in entries + if marker is None + or any(_extra_marker_applies(marker._markers, extra) for extra in {""} | extras) ] diff --git a/tests/test_apps/test_build.py b/tests/test_apps/test_build.py index 56a80618c..8750b0f44 100644 --- a/tests/test_apps/test_build.py +++ b/tests/test_apps/test_build.py @@ -6,6 +6,7 @@ from unittest.mock import patch import pytest +from packaging.markers import default_environment from runpod.apps import App from runpod.apps.app import _clear_registry @@ -14,6 +15,7 @@ BuildError, build_environment, collect_requirements, + installed_closure, prune_excluded, requirement_name, runpod_requirement, @@ -489,6 +491,44 @@ def vendor(env_dir, requirements, python_version, **kwargs): result = build_environment(app, tmp_path, scratch_dir=tmp_path / "s") assert result.excluded == ["torch"] + @pytest.mark.parametrize( + "marker", + [ + '(sys_platform == "linux" or sys_platform == "win32") and extra == "torch"', + 'extra == "torch" and (sys_platform == "linux" or sys_platform == "win32")', + '"torch" == extra', + 'extra == "torch" and python_version < "3.12"', + ], + ) + @pytest.mark.parametrize("requested", [False, True]) + @pytest.mark.parametrize("builder_platform", ["darwin", "linux"]) + def test_mixed_app_evaluates_extra_markers( + self, tmp_path, marker, requested, builder_platform + ): + app = self._app(gpu="AMPERE_16") + + @app.queue( + name="prep", + cpu="cpu3c-1-2", + dependencies=["transformers[torch]" if requested else "transformers"], + ) + def prep(): + pass + + def vendor(env_dir, requirements, python_version, **kwargs): + self._fake_vendor(env_dir, requirements, python_version) + _wheel(env_dir, "transformers", requires=[f"torch; {marker}"]) + + builder = default_environment() + builder.update(sys_platform=builder_platform, python_version="3.14") + with ( + patch("runpod.apps.build.vendor", side_effect=vendor), + patch("packaging.markers.default_environment", return_value=builder), + ): + result = build_environment(app, tmp_path, scratch_dir=tmp_path / "s") + + assert result.excluded == (["torch"] if requested else []) + def test_mixed_app_records_torch_when_a_marker_branch_skips_the_extra( self, tmp_path ): @@ -539,3 +579,37 @@ def test_cpu_keeps_transitive_torch(self, tmp_path): ) assert result.excluded == [] assert _installed(result.env_dir) == ["peft", "sympy", "torch"] + + +class TestInstalledClosure: + def test_transitive_extras_activate_grouped_marker(self, tmp_path): + _wheel(tmp_path, "wrapper", requires=["models[torch]"]) + _wheel( + tmp_path, + "models", + requires=[ + 'torch; (sys_platform == "linux" or sys_platform == "win32") and extra == "torch"' + ], + ) + _wheel(tmp_path, "torch") + assert installed_closure(tmp_path, ["wrapper"]) == {"wrapper", "models", "torch"} + + def test_each_extra_is_a_separate_marker_context(self, tmp_path): + _wheel(tmp_path, "models", requires=['torch; extra == "a" and extra == "b"']) + _wheel(tmp_path, "torch") + assert installed_closure(tmp_path, ["models[a,b]"]) == {"models"} + + @pytest.mark.parametrize("requested", [False, True]) + def test_or_in_extra_name_is_not_an_operator(self, tmp_path, requested): + _wheel(tmp_path, "models", requires=['torch; extra == "torch-or-gpu"']) + _wheel(tmp_path, "torch") + roots = ["models[torch-or-gpu]" if requested else "models"] + assert installed_closure(tmp_path, roots) == ( + {"models", "torch"} if requested else {"models"} + ) + + def test_invalid_marker_fails_closed(self, tmp_path): + _wheel(tmp_path, "models", requires=['torch; extra ==']) + _wheel(tmp_path, "torch") + _wheel(tmp_path, "other") + assert installed_closure(tmp_path, ["models"]) == {"models", "torch", "other"} From 00d85cc8c2315385d72b0a8d589e577e1ce2274e Mon Sep 17 00:00:00 2001 From: zeke <40004347+KAJdev@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:16:33 -0700 Subject: [PATCH 7/7] refactor(apps): reuse parsed package requirements --- runpod/apps/build.py | 137 +++++++++++++--------------------- tests/test_apps/test_build.py | 31 ++++++++ 2 files changed, 81 insertions(+), 87 deletions(-) diff --git a/runpod/apps/build.py b/runpod/apps/build.py index 6e63241fb..efb91dac6 100644 --- a/runpod/apps/build.py +++ b/runpod/apps/build.py @@ -15,7 +15,9 @@ from pathlib import Path from typing import Dict, List, Optional, Tuple -from packaging.markers import InvalidMarker, Marker +from packaging.markers import Marker +from packaging.requirements import InvalidRequirement, Requirement +from packaging.utils import canonicalize_name from .errors import AppError from .images import DEFAULT_PYTHON_VERSION, SUPPORTED_PYTHON_VERSIONS @@ -64,13 +66,6 @@ _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*(?:$|[<>=!~;@(])" -) -_REQ_EXTRAS_RE = re.compile(r"\s*[A-Za-z0-9][A-Za-z0-9._-]*\s*\[([^\]]*)\]") - class BuildError(AppError): """environment vendoring failed.""" @@ -152,22 +147,17 @@ 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 +def _parse_requirement(requirement: str) -> Optional[Requirement]: + """parse a named requirement, or None for unnamed or invalid inputs.""" + try: + parsed = Requirement(requirement) + except InvalidRequirement: + return None + if requirement.strip() == parsed.name and re.search( + r"\.(whl|zip|tar\.gz|tgz)$", parsed.name, re.I ): - return "" - return _canonical_name(match.group(1)) + return None + return parsed def _dist_info_name(dist_info_dir: str) -> str: @@ -203,40 +193,24 @@ def _read_record(dist_info: Path) -> Optional[set]: def _installed_dist_infos(env_dir: Path) -> Dict[str, Path]: """canonical distribution name -> its .dist-info dir in env_dir.""" return { - _canonical_name(_dist_info_name(d.name)): d + canonicalize_name(_dist_info_name(d.name)): d for d in env_dir.glob("*.dist-info") if d.is_dir() } -def _requires_dist(dist_info: Path) -> Optional[List[str]]: - """a wheel's Requires-Dist lines, or None when METADATA is missing or - a line does not start with a name.""" +def _requires_dist(dist_info: Path) -> Optional[List[Requirement]]: + """a wheel's requirements, or None for missing or invalid metadata.""" metadata = dist_info / "METADATA" if not metadata.is_file(): return None headers = HeaderParser().parsestr( metadata.read_text(encoding="utf-8", errors="replace") ) - lines = headers.get_all("Requires-Dist") or [] - return None if any(not _strict_requirement_name(line) for line in lines) else lines - - -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. - """ - lines = _requires_dist(dist_info) - return None if lines is None else {_strict_requirement_name(line) for line in lines} - - -def _requirement_extras(requirement: str) -> set: - """canonical extras a requirement asks for ("transformers[torch]").""" - match = _REQ_EXTRAS_RE.match(requirement) - if not match: - return set() - return {_canonical_name(e.strip()) for e in match.group(1).split(",") if e.strip()} + requirements = [ + _parse_requirement(line) for line in headers.get_all("Requires-Dist") or [] + ] + return None if None in requirements else requirements def _extra_marker_applies(markers: list, extra: str) -> bool: @@ -259,38 +233,16 @@ def _extra_marker_applies(markers: list, extra: str) -> bool: return any(groups) -def _read_requires_by_extra(dist_info: Path) -> Optional[list]: - """(dependency, its extras, parsed marker) per Requires-Dist - line, or None when unreadable. other markers are ignored, which can - only keep more wheels.""" - lines = _requires_dist(dist_info) - if lines is None: - return None - entries = [] - for line in lines: - requirement, _, marker = line.partition(";") - try: - parsed = Marker(marker) if marker.strip() else None - except InvalidMarker: - return None - entries.append( - ( - _strict_requirement_name(line), - _requirement_extras(requirement), - parsed, - ) - ) - return entries - - -def _active_requirements(entries: list, extras: set) -> list: +def _active_requirements(requirements: List[Requirement], extras: set) -> list: """(dependency, its extras) for each Requires-Dist entry that an install with these extras pulls in. non-extra markers are ignored.""" return [ - (dep, dep_extras) - for dep, dep_extras, marker in entries - if marker is None - or any(_extra_marker_applies(marker._markers, extra) for extra in {""} | extras) + (canonicalize_name(req.name), {canonicalize_name(e) for e in req.extras}) + for req in requirements + if req.marker is None + or any( + _extra_marker_applies(req.marker._markers, extra) for extra in {""} | extras + ) ] @@ -301,13 +253,18 @@ def installed_closure(env_dir: Path, roots: List[str]) -> set: transformers[torch] does). unreadable metadata counts every installed distribution as pulled in.""" requires = { - name: _read_requires_by_extra(d) + name: _requires_dist(d) for name, d in _installed_dist_infos(env_dir).items() } - stack = [(_strict_requirement_name(r), _requirement_extras(r)) for r in roots] + parsed_roots = [_parse_requirement(r) for r in roots] + stack = [ + (canonicalize_name(r.name), {canonicalize_name(e) for e in r.extras}) + for r in parsed_roots + if r is not None + ] names = {name for name, _ in stack} - if "" in names or any(v is None for v in requires.values()): - return set(requires) | (names - {""}) + if None in parsed_roots or None in requires.values(): + return set(requires) | names activated: Dict[str, set] = {} while stack: node, extras = stack.pop() @@ -385,7 +342,7 @@ def prune_excluded( returns the canonical names of removed distributions. """ dist_infos = _installed_dist_infos(env_dir) - excluded_set = {_canonical_name(n) for n in excluded} + excluded_set = {canonicalize_name(n) for n in excluded} if not set(dist_infos) & excluded_set: return [] records = {name: _read_record(d) for name, d in dist_infos.items()} @@ -408,12 +365,18 @@ def _distributions_to_drop( or a root has no name.""" installed = set(dist_infos) targets = installed & excluded - requires = {name: _read_requires(d) for name, d in dist_infos.items()} - root_names = {_strict_requirement_name(r) for r in roots} - if None in records.values() or None in requires.values() or "" in root_names: + parsed_roots = [_parse_requirement(r) for r in roots] + if None in records.values() or None in parsed_roots: return targets - - provided = {_canonical_name(n) for n in image_provided} + root_names = {canonicalize_name(r.name) for r in parsed_roots} + requires = {} + for name, dist_info in dist_infos.items(): + dependencies = _requires_dist(dist_info) + if dependencies is None: + return targets + requires[name] = {canonicalize_name(r.name) for r in dependencies} + + provided = {canonicalize_name(n) for n in image_provided} pulled_by_provided = _dependency_closure(requires, installed & provided, set()) # keep what the roots, unrelated wheels, and user exclusions (which # remove only the named package) need, without walking into @@ -636,7 +599,7 @@ def _exclusions_for_cpu_workers( cold start. the gpu images ship the torch family, so only cpu workers act on those entries: keep one only when the cpu resources need it. """ - user_excluded = {_canonical_name(requirement_name(e)) for e in user_exclude} + user_excluded = {canonicalize_name(requirement_name(e)) for e in user_exclude} kept = [ name for name in excluded @@ -744,7 +707,7 @@ def build_environment( ", ".join(pruned), ) excluded = sorted( - set(excluded) | ({_canonical_name(n) for n in prune_names} & set(pruned)) + set(excluded) | ({canonicalize_name(n) for n in prune_names} & set(pruned)) ) if cpu_needs is not None: excluded = _exclusions_for_cpu_workers(excluded, exclude or [], cpu_needs) diff --git a/tests/test_apps/test_build.py b/tests/test_apps/test_build.py index 8750b0f44..dde4e5aa6 100644 --- a/tests/test_apps/test_build.py +++ b/tests/test_apps/test_build.py @@ -582,6 +582,37 @@ def test_cpu_keeps_transitive_torch(self, tmp_path): class TestInstalledClosure: + @pytest.mark.parametrize("requested", [False, True]) + def test_named_url_requirements_preserve_extras_and_markers(self, tmp_path, requested): + _wheel( + tmp_path, + "models", + requires=[ + 'torch @ https://example.com/torch.whl;download=1 ; extra == "gpu-mode"' + ], + ) + _wheel(tmp_path, "torch") + root = "models[GPU_Mode]" if requested else "models" + assert installed_closure( + tmp_path, [f"{root} @ https://example.com/models.whl"] + ) == ({"models", "torch"} if requested else {"models"}) + + @pytest.mark.parametrize( + "root", + [ + "https://example.com/models.whl", + "./models.whl", + "models.whl", + "models.zip", + "models.tar.gz", + "models.tgz", + ], + ) + def test_unnamed_roots_fail_closed(self, tmp_path, root): + _wheel(tmp_path, "models") + _wheel(tmp_path, "other") + assert installed_closure(tmp_path, [root]) == {"models", "other"} + def test_transitive_extras_activate_grouped_marker(self, tmp_path): _wheel(tmp_path, "wrapper", requires=["models[torch]"]) _wheel(