Context Manager & API Integration
Relevant source files
The following files were used as context for generating this wiki page:
NeuralDBG utilizes a non-invasive integration pattern that allows researchers to monitor deep learning training dynamics without modifying their existing model code. The system is designed as a "Hub & Spokes" architecture, where the core logic remains stable while external frameworks like PyTorch interact through dedicated adapters.
The primary entry point for users is a Python context manager that wraps standard training loops. This interface manages the lifecycle of diagnostic hooks, captures semantic events such as gradient vanishing or activation saturation, and provides a post-mortem API for causal reasoning and failure explanation.
Sources: copilot-instructions.md, README.md:21-25, infrastructure_planning/CODEBASE_GUIDE.md:65-68
NeuralDbg Context Manager
The NeuralDbg class implements the context manager protocol (__enter__ and __exit__), serving as the primary interface for instrumenting PyTorch models. When entered, it attaches forward and backward hooks to module boundaries to monitor tensor statistics.
Lifecycle Management
The context manager ensures that monitoring is scoped strictly to the training session. Upon exit, hooks are removed to prevent memory leaks and ensure the model returns to its original state.
The diagram shows how NeuralDbg intercepts the training process by installing hooks at the start of the context and cleaning them up at the end.
Sources: infrastructure_planning/CODEBASE_GUIDE.md:38-41, plans/reasoning_engine_analysis.md:310-315
Key Context Methods
| Method | Description |
|---|---|
__enter__ | Initializes the monitoring session and installs PyTorch hooks. |
__exit__ | Finalizes event collection, removes hooks, and triggers pruning if necessary. |
_install_hooks | Attaches callbacks to every layer to observe training dynamics. |
Sources: infrastructure_planning/CODEBASE_GUIDE.md:38-41, plans/reasoning_engine_analysis.md:310-313
Inference & Diagnostic API
After a training run completes—or fails—NeuralDBG provides a set of API methods to query the extracted semantic trace. This API focuses on "Causal Compression," reducing raw event data into high-level explanations.
Explanation and Root Cause Analysis
The API allows users to query the engine for ranked causal hypotheses. This moves beyond simple metric tracking to identify the specific layer and step where a failure originated.
This flowchart illustrates the data flow from raw semantic events to the various diagnostic outputs provided by the API.
Sources: README.md:55-66, infrastructure_planning/CODEBASE_GUIDE.md:42-46
API Reference Table
| Endpoint / Method | Return Type | Description |
|---|---|---|
explain_failure() | List[str] | Main entry point for natural language causal reasoning. |
get_root_causes() | List[SemanticEvent] | Identifies the earliest events in a failure chain. |
detect_coupled_failures() | List[Coupling] | Identifies links between different failure types (e.g., saturation leading to vanishing). |
export_mermaid_causal_graph() | str | Generates a Mermaid-compatible string to visualize the failure trace. |
Sources: README.md:60-66, infrastructure_planning/CODEBASE_GUIDE.md:42-46
Integration Patterns
NeuralDBG is designed to be "compiler-aware," meaning it operates at module boundaries. This ensures compatibility with torch.compile by avoiding tensor inspection within the inner loops of optimized kernels.
Hook Mechanism
The system utilizes PyTorch hooks to observe training without altering the business logic.
- Forward Hooks: Monitor activation health (e.g., dead neurons, saturation).
- Backward Hooks: Monitor gradient health (e.g., vanishing or exploding gradients).
The diagram depicts the Hub and Spokes architecture where the PyTorch adapter feeds semantic events into the core reasoning logic.
Sources: copilot-instructions.md:46-52, infrastructure_planning/CODEBASE_GUIDE.md:104-115, ROADMAP.md:46-50
Performance & Resource Management
To maintain a low overhead, the API includes mechanisms for event pruning and garbage collection. This prevents the event list from growing unbounded during long training sessions.
- Event Pruning: Keeps only high-confidence or recent events based on a configurable maximum limit.
- Weak References: Uses
weakrefto point to the model, ensuring the debugger does not prevent the model from being garbage collected by the Python runtime.
Sources: plans/reasoning_engine_analysis.md:300-310, plans/reasoning_engine_analysis.md:320-325
The Context Manager and API Integration serve as the bridge between raw neural network training and structured causal analysis, providing a non-invasive yet deep diagnostic capability for ML researchers.