From bb3ce1ead4d53a9fd2dfca7f78d7737d237d4e0b Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 11 Nov 2026 05:33:13 +0000 Subject: [PATCH 1/2] Retain executable graphs for repeated serving replays The graph export change in b455be5e replaced an instantiated executable with reset template metadata. A serving package produced by that path fails on its second full-model request: `cuGraphLaunch` reports CUDA_ERROR_ILLEGAL_ADDRESS after the first request completes. Exporting the same model, layouts, and weights through the executable-retaining path preserves valid replay. Restore `instantiate_graph` in `record_export_graph` and return the executable graph while retaining the existing graph-pool ownership validation. This keeps the graph and its external storage alive for repeated launches instead of resetting them into template-only metadata. Update the metadata tests to assert that the returned object is the original instantiated graph, that no `record_template` call is made, and that collective consumers keep their associations. The CUDA export fixture validates executable graph metadata and requires that graph-pool storage remains valid after `empty_cache`. Validation: all 9 graph metadata tests and all 23 runtime-core/capture-readiness tests pass with CUDA disabled. The real CUDA export fixture passes. As an end-to-end comparison, a C4 export through the reset-template path fails on the second long request, while the same model through this fix completes four sequential long requests (frames 831, 1777, 2896, and 4110) with HTTP 200 on TP16. No claim is made about model quality or performance. Design-Signed-Off-By: Kimi (Kimi); design-review=cfcf Co-Authored-By: Kimi --- ignition/preparation/runtime_core.py | 21 +++++---------- .../test_capture_graph_metadata.py | 26 +++++++------------ .../test_recorded_export_graph_cuda.py | 13 +++++----- 3 files changed, 23 insertions(+), 37 deletions(-) diff --git a/ignition/preparation/runtime_core.py b/ignition/preparation/runtime_core.py index aa1cf1e6..25d7ce57 100644 --- a/ignition/preparation/runtime_core.py +++ b/ignition/preparation/runtime_core.py @@ -190,26 +190,17 @@ 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 + """Keep an instantiated graph valid for every serving replay. - images = ModuleImages.current() - if images is None: - raise RuntimeError("graph metadata recording requires an active native module recorder") + Replacing the executable with reset template metadata breaks repeated + full-model launches, so preserve the executable export path instead. + """ 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 + instantiate_graph(graph) + return graph def capture_graph(*, retain_graphs=False): diff --git a/tests/preparation/test_capture_graph_metadata.py b/tests/preparation/test_capture_graph_metadata.py index a4423fb5..b7320976 100644 --- a/tests/preparation/test_capture_graph_metadata.py +++ b/tests/preparation/test_capture_graph_metadata.py @@ -187,7 +187,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 +196,14 @@ 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 == [] def test_export_rejects_unowned_pool_without_destroying_template(monkeypatch, recorder): @@ -283,5 +276,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)}} diff --git a/tests/preparation/test_recorded_export_graph_cuda.py b/tests/preparation/test_recorded_export_graph_cuda.py index aad112cd..eb081dbf 100644 --- a/tests/preparation/test_recorded_export_graph_cuda.py +++ b/tests/preparation/test_recorded_export_graph_cuda.py @@ -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()] From ec6190c0bcbec9aa4218a8c454693dbc517fda56 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 11 Nov 2026 10:05:43 +0000 Subject: [PATCH 2/2] Retain export graph resource owners through reset and collection Returning an instantiated CUDA graph discarded the explicit kernel and pool owners supplied to record_export_graph. Device pointer capture does not retain those Python objects. Bind their lifetime to the graph with Torch's retained-object API before instantiation, and preserve pool checks. Retained resources survive instantiation failure and are released on reset or collection after tracked Python replay streams drain. Native callers still provide their existing explicit drain. Document the prohibition on owner-to-graph reference cycles and the preparation memory cost of retaining executables. Correct export documentation and misleading test names. This repairs the helper's lifetime contract, without changing serialized core format, model arithmetic or runtime policy. Three ownership regressions fail on the submitted implementation and pass with this change. The focused CPU capture suite passes 73 tests. No CUDA or end-to-end test was run. The reported second-request failure remains causally unqualified because the available comparison also changes the Kimi gather configuration. Design-Signed-Off-By: Codex GPT-6; design-review=cfcf Signed-off-by: Martin Vit --- docs/preparation-development.md | 21 +++--- docs/runtime-core-design.md | 20 +++--- ignition/preparation/capture.py | 2 +- ignition/preparation/graphs.py | 4 +- ignition/preparation/runtime_core.py | 9 ++- .../test_capture_graph_metadata.py | 65 ++++++++++++++++++- .../test_recorded_export_graph_cuda.py | 4 +- 7 files changed, 96 insertions(+), 29 deletions(-) diff --git a/docs/preparation-development.md b/docs/preparation-development.md index 8c5a8144..5cc70b39 100644 --- a/docs/preparation-development.md +++ b/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, diff --git a/docs/runtime-core-design.md b/docs/runtime-core-design.md index c075a2ac..f0c1b8e3 100644 --- a/docs/runtime-core-design.md +++ b/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. diff --git a/ignition/preparation/capture.py b/ignition/preparation/capture.py index ec07f343..437b9880 100644 --- a/ignition/preparation/capture.py +++ b/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. diff --git a/ignition/preparation/graphs.py b/ignition/preparation/graphs.py index a038a182..f6027b70 100644 --- a/ignition/preparation/graphs.py +++ b/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 diff --git a/ignition/preparation/runtime_core.py b/ignition/preparation/runtime_core.py index 25d7ce57..8c95958e 100644 --- a/ignition/preparation/runtime_core.py +++ b/ignition/preparation/runtime_core.py @@ -190,15 +190,18 @@ def reset(self): def record_export_graph(graph, storage, *, owners=()): - """Keep an instantiated graph valid for every serving replay. + """Retain a captured template, its executable and external export resources. - Replacing the executable with reset template metadata breaks repeated - full-model launches, so preserve the executable export path instead. + 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") + # 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 diff --git a/tests/preparation/test_capture_graph_metadata.py b/tests/preparation/test_capture_graph_metadata.py index b7320976..d81b2000 100644 --- a/tests/preparation/test_capture_graph_metadata.py +++ b/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 @@ -206,6 +207,67 @@ def test_export_instantiates_and_retains_executable(monkeypatch, 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): storage = SimpleNamespace(pool=SimpleNamespace(id=(0, 1)), graph_pool=SimpleNamespace(id=(0, 2))) graph = Graph(keep_graph=True) @@ -214,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 @@ -251,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 diff --git a/tests/preparation/test_recorded_export_graph_cuda.py b/tests/preparation/test_recorded_export_graph_cuda.py index eb081dbf..eb5376d7 100644 --- a/tests/preparation/test_recorded_export_graph_cuda.py +++ b/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; "