trainlint 0.1.0__py3-none-any.whl

trainlint/__init__.py

@@ -0,0 +1,3 @@
+"""Static checks for ML code that runs fine and is still wrong."""
+
+__version__ = "0.1.0"

trainlint/__main__.py

@@ -0,0 +1,3 @@
+from trainlint.cli import main
+
+raise SystemExit(main())

trainlint/astutil.py

@@ -0,0 +1,122 @@
+"""Helpers for matching calls by name. Imports are not resolved, so `F.softmax`,
+`torch.softmax` and a bare `softmax` all have the tail `softmax`."""
+
+from __future__ import annotations
+
+import ast
+import re
+from collections.abc import Iterable, Iterator
+
+_NESTED_SCOPES = (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)
+
+# Variable names are the only signal available for which split a value belongs to.
+_TRAIN_NAME = re.compile(r"(?:^|_)(?:x|y)?(?:train|trn|tr)(?:$|_|\d)", re.IGNORECASE)
+_EVAL_NAME = re.compile(
+    r"(?:^|_)(?:x|y)?(?:test|val|valid|validation|eval|dev|holdout)(?:$|_|\d)", re.IGNORECASE
+)
+
+
+def dotted(node: ast.AST) -> str | None:
+    """Return 'a.b.c' for a chain of names and attributes, otherwise None."""
+    parts: list[str] = []
+    while isinstance(node, ast.Attribute):
+        parts.append(node.attr)
+        node = node.value
+    if isinstance(node, ast.Name):
+        parts.append(node.id)
+        return ".".join(reversed(parts))
+    return None
+
+
+def call_tail(call: ast.Call) -> str | None:
+    """Last component of the callee: 'softmax' for both F.softmax(x) and softmax(x)."""
+    func = call.func
+    if isinstance(func, ast.Attribute):
+        return func.attr
+    if isinstance(func, ast.Name):
+        return func.id
+    return None
+
+
+def receiver(call: ast.Call) -> str | None:
+    """Dotted name a method is called on: 'self.opt' for self.opt.step()."""
+    if isinstance(call.func, ast.Attribute):
+        return dotted(call.func.value)
+    return None
+
+
+def root_name(node: ast.AST) -> str | None:
+    """Base variable of X, X[idx], X.values and X[cols].values."""
+    while isinstance(node, ast.Attribute | ast.Subscript):
+        node = node.value
+    return node.id if isinstance(node, ast.Name) else None
+
+
+def target_names(target: ast.AST) -> list[str]:
+    """Names bound by an assignment target, including tuple unpacking."""
+    if isinstance(target, ast.Name):
+        return [target.id]
+    if isinstance(target, ast.Tuple | ast.List):
+        return [name for elt in target.elts for name in target_names(elt)]
+    return []
+
+
+def call_targets(nodes: Iterable[ast.AST]) -> dict[int, set[str]]:
+    """Names each call's result is assigned to, keyed by id(call)."""
+    bound: dict[int, set[str]] = {}
+    for node in nodes:
+        if isinstance(node, ast.Assign) and isinstance(node.value, ast.Call):
+            bound[id(node.value)] = {n for t in node.targets for n in target_names(t)}
+    return bound
+
+
+def kwarg(call: ast.Call, name: str) -> ast.expr | None:
+    for keyword in call.keywords:
+        if keyword.arg == name:
+            return keyword.value
+    return None
+
+
+def is_const(node: ast.AST | None, value: object) -> bool:
+    return (
+        isinstance(node, ast.Constant) and type(node.value) is type(value) and node.value == value
+    )
+
+
+def is_train_name(name: str) -> bool:
+    return bool(_TRAIN_NAME.search(name))
+
+
+def is_eval_name(name: str) -> bool:
+    return bool(_EVAL_NAME.search(name))
+
+
+def by_position(nodes: Iterable[ast.AST]) -> list[ast.AST]:
+    return sorted(nodes, key=lambda n: (n.lineno, n.col_offset))
+
+
+def scopes(tree: ast.Module) -> Iterator[ast.Module | ast.FunctionDef | ast.AsyncFunctionDef]:
+    """The module and every function in it."""
+    yield tree
+    for node in ast.walk(tree):
+        if isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef):
+            yield node
+
+
+def walk_scope(scope: ast.AST) -> Iterator[ast.AST]:
+    """Walk a scope without entering nested functions or classes."""
+    stack = list(ast.iter_child_nodes(scope))
+    while stack:
+        node = stack.pop()
+        yield node
+        if not isinstance(node, _NESTED_SCOPES):
+            stack.extend(ast.iter_child_nodes(node))
+
+
+def calls_in(nodes: Iterable[ast.AST]) -> list[ast.Call]:
+    """Calls among `nodes` in source order."""
+    return by_position(n for n in nodes if isinstance(n, ast.Call))
+
+
+def all_calls(tree: ast.AST) -> list[ast.Call]:
+    return calls_in(ast.walk(tree))

trainlint/cli.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from collections.abc import Sequence
+from pathlib import Path
+
+from trainlint import __version__
+from trainlint.engine import lint_paths, select_rules
+from trainlint.report import render_json, render_text
+from trainlint.rules import ALL_RULES
+from trainlint.skills_install import AGENT_DIRS, default_destination, install
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="trainlint",
+        description="Find ML training and evaluation bugs that do not crash.",
+    )
+    parser.add_argument("--version", action="version", version=f"trainlint {__version__}")
+    commands = parser.add_subparsers(dest="command", required=True)
+
+    check = commands.add_parser("check", help="check Python files or directories")
+    check.add_argument("paths", nargs="*", type=Path, default=[Path()], help="default: .")
+    check.add_argument("--select", metavar="RULES", help="comma-separated rule ids to run")
+    check.add_argument("--ignore", metavar="RULES", help="comma-separated rule ids to skip")
+    check.add_argument("--format", choices=("text", "json"), default="text")
+    check.add_argument(
+        "--color", choices=("auto", "always", "never"), default="auto", help="default: auto"
+    )
+    check.add_argument("--exit-zero", action="store_true", help="exit 0 even if problems are found")
+    check.set_defaults(run=_check)
+
+    rules = commands.add_parser("rules", help="list the available rules")
+    rules.set_defaults(run=_rules)
+
+    skills = commands.add_parser("skills", help="manage the bundled agent skills")
+    skill_commands = skills.add_subparsers(dest="skills_command", required=True)
+    install = skill_commands.add_parser(
+        "install", help="copy the skills into a project or user skills folder"
+    )
+    install.add_argument(
+        "--agent",
+        choices=sorted(AGENT_DIRS),
+        default="claude",
+        help="which agent's skills folder to use (default: claude)",
+    )
+    install.add_argument(
+        "--global", dest="user_wide", action="store_true", help="install for the current user"
+    )
+    install.add_argument("--dest", type=Path, help="install into this folder instead")
+    install.add_argument("--force", action="store_true", help="overwrite skills already there")
+    install.set_defaults(run=_install_skills)
+    return parser
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+    stdout_reconfigure = getattr(sys.stdout, "reconfigure", None)
+    if stdout_reconfigure:
+        stdout_reconfigure(errors="replace")
+    try:
+        return args.run(args)
+    except (FileNotFoundError, ValueError) as error:
+        print(f"trainlint: error: {error}", file=sys.stderr)
+        return 2
+
+
+def _split(value: str | None) -> list[str]:
+    return [part.strip() for part in value.split(",") if part.strip()] if value else []
+
+
+def _check(args: argparse.Namespace) -> int:
+    rules = select_rules(_split(args.select), _split(args.ignore))
+    report = lint_paths(args.paths, rules)
+    for path, reason in report.skipped:
+        print(f"trainlint: skipped {path} ({reason})", file=sys.stderr)
+    if args.format == "json":
+        print(render_json(report))
+    else:
+        print(render_text(report, color=_use_color(args.color)))
+    return 1 if report.findings and not args.exit_zero else 0
+
+
+def _rules(_: argparse.Namespace) -> int:
+    for rule in ALL_RULES:
+        print(f"{rule.id}  {rule.title}")
+    return 0
+
+
+def _install_skills(args: argparse.Namespace) -> int:
+    destination = args.dest or default_destination(args.agent, args.user_wide)
+    installed, skipped = install(destination, force=args.force)
+    for path in installed:
+        print(f"installed {path}")
+    for path in skipped:
+        print(f"skipped {path} (already exists, use --force to replace)")
+    return 0
+
+
+def _use_color(mode: str) -> bool:
+    if mode != "auto":
+        return mode == "always"
+    if "NO_COLOR" in os.environ or not sys.stdout.isatty():
+        return False
+    return os.name != "nt" or _enable_windows_ansi()
+
+
+def _enable_windows_ansi() -> bool:
+    import ctypes
+
+    kernel32 = ctypes.windll.kernel32
+    handle = kernel32.GetStdHandle(-11)
+    mode = ctypes.c_ulong()
+    if not kernel32.GetConsoleMode(handle, ctypes.byref(mode)):
+        return False
+    return bool(kernel32.SetConsoleMode(handle, mode.value | 0x0004))

trainlint/engine.py

@@ -0,0 +1,132 @@
+"""Parse files, run rules, apply inline suppressions."""
+
+from __future__ import annotations
+
+import ast
+import os
+import re
+import tokenize
+from collections.abc import Iterable, Iterator, Sequence
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from trainlint.rules import ALL_RULES, Rule
+
+SKIP_DIRS = frozenset(
+    {
+        ".git",
+        ".hg",
+        ".venv",
+        "venv",
+        "env",
+        "node_modules",
+        "__pycache__",
+        ".tox",
+        ".mypy_cache",
+        ".ruff_cache",
+        ".pytest_cache",
+        "build",
+        "dist",
+        "site-packages",
+    }
+)
+
+_IGNORE_COMMENT = re.compile(r"#\s*trainlint:\s*ignore(?:\[([^\]]*)\])?", re.IGNORECASE)
+_LINE_BREAK = re.compile(r"\r\n|\r|\n")
+
+
+@dataclass(frozen=True, order=True)
+class Finding:
+    path: str
+    line: int
+    col: int
+    end_col: int
+    rule: str
+    message: str
+    hint: str
+    text: str
+
+
+@dataclass
+class Report:
+    findings: list[Finding] = field(default_factory=list)
+    skipped: list[tuple[str, str]] = field(default_factory=list)
+    files_checked: int = 0
+
+
+def select_rules(select: Iterable[str] = (), ignore: Iterable[str] = ()) -> list[Rule]:
+    """Rules to run. An empty `select` means every rule."""
+    known = {rule.id for rule in ALL_RULES}
+    wanted = {code.upper() for code in select}
+    dropped = {code.upper() for code in ignore}
+    unknown = sorted((wanted | dropped) - known)
+    if unknown:
+        raise ValueError(f"unknown rule {', '.join(unknown)} (see `trainlint rules`)")
+    return [r for r in ALL_RULES if (not wanted or r.id in wanted) and r.id not in dropped]
+
+
+def lint_source(source: str, path: str, rules: Sequence[Rule] = ALL_RULES) -> list[Finding]:
+    """Check one file's source. Raises SyntaxError if it does not parse."""
+    tree = ast.parse(source, filename=path)
+    lines = _LINE_BREAK.split(source)
+    findings = []
+    for rule in rules:
+        for hit in rule.check(tree):
+            line = hit.node.lineno
+            text = lines[line - 1]
+            if _suppressed(text, rule.id):
+                continue
+            col, end_col = _columns(hit.node, text)
+            findings.append(Finding(path, line, col, end_col, rule.id, hit.message, hit.hint, text))
+    return sorted(findings)
+
+
+def lint_paths(paths: Iterable[Path], rules: Sequence[Rule] = ALL_RULES) -> Report:
+    report = Report()
+    for file in python_files(paths):
+        name = file.as_posix()
+        try:
+            with tokenize.open(file) as handle:
+                source = handle.read()
+            report.findings.extend(lint_source(source, name, rules))
+        except (SyntaxError, UnicodeDecodeError, ValueError) as error:
+            report.skipped.append((name, f"{type(error).__name__}: {error}"))
+            continue
+        report.files_checked += 1
+    report.findings.sort()
+    return report
+
+
+def python_files(paths: Iterable[Path]) -> Iterator[Path]:
+    """Files named on the command line, plus *.py found under directories."""
+    for path in paths:
+        if path.is_dir():
+            for root, dirs, files in os.walk(path):
+                dirs[:] = sorted(d for d in dirs if d not in SKIP_DIRS)
+                for name in sorted(files):
+                    if name.endswith(".py"):
+                        yield Path(root, name)
+        elif path.is_file():
+            yield path
+        else:
+            raise FileNotFoundError(f"no such file or directory: {path}")
+
+
+def _suppressed(line_text: str, rule_id: str) -> bool:
+    match = _IGNORE_COMMENT.search(line_text)
+    if match is None:
+        return False
+    codes = match.group(1)
+    return codes is None or rule_id in {code.strip().upper() for code in codes.split(",")}
+
+
+def _columns(node: ast.AST, text: str) -> tuple[int, int]:
+    """1-based start and end columns in characters. ast offsets are UTF-8 bytes."""
+    start = node.col_offset
+    end = node.end_col_offset
+    if node.end_lineno != node.lineno:
+        end = len(text.encode("utf-8"))
+    raw = text.encode("utf-8")
+    start_col = len(raw[:start].decode("utf-8", errors="ignore")) + 1
+    end_col = len(raw[:end].decode("utf-8", errors="ignore")) + 1
+    return start_col, max(end_col, start_col + 1)

trainlint/report.py

@@ -0,0 +1,72 @@
+"""Render a Report as terminal text or JSON."""
+
+from __future__ import annotations
+
+import json
+
+from trainlint import __version__
+from trainlint.engine import Finding, Report
+
+_CODES = {"bold": "1", "red": "31", "cyan": "36", "dim": "2"}
+
+
+def _paint(text: str, style: str, enabled: bool) -> str:
+    return f"\x1b[{_CODES[style]}m{text}\x1b[0m" if enabled else text
+
+
+def _caret_line(finding: Finding) -> str:
+    # Reuse tabs from the source line so the carets stay under the right characters.
+    indent = "".join(c if c == "\t" else " " for c in finding.text[: finding.col - 1])
+    return indent + "^" * (finding.end_col - finding.col)
+
+
+def render_text(report: Report, color: bool = False) -> str:
+    out: list[str] = []
+    for f in report.findings:
+        gutter = " " * len(str(f.line))
+        location = _paint(f"{f.path}:{f.line}:{f.col}", "bold", color)
+        out.append(f"{location}: {_paint(f.rule, 'red', color)} {f.message}")
+        out.append(f"{gutter} |")
+        out.append(f"{f.line} | {f.text}")
+        out.append(f"{gutter} | {_paint(_caret_line(f), 'red', color)}")
+        out.append(f"{gutter} = {_paint('help:', 'cyan', color)} {f.hint}")
+        out.append("")
+    out.append(_summary(report))
+    return "\n".join(out)
+
+
+def _summary(report: Report) -> str:
+    checked = _plural(report.files_checked, "file")
+    if not report.findings:
+        return f"No problems found in {checked}."
+    flagged = len({f.path for f in report.findings})
+    return (
+        f"Found {_plural(len(report.findings), 'problem')} in "
+        f"{_plural(flagged, 'file')} (checked {checked})."
+    )
+
+
+def _plural(count: int, noun: str) -> str:
+    return f"{count} {noun}" if count == 1 else f"{count} {noun}s"
+
+
+def render_json(report: Report) -> str:
+    findings = [
+        {
+            "path": f.path,
+            "line": f.line,
+            "column": f.col,
+            "rule": f.rule,
+            "message": f.message,
+            "hint": f.hint,
+        }
+        for f in report.findings
+    ]
+    skipped = [{"path": path, "reason": reason} for path, reason in report.skipped]
+    payload = {
+        "version": __version__,
+        "files_checked": report.files_checked,
+        "findings": findings,
+        "skipped": skipped,
+    }
+    return json.dumps(payload, indent=2)

trainlint/rules/__init__.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from trainlint.rules import checkpoints, evaluation, leakage, mixed_precision, training
+from trainlint.rules.base import Hit, Rule
+
+ALL_RULES: tuple[Rule, ...] = tuple(
+    sorted(
+        [
+            *leakage.RULES,
+            *training.RULES,
+            *mixed_precision.RULES,
+            *checkpoints.RULES,
+            *evaluation.RULES,
+        ],
+        key=lambda rule: rule.id,
+    )
+)
+
+__all__ = ["ALL_RULES", "Hit", "Rule"]

trainlint/rules/base.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+import ast
+from collections.abc import Callable, Iterator
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class Hit:
+    """A spot a rule objects to, before the engine attaches file and line text."""
+
+    node: ast.expr | ast.stmt
+    message: str
+    hint: str
+
+
+@dataclass(frozen=True)
+class Rule:
+    id: str
+    title: str
+    check: Callable[[ast.Module], Iterator[Hit]]

trainlint/rules/checkpoints.py

@@ -0,0 +1,68 @@
+"""TL007: checkpoints that cannot restore a training run."""
+
+from __future__ import annotations
+
+import ast
+from collections.abc import Iterator
+
+from trainlint.astutil import all_calls, call_tail, dotted
+from trainlint.rules.base import Hit, Rule
+
+_OPTIMIZER_CLASSES = frozenset({"sgd", "adam", "adamw", "rmsprop", "adagrad", "adadelta", "lbfgs"})
+
+
+def _makes_optimizer(call: ast.Call) -> bool:
+    func = ast.unparse(call.func).lower()
+    return "optim" in func or (call_tail(call) or "").lower() in _OPTIMIZER_CLASSES
+
+
+def _words(node: ast.AST) -> str:
+    """The identifier or string a node carries, if any."""
+    if isinstance(node, ast.Name):
+        return node.id
+    if isinstance(node, ast.arg):
+        return node.arg
+    if isinstance(node, ast.Attribute):
+        return node.attr
+    if isinstance(node, ast.Constant) and isinstance(node.value, str):
+        return node.value
+    return ""
+
+
+def _mentions_resume(tree: ast.Module) -> bool:
+    words = (_words(node).lower() for node in ast.walk(tree))
+    return any("resume" in w or "start_epoch" in w for w in words)
+
+
+def _saves_optimizer_state(calls: list[ast.Call]) -> bool:
+    return any(
+        call_tail(c) == "state_dict"
+        and isinstance(c.func, ast.Attribute)
+        and "opt" in ast.unparse(c.func.value).lower()
+        for c in calls
+    )
+
+
+def _resume_without_optimizer_state(tree: ast.Module) -> Iterator[Hit]:
+    calls = all_calls(tree)
+    saves = [c for c in calls if dotted(c.func) == "torch.save"]
+    if not saves or not _mentions_resume(tree):
+        return
+    if not any(_makes_optimizer(c) for c in calls) or _saves_optimizer_state(calls):
+        return
+    yield Hit(
+        saves[0],
+        "this script resumes from checkpoints but never saves optimizer state, so momentum "
+        "and adaptive learning rates restart from zero after a resume",
+        "Save model, optimizer, scheduler, GradScaler and epoch together, for example "
+        "{'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': epoch}.",
+    )
+
+
+RULES = (
+    Rule(
+        "TL007",
+        "Resumable script saves checkpoints without optimizer state",
+        _resume_without_optimizer_state,
+    ),
+)

trainlint/rules/evaluation.py

@@ -0,0 +1,117 @@
+"""TL011 and TL012: evaluation code that measures the wrong thing."""
+
+from __future__ import annotations
+
+import ast
+from collections.abc import Iterator
+
+from trainlint.astutil import (
+    by_position,
+    call_tail,
+    call_targets,
+    calls_in,
+    is_const,
+    is_eval_name,
+    is_train_name,
+    kwarg,
+    receiver,
+    root_name,
+    scopes,
+    walk_scope,
+)
+from trainlint.rules.base import Hit, Rule
+
+# Metrics where swapping y_true and y_pred changes the answer.
+_ORDER_SENSITIVE = frozenset(
+    {
+        "f1_score",
+        "fbeta_score",
+        "precision_score",
+        "recall_score",
+        "precision_recall_fscore_support",
+        "balanced_accuracy_score",
+        "confusion_matrix",
+        "classification_report",
+        "r2_score",
+        "explained_variance_score",
+        "log_loss",
+    }
+)
+_NEEDS_SCORES = frozenset({"roc_auc_score", "average_precision_score"})
+_HARD_LABEL_CALLS = frozenset({"predict", "argmax"})
+# Boosting libraries return probabilities from Booster.predict for binary objectives.
+_BOOSTER_HINTS = ("booster", "bst", "gbm", "lgb", "xgb", "catboost")
+
+
+def _swapped_arguments(call: ast.Call) -> bool:
+    if len(call.args) < 2 or call_tail(call) not in _ORDER_SENSITIVE | _NEEDS_SCORES:
+        return False
+    first, second = root_name(call.args[0]), root_name(call.args[1])
+    return bool(first and second and "pred" in first.lower() and "pred" not in second.lower())
+
+
+def _latest_value(nodes: list[ast.AST], name: str, before: ast.AST) -> ast.expr | None:
+    assignments = [
+        n
+        for n in nodes
+        if isinstance(n, ast.Assign)
+        and any(isinstance(t, ast.Name) and t.id == name for t in n.targets)
+        and (n.lineno, n.col_offset) < (before.lineno, before.col_offset)
+    ]
+    return by_position(assignments)[-1].value if assignments else None
+
+
+def _is_hard_labels(value: ast.expr | None) -> bool:
+    if not isinstance(value, ast.Call) or call_tail(value) not in _HARD_LABEL_CALLS:
+        return False
+    owner = (receiver(value) or "").lower()
+    return not any(hint in owner for hint in _BOOSTER_HINTS)
+
+
+def _metric_misuse(tree: ast.Module) -> Iterator[Hit]:
+    for scope in scopes(tree):
+        nodes = list(walk_scope(scope))
+        for call in calls_in(nodes):
+            tail = call_tail(call)
+            if _swapped_arguments(call):
+                yield Hit(
+                    call,
+                    f"`{tail}` takes (y_true, y_pred), but the first argument looks like "
+                    "predictions",
+                    "Pass the ground-truth labels first and the predictions second.",
+                )
+            elif tail in _NEEDS_SCORES and len(call.args) >= 2:
+                scores = call.args[1]
+                if isinstance(scores, ast.Name):
+                    scores = _latest_value(nodes, scores.id, call) or scores
+                if _is_hard_labels(scores):
+                    yield Hit(
+                        call,
+                        f"`{tail}` is given class labels from predict(), which collapses the "
+                        "curve to a single point",
+                        "Pass scores instead: predict_proba(X)[:, 1] or decision_function(X).",
+                    )
+
+
+def _eval_loader_drops_samples(tree: ast.Module) -> Iterator[Hit]:
+    for scope in scopes(tree):
+        nodes = list(walk_scope(scope))
+        bound = call_targets(nodes)
+        for call in calls_in(nodes):
+            if call_tail(call) != "DataLoader" or not is_const(kwarg(call, "drop_last"), True):
+                continue
+            dataset = root_name(call.args[0]) if call.args else None
+            names = [*bound.get(id(call), ()), *([dataset] if dataset else [])]
+            if any(is_eval_name(n) and not is_train_name(n) for n in names):
+                yield Hit(
+                    call,
+                    "this evaluation DataLoader uses drop_last=True, so the last partial "
+                    "batch is never scored",
+                    "Use drop_last=False for validation and test loaders.",
+                )
+
+
+RULES = (
+    Rule("TL011", "Metric called with swapped arguments or hard labels", _metric_misuse),
+    Rule("TL012", "Evaluation DataLoader drops the last batch", _eval_loader_drops_samples),
+)