feat(jit): add explicit module registration

.github/workflows/ci.yml

@@ -317,7 +317,8 @@ jobs:
       ${{ !cancelled() && needs.wheel.result == 'success' &&
           (needs.pytorch-image.result == 'success' || needs.pytorch-image.result == 'skipped') }}
     runs-on: prod-aitune-tester-rtx-pro-4500-v1
-    timeout-minutes: 40
+    # Fresh-process JIT cases include native backend import time for every test.
+    timeout-minutes: 60
     permissions:
       contents: read
       packages: read

AGENTS.md

@@ -86,7 +86,7 @@ Key fixtures in `conftest.py` (all autouse unless noted):
 
 - `module_registry_cleanup` — clears the global module registry after each test.
 - `aitune_cache_dir` — isolates the AITune cache in a temporary path per test.
-- `jit_cleanup` — calls `jit_reset()` before and after each test.
+- JIT-mutating tests and JIT doctests run in fresh Python processes; there is no production JIT reset API.
 - `aitune_logging_setup` — initializes logging; set `AITUNE_TESTS_LOG_LEVEL=DEBUG`
   for verbose test logs.
 - `torch_device` (not autouse) — selects the test device; override with

CHANGELOG.md

@@ -7,6 +7,8 @@ title: "Changelog"
 
 ## 0.7.0 (unreleased)
 
+- feat: explicitly register existing module instances as targeted JIT tuning heads
+- breaking change: JIT registration is one complete batch per worker; remove JIT reset and registration rollback
 - feat: support tuning existing ONNX models with ONNX Runtime and TensorRT
 - feat: expose portable deployment artifacts for ONNX Runtime, TensorRT, and TorchInductor AOT models, including required additional files and representative inputs
 - feat: publish tuned artifacts or existing ONNX, TensorRT, and PT2 files as Triton model repositories with backend-specific configurations

README.md

@@ -100,6 +100,20 @@ A new process starts tuning again; use the explicit workflow below for saved art
 See the [JIT guide](https://docs.nvidia.com/aitune/guides/just-in-time-tuning/) for configuration and deferred tuning
 for pipelines.
 
+If another library or runtime has already constructed the modules you want to tune, register the final instances
+directly:
+
+```python
+from aitune.torch import register_for_jit_tuning
+
+registration = register_for_jit_tuning([pipe.unet])
+```
+
+Each registered module is an independent top-level JIT target. Do not combine explicit registration with automatic
+constructor interception in the same live JIT session. The returned registration exposes live inspection reports,
+state counts, and an `all_tuned` snapshot. See the
+[JIT tuning guide](docs/guides/jit_tuning.md#registering-existing-modules) for the complete contract.
+
 ### Ahead-of-time (AOT): prepare a reusable tuning artifact
 
 Ahead-of-time tuning prepares your model for deployment with a consistent flow:

aitune/torch/__init__.py

@@ -28,6 +28,7 @@
 from aitune.torch.jit.config import config as jit_config
 from aitune.torch.jit.patched_module import PatchedModule
 from aitune.torch.jit.patcher import patch_for_jit_tuning
+from aitune.torch.jit.registration import JITRegistration, register_for_jit_tuning
 
 __all__ = [
     "aitune_cache_dir",
@@ -58,6 +59,8 @@
     "LocalTorchStorage",
     "DataLoaderFactory",
     "patch_for_jit_tuning",
+    "register_for_jit_tuning",
+    "JITRegistration",
 ]
 
 if config.enable_diffusers_integration:

aitune/torch/jit/inspect_module.py

@@ -90,6 +90,9 @@ def _forward_init(self, wrapped, instance, args, kwargs):
         This method is called when the module is first called.
         It initializes the module and sets the state to INSPECT.
         """
+        from aitune.torch.jit.patcher import Patcher  # avoid circular deps
+
+        Patcher._check_mode()
         params = count_parameters(self.__wrapped__)
         self._name += f" 📊{params}"
         if self._should_be_skipped():
@@ -123,6 +126,9 @@ def _forward_init(self, wrapped, instance, args, kwargs):
 
     def _forward_inspect(self, wrapped, instance, args, kwargs):
         """Forward call for the inspect time."""
+        from aitune.torch.jit.patcher import Patcher  # avoid circular deps
+
+        Patcher._check_mode()
         self._call_count += 1
         self._gather_args_kwargs(args, kwargs)
         self._restore_original_forward()
@@ -250,13 +256,6 @@ def print_history(sink=print):
         """Prints history to the sink."""
         sink("\n".join(InspectModule.history))
 
-    @staticmethod
-    def reset():
-        """Reset PatchedModule state."""
-        InspectModule.history.clear()
-        InspectModule.heads.clear()
-        InspectModule.stack.clear()
-
     @staticmethod
     def save_report(path: str | Path, model_name: str = "Model"):
         """Save module hierarchy to interactive HTML with folding/unfolding capabilities.