diff --git a/runpod/apps/build.py b/runpod/apps/build.py index b438f0ac..cf250032 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.""" @@ -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) diff --git a/tests/test_apps/test_build.py b/tests/test_apps/test_build.py index b757867e..a5d6d6ac 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,232 @@ 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") + # 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 + + 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"]