Merge reconciled parallel-infer-integration base into PR #140

Author: Essoz

AGENTS.md

@@ -0,0 +1,125 @@
+# TrainCheck engineering guidance
+
+This file applies throughout the repository. Preserve TrainCheck's explicit type
+and runtime contracts: catching malformed data early is part of correctness.
+Follow the user's task scope and preserve unrelated work. Resolve routine choices
+using these rules; ask when an unresolved choice changes public behavior or scope.
+
+## Types and data contracts
+
+- Fully annotate new functions and changed function interfaces, including private
+  helpers, worker payloads, container element types, and return values. When
+  changing an untyped helper's behavior, give it an explicit contract.
+- Represent structured data with appropriate aliases, `TypedDict`, dataclasses,
+  or small protocols. Prefer the simplest representation that states the real
+  contract; do not introduce abstractions solely to satisfy a checker.
+- Do not remove annotations or assertions, widen a type to `Any`, coerce values,
+  or add ignores merely to silence errors. A necessary contract change needs a
+  stated reason, caller review, and tests. Document narrow uses of `Any`, `cast`,
+  `type: ignore`, and `noqa`; existing typing debt is not a reason to add more.
+- Treat external JSON, framework metadata, and dataframe values as untrusted
+  inputs. Validate and narrow at the boundary, using `object` for unknown values
+  where practical. `cast` and `TypedDict` do not perform runtime validation.
+- Required fields must be required in both the type and implementation. Validate
+  them and use direct access, such as `record["func_call_id"]`; do not substitute
+  a default when a required field is absent or malformed.
+- `.get()`, `getattr()`, and optional types are appropriate when absence has a
+  defined meaning, such as a cache miss or an unavailable framework capability.
+  State that meaning and preserve it in the result. Do not blanket-ban these APIs.
+- Keep missing, empty, false, unobserved, skipped, unsupported, and failed states
+  distinct where they mean different things. Never turn an unknown result into
+  success, zero violations, or an empty dataset through a permissive default.
+- Types must describe accepted values, including legacy dataframe scalars. Do
+  not claim all IDs are integers or silently coerce them unless that is the
+  actual validated schema.
+
+## Failures and compatibility
+
+- Retain explicit assertions for internal invariants. Reject malformed external
+  inputs with explicit validation and useful errors; assertions alone disappear
+  under optimized Python execution.
+- Catch exceptions narrowly. A broad catch is justified for cleanup followed by
+  re-raising, or an explicit failure result whose caller handles it. Do not swallow
+  errors, continue with fabricated data, or silently fall back after corruption.
+- Keep compatibility fallbacks at clear boundaries. Distinguish an absent optional
+  API from a present API with malformed metadata or unexpected runtime failure.
+  Expose unsupported behavior and test supported and fallback paths.
+- Annotate context-manager exception and traceback parameters. If `__exit__`
+  never suppresses an exception, use `Literal[False]` for its return type.
+- Prefer small, cohesive changes. Separate deliberate semantic changes from
+  performance changes so reviewers can assess each independently.
+
+## Tracing, inference, and performance
+
+- Instrumentation must preserve application behavior: RNG streams, model state,
+  return values, exceptions, and checkpoint compatibility. Test the relevant
+  contract when changing wrappers, serialization, suppression, or sampling.
+- Tracer bookkeeping and serialization must not recursively generate user trace
+  events. Check record completeness and multiplicity, including failure paths
+  and concurrency where applicable.
+- Optimizations must preserve row order, dtypes, duplicates, identity, boundary
+  conditions, and error behavior. Document cache ownership, data stability, and
+  invalidation; test mutation or reuse when supported. A fast path needs the same
+  semantics as the fallback. Inputs unsupported by the fast path must reach the
+  fallback with unchanged semantics.
+- Compare outputs without normalizing away meaningful differences. Dropping
+  records, IDs, or array contents to obtain equality does not establish parity.
+- Invariant validation must report what actually executed: activation/coverage,
+  dormant rules, skips, failures, and violations. Zero violations alone is not
+  evidence that the intended invariants were checked.
+- For performance claims, record the baseline and candidate revisions, workload,
+  data/trace identity, configuration, environment, and reproducible commands.
+  Use matched conditions and repeats/order controls when variability matters.
+  Separate profiled timings from ordinary runs and component timings from total
+  inference time. Give the denominator for percentages, memory cost, output
+  equivalence, and material limitations. Label single-run results accordingly.
+- Keep large traces, datasets, checkpoints, credentials, and machine-specific
+  scratch files out of commits. Store durable results in the agreed location and
+  record how to find and reproduce them. Use isolated experiment directories and
+  environments on shared machines; heavy cluster work belongs on allocated
+  compute resources, not login nodes.
+
+## Validation
+
+- Use the repository's Python requirement in `pyproject.toml` and the tool versions
+  and options in `.pre-commit-config.yaml` and `.github/workflows/`. Do not disable
+  hooks, signing, assertions, or checks to make a change pass.
+- Run the configured checks for changed code, for example
+  `pre-commit run --files path/to/changed.py`, and relevant CI checks before
+  claiming readiness. Run `pre-commit run --all-files` in an isolated clean
+  checkout when broader validation is needed. Hooks can edit files: inspect and
+  retain only changes belonging to the task.
+- CI is a minimum, not complete coverage. Currently the mypy hook is scoped to
+  `traincheck/`, untyped bodies are not fully checked, tests have lint/format
+  exclusions, and tools-only changes may skip CI. Explicitly type-check changed
+  reusable tools and run relevant tests even when automation does not select them.
+- Fully check new modules and reusable helpers; use stricter mypy options such as
+  `--check-untyped-defs`, `--disallow-untyped-defs`, and `--disallow-any-generics`
+  for their scope as appropriate. For legacy modules, compare with the base and
+  introduce no new errors. Do not expand a focused task into a repository-wide
+  strictness migration. Report any remaining errors and checker coverage limits.
+- Behavioral bug fixes need focused regression tests that fail before and pass
+  after the fix. Exercise meaningful negative paths: missing required fields, wrong types,
+  malformed records, propagated exceptions, and affected fallback behavior.
+  Add state, RNG, cache, or concurrency tests when those contracts change.
+- Run existing tests relevant to affected callers and compare outputs for
+  optimizations. Do not add tests that merely repeat implementation details.
+  Documentation-only edits normally need content/link and diff checks, not a
+  runtime test suite. Report exactly what ran and what could not run.
+
+## Git and review
+
+- Inspect branch, worktree status, and applicable guidance before editing. Preserve
+  unrelated changes and shared environments. Use an isolated branch/worktree when
+  needed; stage named files or hunks, never sweep unrelated work into a commit.
+- Checkpoint coherent progress with small, signed commits and focused PRs as
+  authorized. Preserve configured signing; if signing fails, report the failure
+  instead of silently producing an unsigned commit. Do not merge without explicit
+  user authorization.
+- PR descriptions should state the problem, final approach, behavior/contracts
+  preserved or changed, validation, and limitations. Include measured before/after
+  numbers for performance work, with units, baseline, and end-to-end context.
+  Do not invent improvement numbers for documentation or correctness-only work.
+- State dependencies between stacked PRs. Recheck the composed result after
+  rebasing or integrating changes; a branch-local pass does not prove the final
+  stack works. Keep experiment evidence separate from unverified expectations.

README.md

@@ -2,73 +2,102 @@
 <picture>
   <img alt="TrainCheck logo" width="55%" src="https://raw.githubusercontent.com/OrderLab/TrainCheck/main/docs/assets/images/traincheck_logo.png">
 </picture>
-<h1>TrainCheck: Invariant Checking & Observability for AI Training</h1>
+<h1>TrainCheck: Invariant Checking for AI Training</h1>
 
 [![Chat on Discord](https://img.shields.io/badge/Discord-Join%20us-5865F2?logo=discord&logoColor=white)](https://discord.gg/ZvYewjsQ9D)
 [![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/OrderLab/TrainCheck)
 
 </div>
 
+TrainCheck catches silent training bugs by learning what a healthy run does, then checking a new run against those learned invariants. It works by tracing PyTorch API calls and model state changes, so you can inspect training behavior before a loss curve or final metric tells you something went wrong.
 
-**Stop flying blind.** TrainCheck gives you deep visibility into your training dynamics, continuously validating correctness and stability where standard metrics fail.
+## Install
 
----
+Install TrainCheck in the same Python environment that runs your training script:
 
-### Why TrainCheck?
+```bash
+pip3 install traincheck
+```
 
-✅ **Continuous Invariant Checking**
-TrainCheck validates the "physics" of your training process in real-time. It ensures your model adheres to learned invariants—such as gradient norms, tensor shapes, and update magnitudes—effectively catching silent corruption before it wastes GPU hours.
+For CUDA, conda, and source-install details, see the [Installation Guide](https://orderlab.io/TrainCheck/installation-guide/).
 
-🚀 **Holistic Observability**
-Traditional tools only show you *if* your model crashed. TrainCheck shows you *why* it's degrading, analyzing internal state dynamics that loss curves miss.
+## Use TrainCheck
 
-🧠 **Zero-Config Validation**
-No manual tests required. TrainCheck automatically learns the invariants of your specific model from healthy runs and flags deviations instantly.
+TrainCheck has four main steps.
 
-⚡ **Universal Compatibility**
-Drop-in support for PyTorch, Hugging Face, and industry-class workloads using DeepSpeed/Megatron and more.
+### 1. Collect a Reference Trace
 
----
-## Installation
+Run `traincheck-collect` on a known-good training script. This should be a short run that covers the training behavior you want TrainCheck to learn.
+
+```bash
+traincheck-collect \
+  --pyscript reference.py \
+  --models-to-track model \
+  --output-dir reference_trace
+```
 
-Install TrainCheck in the Python environment where you will run your training script:
+### 2. Infer Invariants
+
+Turn the reference trace into invariants:
 
 ```bash
-pip3 install traincheck
+traincheck-infer -f reference_trace -o invariants.json
+```
+
+### 3. Collect a Target Trace
+
+Run the target training script with the inferred invariants. Passing `--invariants` lets TrainCheck trace only the APIs and variables needed for those checks.
+
+```bash
+traincheck-collect \
+  --pyscript target.py \
+  --models-to-track model \
+  --invariants invariants.json \
+  --output-dir target_trace
 ```
 
-For detailed setup (CUDA configuration, UV, conda environments), see the [Installation Guide](https://orderlab.io/TrainCheck/installation-guide/).
+For long target runs, trace fewer steps:
 
+```bash
+traincheck-collect \
+  --pyscript target.py \
+  --models-to-track model \
+  --invariants invariants.json \
+  --sampling-interval 10 \
+  --warm-up-steps 10 \
+  --output-dir target_trace
+```
 
-### How It Works
+### 4. Check the Target Run
 
-1. **Instrument**: We wrap your training loop with lightweight probes—no code changes needed.
-2. **Learn**: We analyze correct runs to infer *invariants* (mathematical rules of healthy training).
-3. **Check**: We monitor new runs in real-time, verifying every step against learned invariants to catch silent logic bugs and hardware faults.
+For live checking, start `traincheck-onlinecheck` while the target run is writing traces:
 
-![Workflow](https://raw.githubusercontent.com/OrderLab/TrainCheck/main/docs/assets/images/workflow.png)
+```bash
+traincheck-onlinecheck -f target_trace -i invariants.json
+```
 
-## 🔥 Try TrainCheck
+The easier offline path is to wait for trace collection to finish, then run:
 
-Work through [5‑Minute Experience with TrainCheck](./docs/5-min-tutorial.md). You’ll learn how to:
-   - Instrument a training script and collect a trace  
-   - Automatically infer invariants  
-   - Uncover silent bugs in the training script
+```bash
+traincheck-check -f target_trace -i invariants.json
+```
+
+Both checkers write a results directory with failure logs and a `report.html` summary.
 
-## Documentation
+## Learn More
 
-- **[Installation Guide](https://orderlab.io/TrainCheck/installation-guide/)**
-- **[Usage Guide: Scenarios and Limitations](https://orderlab.io/TrainCheck/usage-guide/)**
-- **[TrainCheck Technical Doc](https://orderlab.io/TrainCheck/technical-doc/)**
+- [Use TrainCheck](https://orderlab.io/TrainCheck/usage-guide/) explains the full workflow and output files.
+- [5-Minute Tutorial](./docs/5-min-tutorial.md) walks through a real silent training issue.
+- [Installation Guide](https://orderlab.io/TrainCheck/installation-guide/) covers environment setup.
+- [Technical Documentation](https://orderlab.io/TrainCheck/technical-doc/) describes invariants, trace representation, and implementation details.
 
 ## Status
 
-TrainCheck is under active development. Please join our 💬 [Discord server](https://discord.gg/VwxpJDvB) or file a GitHub issue for support. You can also reach the team at [traincheck@umich.edu](mailto:traincheck@umich.edu).
-We welcome feedback and contributions from early adopters.
+TrainCheck is under active development. Please join our [Discord server](https://discord.gg/VwxpJDvB), file a GitHub issue, or email [traincheck@umich.edu](mailto:traincheck@umich.edu).
 
 ## Contributing
 
-We welcome and value any contributions and collaborations. Please check out [Contributing to TrainCheck](./CONTRIBUTING.md) for how to get involved.
+We welcome contributions. See [Contributing to TrainCheck](./CONTRIBUTING.md) for setup and contribution guidance.
 
 ## License
 
@@ -77,6 +106,7 @@ TrainCheck is licensed under the [Apache License 2.0](./LICENSE).
 ## Citation
 
 If TrainCheck is relevant to your work, please cite our paper:
+
 ```bib
 @inproceedings{TrainCheckOSDI2025,
   author = {Jiang, Yuxuan and Zhou, Ziming and Xu, Boyu and Liu, Beijie and Xu, Runhui and Huang, Peng},
@@ -90,7 +120,6 @@ If TrainCheck is relevant to your work, please cite our paper:
 }
 ```
 
-
 ## Artifact Evaluation
 
-🕵️‍♀️ OSDI AE members, please see [TrainCheck AE Guide](./docs/ae.md).
+OSDI AE members should use the [TrainCheck AE Guide](./docs/ae.md).

SUGGESTION.md

@@ -0,0 +1,59 @@
+# Suggestion for PR #137
+
+The current fix using `rsplit()` is better than the original `split()`, but it can still fail in edge cases where directory names repeat patterns.
+
+## The Problem
+Consider a file path like `/abc/abcd/abc/abcd/file.py` where `root_module = "abc"`:
+- Current approach: `rsplit(f"/abc/", 1)` splits on the last `/abc/`, producing `abcd/file.py`
+- Result: `abc.abcd.file` instead of the correct `abc.abcd.abc.abcd.file`
+
+## More Robust Solution
+Instead of relying on string parsing with the module name, get the actual filesystem location of the root module and remove it from the file path:
+
+```python
+def get_module_path_from_file_path(file_path: str, root_module: str) -> str | None:
+    import importlib.util
+    
+    # Validate inputs
+    if (
+        not file_path.endswith(".py")
+        or not os.path.exists(file_path)
+        or f"/{root_module}/" not in file_path
+    ):
+        return None
+    
+    # Get the actual location of the root module
+    try:
+        spec = importlib.util.find_spec(root_module)
+        if spec is None or spec.origin is None:
+            return None
+        
+        module_dir = os.path.dirname(spec.origin)
+        # Go up one level if it's an __init__.py
+        if os.path.basename(spec.origin) == "__init__.py":
+            module_dir = os.path.dirname(module_dir)
+    except (ImportError, AttributeError, ValueError):
+        return None
+    
+    # Get relative path from module directory
+    try:
+        relative_path = os.path.relpath(file_path, module_dir)
+    except ValueError:
+        # Paths on different drives on Windows
+        return None
+    
+    # Convert to module path
+    if relative_path.endswith(".py"):
+        relative_path = relative_path[:-3]
+    module_path = root_module + "." + relative_path.replace(os.sep, ".")
+    
+    return module_path
+```
+
+## Benefits
+- **No string parsing vulnerability** - Works regardless of repeated naming patterns
+- **Uses actual filesystem paths** - More reliable and maintainable
+- **Handles edge cases** - Works with __init__.py files and package structures
+- **More explicit** - The intent is clearer to future maintainers
+
+This approach is immune to all naming collision edge cases!

docs/ae-eval-s5.5-perf-overhead.md

@@ -21,7 +21,7 @@ This evaluation measures the runtime overhead introduced by TrainCheck’s instr
     - Located in [overhead-e2e](../eval_scripts/perf_benchmark/overhead-e2e)
 
 - The deployed 100 invariants:
-    [eval_scripts/perf_benchmark/overhead-e2e/sampled_100_invariants.json](../eval_scripts/perf_benchmark/overhead-e2e/sampled_100_invariants.json)
+    `eval_scripts/perf_benchmark/overhead-e2e/sampled_100_invariants.json`
 
 
 ## 🛠 How to Run

docs/assets/blog/traincheck-launch/02-bloom-timeline.png

@@ -0,0 +1,4 @@
+# TrainCheck Blog
+
+Research notes, engineering lessons, and project updates about dependable
+machine learning training.