2424class 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
216271def 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 )}}
0 commit comments