Skip to content

Commit 3919c1d

Browse files
committed
Retain CUDA graph resources throughout core export
Captured node pointers do not retain the Python owners of pools and kernel resources. Releasing graph handles before serialization can therefore leave export without the complete resource lifetimes used during capture. Retain each captured template and instantiated executable through export. Attach external pool and kernel owners to the graph's synchronized release path. Extend ownership tests for reset, collection, and instantiation failure, and document the executable memory retained during preparation. This keeps export resources alive until graph release without changing the serialized core format or native loading contract. The CPU capture metadata suite passes all 12 tests. The CUDA export test is skipped, and a full Kimi export has not been run with this change. Applies NVIDIA/hyperflare-inference PR #262 commits bb3ce1e and ec6190c. Design-Signed-Off-By: Codex GPT-6; design-review=ed22
1 parent 5fb495d commit 3919c1d

7 files changed

Lines changed: 116 additions & 63 deletions

File tree

‎docs/preparation-development.md‎

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -96,16 +96,17 @@ storage, downloads retained bytes, serializes resources, and owns asset writes,
9696
hashing, zstd level-3 executable compression, signing, conversion, rank payload
9797
transfer, and publication. Each executable resource is one zstd level-3 frame
9898
with a signed decoded length and compressed-chunk commitments.
99-
Export capture copies each graph's nodes, argument bytes, attributes, and
100-
annotated edges into native storage without instantiating an executable.
101-
Ignition retains the external CUDA memory pools and captured kernel owners,
102-
then resets the capture object to release its template. Diagnostic captures
103-
retain executables; their post-instantiation hooks can copy the same metadata
104-
before PyTorch releases their templates. Native export resolves allocations
105-
and validates bindings against the copied metadata. Keep captured buffers,
106-
contexts, copied metadata, pool owners, kernel owners, and any diagnostic
107-
executables alive until native work confirms a drain; an unsuccessful drain
108-
retains their owners.
99+
Ordinary export retains CUDA templates and instantiated executables, including
100+
their external memory pools and captured kernel owners, until native export
101+
has drained. Budget preparation memory for those executables across the serving
102+
graph catalog. Native export copies nodes, argument bytes, attributes, and
103+
annotated edges, then resolves allocations and validates their bindings.
104+
Diagnostic captures can copy the same metadata through post-instantiation hooks
105+
before PyTorch releases their templates. The core stores graph descriptions,
106+
not the Python executables; Helios creates executable instances at load time.
107+
Keep captured buffers, contexts, metadata, pool owners, kernel owners, and graph
108+
handles alive until native work confirms a drain; an unsuccessful drain retains
109+
their owners.
109110

110111
Native transfer pools size file-buffer concurrency from available host memory
111112
and preparation rank count, accounting for pinned CUDA download storage. Each

‎docs/runtime-core-design.md‎

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

‎ignition/preparation/capture.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ def capture_forwards(functions, signatures, *, session, device, stream, pool,
4848
serving replays variants serially and may choose them in any order.
4949
All eager warmups precede capture so their transient workspaces do not
5050
coexist with retained graph intermediates.
51-
``retain_graphs`` records export metadata without executables when
51+
``retain_graphs`` keeps templates and executables for export when
5252
``warmup_storage`` supplies external CUDA pool owners.
5353
``before_run`` restores caller-owned trial inputs outside compiled warmup
5454
and graph capture; it is unused for the ordinary serving catalog.

‎ignition/preparation/graphs.py‎

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

‎ignition/preparation/runtime_core.py‎

Lines changed: 9 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -190,26 +190,20 @@ def reset(self):
190190

191191

192192
def record_export_graph(graph, storage, *, owners=()):
193-
"""Copy an uninstantiated graph before releasing its CUDA template."""
194-
import flare_heliopause
195-
from .module_images import ModuleImages
193+
"""Retain a captured template, its executable and external export resources.
196194
197-
images = ModuleImages.current()
198-
if images is None:
199-
raise RuntimeError("graph metadata recording requires an active native module recorder")
195+
Owners must not reference the graph: Torch's destruction finalizer retains
196+
them independently. Native replay callers must drain before resetting graphs.
197+
"""
200198
pools = (storage.pool, storage.graph_pool)
201199
owned = {pool.id for pool in pools}
202200
if not set(graph.pools()) <= owned:
203201
raise ValueError("export graph uses a CUDA pool without an external owner")
204-
context, _ = flare_heliopause.current_context()
205-
recorded = images.native.record_template(int(graph.raw_cuda_graph()), 1)
206-
owner = RecordedExportGraph(images.native, recorded, context, pools, owners)
207-
try:
208-
graph.reset()
209-
except BaseException:
210-
owner.reset()
211-
raise
212-
return owner
202+
# Device pointers in captured nodes do not retain their Python resource owners.
203+
# Release these references only when the graph is reset or collected.
204+
graph.retain_object((pools, tuple(owners)), synchronize_before_release=True)
205+
instantiate_graph(graph)
206+
return graph
213207

214208

215209
def capture_graph(*, retain_graphs=False):

‎tests/preparation/test_capture_graph_metadata.py‎

Lines changed: 74 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
class Graph:
2525
register_post_instantiate_hook = CUDAGraph.register_post_instantiate_hook
2626
register_destroy_callback = CUDAGraph.register_destroy_callback
27+
retain_object = CUDAGraph.retain_object
2728

2829
def __init__(self, keep_graph=False):
2930
self.keep_graph = keep_graph
@@ -187,7 +188,7 @@ def graph_scope(graph, **kwargs):
187188
assert recorder.events == ["record", "record"]
188189

189190

190-
def test_export_metadata_outlives_template_and_retains_external_owners(monkeypatch, recorder):
191+
def test_export_instantiates_and_retains_executable(monkeypatch, recorder):
191192
import flare_heliopause
192193

193194
monkeypatch.setattr(flare_heliopause, "current_context", lambda: (100, 0))
@@ -196,21 +197,75 @@ def test_export_metadata_outlives_template_and_retains_external_owners(monkeypat
196197
graph.pools = lambda: [(0, 2)]
197198
graph.capture_end()
198199
kernel_owner = SimpleNamespace()
199-
owner = record_export_graph(graph, storage, owners=(kernel_owner,))
200-
assert graph.template == graph.executable == graph.instantiations == 0
201-
assert owner.pools == (storage.pool, storage.graph_pool)
202-
assert owner.owners == (kernel_owner,)
203-
assert _graphs({"decode": owner}, recorder=recorder) == [dict(name="decode", recorded=1)]
204-
with pytest.raises(ValueError, match="another native module recorder"):
205-
_graphs({"decode": owner}, recorder=Recorder())
206-
assert recorder.templates == {1: 11}
207-
owner.reset()
208-
owner.reset()
209-
assert owner.pools == owner.owners == ()
200+
returned = record_export_graph(graph, storage, owners=(kernel_owner,))
201+
assert returned is graph
202+
assert graph.instantiations == 1
203+
assert graph.template == 11
204+
assert graph.executable == 12
205+
assert _graphs({"decode": graph}) == [dict(name="decode", graph=11, executable=12)]
210206
assert not recorder.templates
211-
assert recorder.events == ["record-template", "release-template"]
212-
with pytest.raises(ValueError, match="released"):
213-
_graphs({"decode": owner}, recorder=recorder)
207+
assert recorder.events == []
208+
209+
210+
@pytest.mark.parametrize("release", ("reset", "collection"))
211+
def test_export_graph_retains_external_owners_until_release(recorder, release):
212+
class Owner:
213+
pass
214+
215+
first_pool, second_pool, kernel_owner = Owner(), Owner(), Owner()
216+
first_pool.id, second_pool.id = (0, 1), (0, 2)
217+
storage = SimpleNamespace(pool=first_pool, graph_pool=second_pool)
218+
references = tuple(weakref.ref(owner) for owner in (first_pool, second_pool, kernel_owner))
219+
graph = Graph(keep_graph=True)
220+
graph.pools = lambda: [(0, 2)]
221+
graph.capture_end()
222+
returned = record_export_graph(graph, storage, owners=(kernel_owner,))
223+
assert returned is graph
224+
graph_reference = weakref.ref(graph)
225+
del first_pool, second_pool, kernel_owner, storage
226+
gc.collect()
227+
assert all(reference() is not None for reference in references)
228+
229+
drained = []
230+
class Stream:
231+
def synchronize(self):
232+
assert all(reference() is not None for reference in references)
233+
drained.append(True)
234+
235+
graph._retained.replay_streams.add(Stream())
236+
if release == "reset":
237+
graph.reset()
238+
assert all(reference() is None for reference in references)
239+
assert graph_reference() is graph
240+
del graph, returned
241+
gc.collect()
242+
assert graph_reference() is None
243+
assert all(reference() is None for reference in references)
244+
assert drained == [True]
245+
246+
247+
def test_export_instantiation_failure_retains_owners_until_reset(monkeypatch, recorder):
248+
class Owner:
249+
pass
250+
251+
owner = Owner()
252+
reference = weakref.ref(owner)
253+
storage = SimpleNamespace(pool=SimpleNamespace(id=(0, 1)), graph_pool=SimpleNamespace(id=(0, 2)))
254+
graph = Graph(keep_graph=True)
255+
graph.pools = lambda: [(0, 2)]
256+
graph.capture_end()
257+
258+
def fail_instantiation():
259+
raise RuntimeError("graph instantiation failed")
260+
261+
monkeypatch.setattr(graph, "instantiate", fail_instantiation)
262+
with pytest.raises(RuntimeError, match="graph instantiation failed"):
263+
record_export_graph(graph, storage, owners=(owner,))
264+
del owner
265+
gc.collect()
266+
assert reference() is not None
267+
graph.reset()
268+
assert reference() is None
214269

215270

216271
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
221276
with pytest.raises(ValueError, match="without an external owner"):
222277
record_export_graph(graph, storage)
223278
assert graph.template == 11
279+
assert graph.instantiations == 0 and not graph._retained.objects
224280
assert not recorder.templates
225281

226282

@@ -258,7 +314,7 @@ def test_export_resources_accept_recorded_graphs_and_executable_tuners(monkeypat
258314
bundle._validate_resources(export_graphs=True)
259315

260316

261-
def test_export_preserves_collective_consumers_after_replacing_cuda_graph(monkeypatch, recorder):
317+
def test_export_preserves_collective_consumer_associations(monkeypatch, recorder):
262318
import flare_heliopause
263319
from ignition.preparation.startup_tuning import record_consumer
264320

@@ -283,5 +339,6 @@ def graph_scope(graph, **kwargs):
283339
session=session, device="cpu", stream=stream, pool=storage.graph_pool.id, compile_decoder=False,
284340
retain_graphs=True, warmup_storage=storage)
285341
graph = graphs["decode"]
286-
assert isinstance(graph, RecordedExportGraph)
342+
assert isinstance(graph, Graph)
343+
assert graph.instantiations == 1
287344
assert registry.captures == {id(graph): {id(collective)}}

‎tests/preparation/test_recorded_export_graph_cuda.py‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
# SPDX-FileCopyrightText: Copyright (c) 2026 Flare community contributors
22
# SPDX-License-Identifier: Apache-2.0
33

4-
"""Export graph metadata retains CUDA storage without graph handles."""
4+
"""Retained export graphs preserve storage and serialize their execution metadata."""
55

66
import os
77
from pathlib import Path
@@ -12,7 +12,7 @@
1212

1313

1414
@pytest.mark.skipif(os.environ.get("FLARE_RUN_CUDA_TESTS") != "1", reason="requires explicit CUDA test opt-in")
15-
def test_export_capture_releases_cuda_handles_and_preserves_storage(tmp_path):
15+
def test_export_capture_retains_cuda_handles_and_preserves_storage(tmp_path):
1616
subprocess.run([sys.executable, "-c",
1717
"from pathlib import Path; import sys; "
1818
"from tests.preparation.test_recorded_export_graph_cuda import check_export; "
@@ -30,7 +30,7 @@ def check_export(tmp_path):
3030
from ignition.preparation.core_file import read_sources
3131
from ignition.preparation.core_storage import CoreStorage
3232
from ignition.preparation.runtime_core import (
33-
Allocation, RecordedExportGraph, RuntimeCore, capture_graph, instantiate_graph,
33+
Allocation, RuntimeCore, capture_graph, instantiate_graph,
3434
)
3535

3636
with ModuleImages() as images:
@@ -51,9 +51,9 @@ def forward():
5151
session=SimpleNamespace(capture=nullcontext), device=0, stream=stream,
5252
pool=storage.graph_pool.id, warmup_storage=storage, retain_graphs=True, compile_decoder=False)
5353
graph = graphs["decode"]
54-
assert isinstance(graph, RecordedExportGraph)
55-
assert graph.owners == (forward,)
56-
assert storage.graph_pool.use_count() == 1
54+
assert isinstance(graph, torch.cuda.CUDAGraph)
55+
assert graph._keep_graph
56+
assert storage.graph_pool.use_count() >= 1
5757
before = [(segment["address"], segment["total_size"])
5858
for segment in storage.graph_pool.snapshot(include_traces=False)]
5959
torch.cuda.empty_cache()
@@ -70,8 +70,9 @@ def forward():
7070
tuning_graphs={"tune": TuningGraph(tuner, ("selector",), "status")})
7171
manifest = bundle.export_manifest()
7272
assert all("handle" not in value for value in (*manifest["graphs"], *manifest["tuning_graphs"]))
73-
with pytest.raises(TypeError, match="CUDAGraph"):
74-
bundle.manifest()
73+
runtime_manifest = bundle.manifest()
74+
assert runtime_manifest["graphs"][0]["kind"] == "decode"
75+
assert "handle" in runtime_manifest["graphs"][0]
7576
path = tmp_path / "rank.json"
7677
allocations = [Allocation(name, tensor.data_ptr(), tensor.nbytes)
7778
for name, tensor in buffers.items()]

0 commit comments

Comments
 (0)