File size: 2,440 Bytes
b4f32d3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | from __future__ import annotations
import importlib.util
import os
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path
PATCH_DIR = Path(os.environ.get("MODAL_PATCH_DIR", "/root/patches"))
@dataclass(frozen=True)
class PatchSpec:
name: str
module: str
patch_file: str
strip: int
includes: tuple[str, ...] = ()
PATCHES = (
PatchSpec(
name="flashinfer-pr-3312",
module="flashinfer",
patch_file="flashinfer-pr-3312.patch",
strip=1,
),
)
def _package_parent(module_name: str) -> Path:
spec = importlib.util.find_spec(module_name)
if spec is None or spec.submodule_search_locations is None:
raise RuntimeError(f"Could not find installed package {module_name!r}.")
locations = list(spec.submodule_search_locations)
if not locations:
raise RuntimeError(f"Installed package {module_name!r} has no package path.")
return Path(locations[0]).resolve().parent
def _git_apply_command(spec: PatchSpec, patch_path: Path) -> list[str]:
cmd = ["git", "apply", f"-p{spec.strip}"]
for include in spec.includes:
cmd.append(f"--include={include}")
cmd.append(str(patch_path))
return cmd
def _check(cmd: list[str], *, cwd: Path) -> subprocess.CompletedProcess[str]:
return subprocess.run(
cmd,
cwd=cwd,
text=True,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
)
def _apply_patch(spec: PatchSpec) -> None:
patch_path = PATCH_DIR / spec.patch_file
if not patch_path.exists():
raise RuntimeError(f"Missing patch file: {patch_path}")
cwd = _package_parent(spec.module)
base_cmd = _git_apply_command(spec, patch_path)
reverse_cmd = [*base_cmd[:2], "--reverse", "--check", *base_cmd[2:]]
check_cmd = [*base_cmd[:2], "--check", *base_cmd[2:]]
reverse = _check(reverse_cmd, cwd=cwd)
if reverse.returncode == 0:
print(f"[patch] {spec.name} already applied under {cwd}")
return
check = _check(check_cmd, cwd=cwd)
if check.returncode != 0:
print(check.stdout, file=sys.stderr)
raise RuntimeError(f"Patch {spec.name} does not apply under {cwd}.")
print(f"[patch] applying {spec.name} under {cwd}")
subprocess.run(base_cmd, cwd=cwd, check=True)
def main() -> None:
for patch in PATCHES:
_apply_patch(patch)
if __name__ == "__main__":
main()
|