File size: 4,024 Bytes
8e9f35a | 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 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | #!/usr/bin/env python3
"""Verify that the packaged sources only relocate internal imports."""
from __future__ import annotations
import ast
import importlib.util
import subprocess
import sys
from pathlib import Path
PINNED_COMMIT = "8a26fb0ec9e353125ead798cb2e312d5ce48cded"
UPSTREAM_SUBDIR = Path("techniques/sparse_backends/sol_attn")
IGNORED_PARTS = {"__pycache__"}
def source_files(root: Path) -> dict[Path, Path]:
return {
path.relative_to(root): path
for path in root.rglob("*")
if path.is_file() and not (set(path.parts) & IGNORED_PARTS)
}
def module_package(relative_path: Path) -> str:
parent_parts = relative_path.with_suffix("").parts[:-1]
return ".".join(("sol_attn", *parent_parts))
def normalized_imports(tree: ast.AST, package: str):
imports = []
for node in ast.walk(tree):
if isinstance(node, ast.Import):
imports.extend(
(alias.name, alias.asname or alias.name.split(".")[0])
for alias in node.names
)
elif isinstance(node, ast.ImportFrom):
module = node.module or ""
if node.level:
relative = "." * node.level + module
module = importlib.util.resolve_name(relative, package)
imports.extend(
(
f"{module}.{alias.name}" if module else alias.name,
alias.asname or alias.name,
)
for alias in node.names
)
return sorted(imports)
def without_imports(text: str) -> str:
tree = ast.parse(text)
lines = text.splitlines(keepends=True)
for node in ast.walk(tree):
if isinstance(node, (ast.Import, ast.ImportFrom)):
for index in range(node.lineno - 1, node.end_lineno):
lines[index] = "\n" if lines[index].endswith("\n") else ""
return "".join(lines)
def main() -> int:
if len(sys.argv) != 2:
print("usage: verify_upstream.py /path/to/Sana", file=sys.stderr)
return 2
checkout = Path(sys.argv[1]).resolve()
commit = subprocess.check_output(
["git", "-C", str(checkout), "rev-parse", "HEAD"],
text=True,
).strip()
if commit != PINNED_COMMIT:
raise SystemExit(f"expected Sana {PINNED_COMMIT}, found {commit}")
upstream_root = checkout / UPSTREAM_SUBDIR
packaged_root = Path(__file__).resolve().parents[1] / "torch-ext/sol_attn"
upstream = source_files(upstream_root)
packaged = source_files(packaged_root)
if upstream.keys() != packaged.keys():
missing = sorted(upstream.keys() - packaged.keys())
extra = sorted(packaged.keys() - upstream.keys())
raise SystemExit(f"source inventory mismatch; missing={missing}, extra={extra}")
relocated = []
for relative_path in sorted(upstream):
upstream_bytes = upstream[relative_path].read_bytes()
packaged_bytes = packaged[relative_path].read_bytes()
if upstream_bytes == packaged_bytes:
continue
if relative_path.suffix != ".py":
raise SystemExit(f"non-Python source differs: {relative_path}")
upstream_text = upstream_bytes.decode()
packaged_text = packaged_bytes.decode()
if without_imports(upstream_text) != without_imports(packaged_text):
raise SystemExit(f"non-import source differs: {relative_path}")
package = module_package(relative_path)
upstream_imports = normalized_imports(ast.parse(upstream_text), package)
packaged_imports = normalized_imports(ast.parse(packaged_text), package)
if upstream_imports != packaged_imports:
raise SystemExit(f"import semantics differ: {relative_path}")
relocated.append(relative_path)
print(
"Verified pinned NVlabs/Sana source; only semantically equivalent "
f"import relocation is present in {len(relocated)} files."
)
return 0
if __name__ == "__main__":
raise SystemExit(main())
|