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 aa1cf1e6..8c95958e 100644 --- a/ignition/preparation/runtime_core.py +++ b/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): diff --git a/tests/preparation/test_capture_graph_metadata.py b/tests/preparation/test_capture_graph_metadata.py index a4423fb5..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 @@ -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)}} diff --git a/tests/preparation/test_recorded_export_graph_cuda.py b/tests/preparation/test_recorded_export_graph_cuda.py index aad112cd..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; " @@ -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()]