Retain CUDA graph executables and resource owners during export

docs/preparation-development.md

@@ -94,16 +94,17 @@ storage, downloads retained bytes, serializes resources, and owns asset writes,
 hashing, zstd level-3 executable compression, signing, conversion, rank payload
 transfer, and publication. Each executable resource is one zstd level-3 frame
 with a signed decoded length and compressed-chunk commitments.
-Export capture copies each graph's nodes, argument bytes, attributes, and
-annotated edges into native storage without instantiating an executable.
-Ignition retains the external CUDA memory pools and captured kernel owners,
-then resets the capture object to release its template. Diagnostic captures
-retain executables; their post-instantiation hooks can copy the same metadata
-before PyTorch releases their templates. Native export resolves allocations
-and validates bindings against the copied metadata. Keep captured buffers,
-contexts, copied metadata, pool owners, kernel owners, and any diagnostic
-executables alive until native work confirms a drain; an unsuccessful drain
-retains their owners.
+Ordinary export retains CUDA templates and instantiated executables, including
+their external memory pools and captured kernel owners, until native export
+has drained. Budget preparation memory for those executables across the serving
+graph catalog. Native export copies nodes, argument bytes, attributes, and
+annotated edges, then resolves allocations and validates their bindings.
+Diagnostic captures can copy the same metadata through post-instantiation hooks
+before PyTorch releases their templates. The core stores graph descriptions,
+not the Python executables; Helios creates executable instances at load time.
+Keep captured buffers, contexts, metadata, pool owners, kernel owners, and graph
+handles alive until native work confirms a drain; an unsuccessful drain retains
+their owners.
 
 Native transfer pools bound file buffers and pinned CUDA download storage from
 available host memory and preparation rank count. Preserve direct bulk I/O,

docs/runtime-core-design.md

@@ -71,16 +71,16 @@ checkpoint ranges; generated tokenizer and template assets supply text.
 Heliopause owns CUPTI callbacks, graph introspection, pointer reachability,
 device-to-host transfers, graph and core serialization, asset materialization,
 cubin debug stripping, hashing, compression, signing, conversion, rank payload channels,
-and publication. Export capture copies graph nodes, argument bytes, attributes,
-and annotated edges into native storage without instantiating an executable.
-Ignition retains external CUDA pool and kernel owners, then resets the capture
-object to release its template. Diagnostic executable captures can record the
-same metadata through post-instantiation hooks before releasing their templates.
-Export performs function discovery,
-allocation binding, relocation checks, and serialization against that copy.
-Callers supplying live templates use the same extraction and validation path.
-Export retains copied metadata, storage, contexts, external pool owners, kernel
-owners, and any diagnostic executables until native CUDA work has drained.
+and publication. Ordinary export retains each captured CUDA template and an
+instantiated executable, together with its external CUDA pool and kernel owners.
+This keeps one executable per exported graph resident during preparation.
+Native export copies graph nodes, argument bytes, attributes, and annotated edges;
+it performs function discovery, allocation binding, relocation checks, and
+serialization while those resources remain alive. The Python executable is not
+stored in the core: Helios reconstructs and instantiates graphs when loading it.
+Diagnostic captures can also record metadata through post-instantiation hooks
+before releasing their templates. Export retains graph resources, copied metadata,
+storage, contexts, pool owners, and kernel owners until native CUDA work has drained.
 A failed drain retains those
 owners. Final assembly consumes rank indexes and open source file ranges
 independently of the prepared model.

ignition/preparation/capture.py

@@ -31,7 +31,7 @@ def capture_forwards(functions, signatures, *, session, device, stream, pool,
     serving replays variants serially and may choose them in any order.
     All eager warmups precede capture so their transient workspaces do not
     coexist with retained graph intermediates.
-    ``retain_graphs`` records export metadata without executables when
+    ``retain_graphs`` keeps templates and executables for export when
     ``warmup_storage`` supplies external CUDA pool owners.
     ``before_run`` restores caller-owned trial inputs outside compiled warmup
     and graph capture; it is unused for the ordinary serving catalog.

ignition/preparation/graphs.py

@@ -32,8 +32,8 @@ class PreparedGraphBundle:
     does not expose a graph's capture device, so capture on that device remains
     a preparation invariant. Registration checks its capability and context,
     every declared buffer's device and extent, and instantiated graph handles.
-    Export-only graphs retain native metadata and external pool owners after
-    releasing their CUDA handles. Native replay still requires executables.
+    Export accepts live CUDA graphs or recorded metadata with retained external
+    pool owners. Native replay requires instantiated executables.
 
     Graphs must not use PyTorch generator-backed random operations: raw native
     launch bypasses PyTorch RNG bookkeeping. After handoff, callers must not

ignition/preparation/runtime_core.py

@@ -190,26 +190,20 @@ def reset(self):
 
 
 def record_export_graph(graph, storage, *, owners=()):
-    """Copy an uninstantiated graph before releasing its CUDA template."""
-    import flare_heliopause
-    from .module_images import ModuleImages
+    """Retain a captured template, its executable and external export resources.
 
-    images = ModuleImages.current()
-    if images is None:
-        raise RuntimeError("graph metadata recording requires an active native module recorder")
+    Owners must not reference the graph: Torch's destruction finalizer retains
+    them independently. Native replay callers must drain before resetting graphs.
+    """
     pools = (storage.pool, storage.graph_pool)
     owned = {pool.id for pool in pools}
     if not set(graph.pools()) <= owned:
         raise ValueError("export graph uses a CUDA pool without an external owner")
-    context, _ = flare_heliopause.current_context()
-    recorded = images.native.record_template(int(graph.raw_cuda_graph()), 1)
-    owner = RecordedExportGraph(images.native, recorded, context, pools, owners)
-    try:
-        graph.reset()
-    except BaseException:
-        owner.reset()
-        raise
-    return owner
+    # Device pointers in captured nodes do not retain their Python resource owners.
+    # Release these references only when the graph is reset or collected.
+    graph.retain_object((pools, tuple(owners)), synchronize_before_release=True)
+    instantiate_graph(graph)
+    return graph
 
 
 def capture_graph(*, retain_graphs=False):

tests/preparation/test_capture_graph_metadata.py

@@ -24,6 +24,7 @@
 class Graph:
     register_post_instantiate_hook = CUDAGraph.register_post_instantiate_hook
     register_destroy_callback = CUDAGraph.register_destroy_callback
+    retain_object = CUDAGraph.retain_object
 
     def __init__(self, keep_graph=False):
         self.keep_graph = keep_graph
@@ -187,7 +188,7 @@ def graph_scope(graph, **kwargs):
     assert recorder.events == ["record", "record"]
 
 
-def test_export_metadata_outlives_template_and_retains_external_owners(monkeypatch, recorder):
+def test_export_instantiates_and_retains_executable(monkeypatch, recorder):
     import flare_heliopause
 
     monkeypatch.setattr(flare_heliopause, "current_context", lambda: (100, 0))
@@ -196,21 +197,75 @@ def test_export_metadata_outlives_template_and_retains_external_owners(monkeypat
     graph.pools = lambda: [(0, 2)]
     graph.capture_end()
     kernel_owner = SimpleNamespace()
-    owner = record_export_graph(graph, storage, owners=(kernel_owner,))
-    assert graph.template == graph.executable == graph.instantiations == 0
-    assert owner.pools == (storage.pool, storage.graph_pool)
-    assert owner.owners == (kernel_owner,)
-    assert _graphs({"decode": owner}, recorder=recorder) == [dict(name="decode", recorded=1)]
-    with pytest.raises(ValueError, match="another native module recorder"):
-        _graphs({"decode": owner}, recorder=Recorder())
-    assert recorder.templates == {1: 11}
-    owner.reset()
-    owner.reset()
-    assert owner.pools == owner.owners == ()
+    returned = record_export_graph(graph, storage, owners=(kernel_owner,))
+    assert returned is graph
+    assert graph.instantiations == 1
+    assert graph.template == 11
+    assert graph.executable == 12
+    assert _graphs({"decode": graph}) == [dict(name="decode", graph=11, executable=12)]
     assert not recorder.templates
-    assert recorder.events == ["record-template", "release-template"]
-    with pytest.raises(ValueError, match="released"):
-        _graphs({"decode": owner}, recorder=recorder)
+    assert recorder.events == []
+
+
+@pytest.mark.parametrize("release", ("reset", "collection"))
+def test_export_graph_retains_external_owners_until_release(recorder, release):
+    class Owner:
+        pass
+
+    first_pool, second_pool, kernel_owner = Owner(), Owner(), Owner()
+    first_pool.id, second_pool.id = (0, 1), (0, 2)
+    storage = SimpleNamespace(pool=first_pool, graph_pool=second_pool)
+    references = tuple(weakref.ref(owner) for owner in (first_pool, second_pool, kernel_owner))
+    graph = Graph(keep_graph=True)
+    graph.pools = lambda: [(0, 2)]
+    graph.capture_end()
+    returned = record_export_graph(graph, storage, owners=(kernel_owner,))
+    assert returned is graph
+    graph_reference = weakref.ref(graph)
+    del first_pool, second_pool, kernel_owner, storage
+    gc.collect()
+    assert all(reference() is not None for reference in references)
+
+    drained = []
+    class Stream:
+        def synchronize(self):
+            assert all(reference() is not None for reference in references)
+            drained.append(True)
+
+    graph._retained.replay_streams.add(Stream())
+    if release == "reset":
+        graph.reset()
+        assert all(reference() is None for reference in references)
+        assert graph_reference() is graph
+    del graph, returned
+    gc.collect()
+    assert graph_reference() is None
+    assert all(reference() is None for reference in references)
+    assert drained == [True]
+
+
+def test_export_instantiation_failure_retains_owners_until_reset(monkeypatch, recorder):
+    class Owner:
+        pass
+
+    owner = Owner()
+    reference = weakref.ref(owner)
+    storage = SimpleNamespace(pool=SimpleNamespace(id=(0, 1)), graph_pool=SimpleNamespace(id=(0, 2)))
+    graph = Graph(keep_graph=True)
+    graph.pools = lambda: [(0, 2)]
+    graph.capture_end()
+
+    def fail_instantiation():
+        raise RuntimeError("graph instantiation failed")
+
+    monkeypatch.setattr(graph, "instantiate", fail_instantiation)
+    with pytest.raises(RuntimeError, match="graph instantiation failed"):
+        record_export_graph(graph, storage, owners=(owner,))
+    del owner
+    gc.collect()
+    assert reference() is not None
+    graph.reset()
+    assert reference() is None
 
 
 def test_export_rejects_unowned_pool_without_destroying_template(monkeypatch, recorder):
@@ -221,6 +276,7 @@ def test_export_rejects_unowned_pool_without_destroying_template(monkeypatch, re
     with pytest.raises(ValueError, match="without an external owner"):
         record_export_graph(graph, storage)
     assert graph.template == 11
+    assert graph.instantiations == 0 and not graph._retained.objects
     assert not recorder.templates
 
 
@@ -258,7 +314,7 @@ def test_export_resources_accept_recorded_graphs_and_executable_tuners(monkeypat
         bundle._validate_resources(export_graphs=True)
 
 
-def test_export_preserves_collective_consumers_after_replacing_cuda_graph(monkeypatch, recorder):
+def test_export_preserves_collective_consumer_associations(monkeypatch, recorder):
     import flare_heliopause
     from ignition.preparation.startup_tuning import record_consumer
 
@@ -283,5 +339,6 @@ def graph_scope(graph, **kwargs):
         session=session, device="cpu", stream=stream, pool=storage.graph_pool.id, compile_decoder=False,
         retain_graphs=True, warmup_storage=storage)
     graph = graphs["decode"]
-    assert isinstance(graph, RecordedExportGraph)
+    assert isinstance(graph, Graph)
+    assert graph.instantiations == 1
     assert registry.captures == {id(graph): {id(collective)}}

tests/preparation/test_recorded_export_graph_cuda.py

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 Flare community contributors
 # SPDX-License-Identifier: Apache-2.0
 
-"""Export graph metadata retains CUDA storage without graph handles."""
+"""Retained export graphs preserve storage and serialize their execution metadata."""
 
 import os
 from pathlib import Path
@@ -12,7 +12,7 @@
 
 
 @pytest.mark.skipif(os.environ.get("FLARE_RUN_CUDA_TESTS") != "1", reason="requires explicit CUDA test opt-in")
-def test_export_capture_releases_cuda_handles_and_preserves_storage(tmp_path):
+def test_export_capture_retains_cuda_handles_and_preserves_storage(tmp_path):
     subprocess.run([sys.executable, "-c",
         "from pathlib import Path; import sys; "
         "from tests.preparation.test_recorded_export_graph_cuda import check_export; "
@@ -30,7 +30,7 @@ def check_export(tmp_path):
     from ignition.preparation.core_file import read_sources
     from ignition.preparation.core_storage import CoreStorage
     from ignition.preparation.runtime_core import (
-        Allocation, RecordedExportGraph, RuntimeCore, capture_graph, instantiate_graph,
+        Allocation, RuntimeCore, capture_graph, instantiate_graph,
     )
 
     with ModuleImages() as images:
@@ -51,9 +51,9 @@ def forward():
                 session=SimpleNamespace(capture=nullcontext), device=0, stream=stream,
                 pool=storage.graph_pool.id, warmup_storage=storage, retain_graphs=True, compile_decoder=False)
             graph = graphs["decode"]
-            assert isinstance(graph, RecordedExportGraph)
-            assert graph.owners == (forward,)
-            assert storage.graph_pool.use_count() == 1
+            assert isinstance(graph, torch.cuda.CUDAGraph)
+            assert graph._keep_graph
+            assert storage.graph_pool.use_count() >= 1
             before = [(segment["address"], segment["total_size"])
                       for segment in storage.graph_pool.snapshot(include_traces=False)]
             torch.cuda.empty_cache()
@@ -70,8 +70,9 @@ def forward():
                 tuning_graphs={"tune": TuningGraph(tuner, ("selector",), "status")})
             manifest = bundle.export_manifest()
             assert all("handle" not in value for value in (*manifest["graphs"], *manifest["tuning_graphs"]))
-            with pytest.raises(TypeError, match="CUDAGraph"):
-                bundle.manifest()
+            runtime_manifest = bundle.manifest()
+            assert runtime_manifest["graphs"][0]["kind"] == "decode"
+            assert "handle" in runtime_manifest["graphs"][0]
             path = tmp_path / "rank.json"
             allocations = [Allocation(name, tensor.data_ptr(), tensor.nbytes)
                            for name, tensor in buffers.items()]