| |
| """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()) |
|
|
