@@ -231,36 +231,42 @@ def session_setup():
231231def init_cuda (request ):
232232 # TODO: rename this to e.g. init_context
233233 device = Device (0 )
234+ # Everything that can raise after set_current() runs inside the try, so the
235+ # finally pops this context; the leftover warning below is escalated to an
236+ # error. A failed set_current() leaves nothing for the finally to undo.
237+ # live_at_setup stays None until the setup snapshot is recorded.
234238 device .set_current ()
235-
236- # Set option to avoid spin-waiting on synchronization.
237- if int (os .environ .get ("CUDA_CORE_TEST_BLOCKING_SYNC" , 0 )) != 0 :
238- handle_return (
239- driver .cuDevicePrimaryCtxSetFlags (device .device_id , driver .CUctx_flags .CU_CTX_SCHED_BLOCKING_SYNC )
240- )
241-
242- # Leftover check: growth since the previous init_cuda test means something
243- # in between left an owned resource alive. Collect first to sweep a cyclic
244- # leftover, and before the counts at setup are recorded so they are reflected there.
245- if _mr_expected_live is not None and any (now > exp for now , exp in zip (_mr_live_at_item_start , _mr_expected_live )):
246- gc .collect ()
247- now_counts = _live_resource_counts ()
248- held = tuple (max (n - e , 0 ) for n , e in zip (now_counts , _mr_expected_live ))
249- excess = tuple (s - e for s , e in zip (_mr_live_at_item_start , _mr_expected_live ))
250- seed = getattr (request .config .option , "randomly_seed" , None )
251- seed_msg = f" (pytest-randomly seed: { seed } )" if seed is not None else ""
252- warnings .warn (
253- f"Owned resources were left alive by tests that ran since { _mr_last_init_cuda_nodeid } "
254- f"({ _format_counts (excess )} ; still alive after a collect: { _format_counts (held )} ). "
255- f"A test or fixture that constructs a memory resource must request init_cuda.{ seed_msg } " ,
256- UnreleasedMemoryResourceWarning ,
257- stacklevel = 2 ,
258- )
259-
260- # Counts at setup. The device default pool is a borrowed reference and is
261- # not counted, so no cached-default bookkeeping is needed.
262- live_at_setup = _live_resource_counts ()
239+ live_at_setup = None
263240 try :
241+ # Set option to avoid spin-waiting on synchronization.
242+ if int (os .environ .get ("CUDA_CORE_TEST_BLOCKING_SYNC" , 0 )) != 0 :
243+ handle_return (
244+ driver .cuDevicePrimaryCtxSetFlags (device .device_id , driver .CUctx_flags .CU_CTX_SCHED_BLOCKING_SYNC )
245+ )
246+
247+ # Leftover check: growth since the previous init_cuda test means something
248+ # in between left an owned resource alive. Collect first to sweep a cyclic
249+ # leftover, and before the counts at setup are recorded so they are reflected there.
250+ if _mr_expected_live is not None and any (
251+ now > exp for now , exp in zip (_mr_live_at_item_start , _mr_expected_live )
252+ ):
253+ gc .collect ()
254+ now_counts = _live_resource_counts ()
255+ held = tuple (max (n - e , 0 ) for n , e in zip (now_counts , _mr_expected_live ))
256+ excess = tuple (s - e for s , e in zip (_mr_live_at_item_start , _mr_expected_live ))
257+ seed = getattr (request .config .option , "randomly_seed" , None )
258+ seed_msg = f" (pytest-randomly seed: { seed } )" if seed is not None else ""
259+ warnings .warn (
260+ f"Owned resources were left alive by tests that ran since { _mr_last_init_cuda_nodeid } "
261+ f"({ _format_counts (excess )} ; still alive after a collect: { _format_counts (held )} ). "
262+ f"A test or fixture that constructs a memory resource must request init_cuda.{ seed_msg } " ,
263+ UnreleasedMemoryResourceWarning ,
264+ stacklevel = 2 ,
265+ )
266+
267+ # Counts at setup. The device default pool is a borrowed reference and is
268+ # not counted, so no cached-default bookkeeping is needed.
269+ live_at_setup = _live_resource_counts ()
264270 yield device
265271 finally :
266272 # pytest holds fixture values in item.funcargs (a pytest-internal dict)
@@ -271,21 +277,24 @@ def init_cuda(request):
271277 # Only an owned resource in a reference cycle outlives the test frame, so
272278 # collect only if the counts differ from their values at setup. A lower
273279 # count means an older resource was released, which could mask a leak
274- # in this test: collect in that case too.
275- live_after = _live_resource_counts ()
276- if live_after != live_at_setup :
280+ # in this test: collect in that case too. Skip the comparison when setup
281+ # aborted before the snapshot (the leftover warning raises).
282+ live_after = _live_resource_counts () if live_at_setup is not None else None
283+ if live_after is not None and live_after != live_at_setup :
277284 gc .collect ()
278285 # Force any pool/allocation whose only remaining reference was a local
279286 # in this test's frame to actually get destroyed now, then drain the
280287 # context so the stream-ordered frees that destruction enqueues retire
281288 # before the next test runs (issue #2381). gc.collect() above must run
282289 # first: cuCtxSynchronize alone cannot drain frees that were never
283290 # enqueued because their owning object had not been collected yet.
284- # With pytest-run-parallel this runs after worker join.
291+ # With pytest-run-parallel this runs after worker join. This also runs
292+ # when setup raises, so an escalated leftover warning cannot leave the
293+ # context current.
285294 driver .cuCtxSynchronize ()
286295 _ = _device_unset_current ()
287296 # After the cleanup, so an escalated warning cannot skip it.
288- if any (now < snap for now , snap in zip (live_after , live_at_setup )):
297+ if live_after is not None and any (now < snap for now , snap in zip (live_after , live_at_setup )):
289298 released = tuple (s - n for s , n in zip (live_at_setup , live_after ))
290299 warnings .warn (
291300 f"{ request .node .nodeid } released owned resources that existed before it started "
0 commit comments