@@ -56,7 +56,7 @@ def __init__(self, **kwargs: Any) -> None:
5656 self .replays = 0
5757 super ().__init__ (
5858 _Transformer ().eval (),
59- mode = kwargs .pop ("mode" , "bucketed " ),
59+ mode = kwargs .pop ("mode" , "exact " ),
6060 window = 256 ,
6161 pad_token_id = 0 ,
6262 clock = lambda : self .now ,
@@ -92,18 +92,28 @@ def room(monkeypatch: pytest.MonkeyPatch) -> dict[str, Any]:
9292 return state
9393
9494
95- @pytest .mark .parametrize ("mode" , ["exact" , "bucketed" ])
96- def test_shape_policy_is_bounded_and_preserves_batch_rows (mode : str ) -> None :
95+ def test_shape_policy_is_bounded_and_preserves_exact_dimensions () -> None :
9796 for rows in (1 , 2 , 3 , 8 , 31 , 64 , 1024 ):
9897 for length in (1 , 15 , 31 , 32 , 33 , 100 , 255 , 256 , 257 ):
99- key = graphs .graph_shape (rows , length , mode = mode , window = 256 , max_tokens = 1024 )
98+ key = graphs .graph_shape (rows , length , mode = "exact" , window = 256 , max_tokens = 1024 )
10099 if key is not None :
101- assert key [ 0 ] == rows
102- assert length <= key [ 1 ] <= 256
100+ assert key == ( rows , length )
101+ assert length <= 256
103102 assert key [0 ] * key [1 ] <= 1024
104- assert key [1 ] == (length if mode == "exact" else - (- length // 32 ) * 32 )
105103 for shape in ((0 , 4 ), (1 , 0 ), (1 , 257 ), (9 , 256 )):
106- assert graphs .graph_shape (* shape , mode = mode , window = 256 , max_tokens = 1024 ) is None
104+ assert graphs .graph_shape (* shape , mode = "exact" , window = 256 , max_tokens = 1024 ) is None
105+
106+
107+ def test_bucketed_mode_is_rejected_before_model_or_forward_access () -> None :
108+ class UnavailableModel :
109+ def __getattr__ (self , name : str ) -> Any :
110+ pytest .fail (f"invalid mode accessed model attribute { name } " )
111+
112+ assert graphs .GRAPH_MODES == ("off" , "exact" )
113+ for length in (31 , 32 , 33 , 256 ):
114+ assert graphs .graph_shape (1 , length , mode = "bucketed" , window = 256 , max_tokens = 1024 ) is None
115+ with pytest .raises (ValueError , match = "exact mode" ):
116+ graphs .EncoderGraphRunner (UnavailableModel (), mode = "bucketed" , window = 256 , pad_token_id = 0 )
107117
108118
109119def test_second_sighting_records_and_cache_replays (room : dict [str , Any ]) -> None :
@@ -119,6 +129,7 @@ def test_second_sighting_records_and_cache_replays(room: dict[str, Any]) -> None
119129
120130def test_no_headroom_global_recording_lock_and_credit_limits (room : dict [str , Any ]) -> None :
121131 runner = _PolicyRunner ()
132+ assert runner .offer () is None # exact mode's first sighting
122133 room ["headroom" ] = False
123134 assert runner .offer () is None
124135 room ["headroom" ] = True
@@ -128,12 +139,13 @@ def test_no_headroom_global_recording_lock_and_credit_limits(room: dict[str, Any
128139 assert runner .offer () is None
129140 runner .now = 2
130141 assert runner .offer () is not None
131- assert set (runner .stats .eager ) == {"no_headroom" , "busy" , "recording_paused" }
142+ assert set (runner .stats .eager ) == {"first_shape" , " no_headroom" , "busy" , "recording_paused" }
132143 runner .close ()
133144
134145
135146def test_over_budget_capture_is_not_retained (room : dict [str , Any ]) -> None :
136147 runner = _PolicyRunner ()
148+ assert runner .offer () is None # exact mode's first sighting
137149 room ["budget" ] = 1
138150 assert runner .offer () is None
139151 assert runner .held_bytes == runner .graph_count == 0
@@ -202,6 +214,7 @@ def accounted_capture(key: graphs.Key, device: torch.device) -> graphs._Graph:
202214 return entry
203215
204216 monkeypatch .setattr (runner , "_capture" , accounted_capture )
217+ assert runner .offer () is None # exact mode's first sighting
205218 assert runner .offer () is None
206219 assert runner .graph_count == runner .held_bytes == runner .replays == 0
207220 assert runner .stats .recording_failures == 1
@@ -229,9 +242,10 @@ def accounted_capture(key: graphs.Key, device: torch.device) -> graphs._Graph:
229242
230243 monkeypatch .setattr (runner , "_capture" , accounted_capture )
231244 room ["budget" ] = 4096
245+ assert runner .offer () is None # exact mode's first sighting
232246 assert runner .offer () is None
233247 assert runner .graph_count == runner .held_bytes == runner .replays == 0
234- assert runner .stats .eager == {"budget_full" : 1 }
248+ assert runner .stats .eager == {"first_shape" : 1 , " budget_full" : 1 }
235249 runner .close ()
236250
237251
@@ -240,13 +254,15 @@ def test_capture_failure_restores_eager_and_bounds_retries(room: dict[str, Any],
240254 runner = _PolicyRunner ()
241255 runner .capture_error = torch .OutOfMemoryError ("allocation failed" ) if oom else RuntimeError ("not graphable" )
242256 for length in (1 , 33 , 65 , 97 ):
257+ assert runner .offer (length = length ) is None # first sighting, or still paused/disabled
243258 assert runner .offer (length = length ) is None
244259 assert runner .held_bytes == 0
245260 if oom :
246261 assert runner .stats .drops == 1
247262 assert not runner ._failed
248263 runner .now = 61
249264 runner .capture_error = None
265+ assert runner .offer () is None # first sighting of this shape
250266 assert runner .offer () is not None
251267 else :
252268 assert len (runner ._failed ) == graphs ._MAX_FAILURES
@@ -257,6 +273,7 @@ def test_capture_failure_restores_eager_and_bounds_retries(room: dict[str, Any],
257273
258274def test_existing_replay_oom_releases_cache_and_propagates (room : dict [str , Any ]) -> None :
259275 runner = _PolicyRunner ()
276+ assert runner .offer () is None # exact mode's first sighting
260277 assert runner .offer () is not None
261278 runner .replay_error = torch .OutOfMemoryError ("allocation failed" )
262279 with pytest .raises (torch .OutOfMemoryError ):
@@ -358,6 +375,7 @@ def test_replay_copies_and_slices_outputs_on_the_owned_stream(monkeypatch: pytes
358375
359376def test_pressure_clear_waits_before_dropping_graph_storage (room : dict [str , Any ]) -> None :
360377 runner = _PolicyRunner ()
378+ assert runner .offer () is None # exact mode's first sighting
361379 assert runner .offer () is not None
362380 observed : list [int ] = []
363381 runner ._stream = SimpleNamespace (synchronize = lambda : observed .append (runner .graph_count ))
@@ -368,7 +386,7 @@ def test_pressure_clear_waits_before_dropping_graph_storage(room: dict[str, Any]
368386 runner .close ()
369387
370388
371- @pytest .mark .parametrize ("mode" , ["off" , False , "exact" , "bucketed" ])
389+ @pytest .mark .parametrize ("mode" , ["off" , False , "exact" ])
372390def test_operator_option_and_default_eager_path (mode : str | bool , monkeypatch : pytest .MonkeyPatch ) -> None :
373391 reject_unknown_loadtime_options (GLiFormerAdapter , {"cuda_graphs" : mode }, model_name = "tiny" )
374392 adapter = GLiFormerAdapter ("tiny" , cuda_graphs = mode )
@@ -382,18 +400,19 @@ def test_operator_option_and_default_eager_path(mode: str | bool, monkeypatch: p
382400 assert adapter .release_optional_memory () == 0
383401
384402
385- @pytest .mark .parametrize ("mode" , [True , None , 1 , "on" , "EXACT" , []])
403+ @pytest .mark .parametrize ("mode" , [True , None , 1 , "on" , "EXACT" , "bucketed" , []])
386404def test_invalid_modes_fail_before_loading (mode : Any ) -> None :
387405 with pytest .raises (ValueError , match = "cuda_graphs" ):
388406 GLiFormerAdapter ("tiny" , cuda_graphs = mode )
389407
390408
391409def test_pressure_release_keeps_weights_and_unload_restores_forward (room : dict [str , Any ]) -> None :
392410 runner = _PolicyRunner ()
393- adapter = GLiFormerAdapter ("tiny" , cuda_graphs = "bucketed " )
411+ adapter = GLiFormerAdapter ("tiny" , cuda_graphs = "exact " )
394412 transformer = runner ._transformer
395413 adapter ._model = transformer
396414 adapter ._graphs = runner
415+ assert runner .offer () is None # exact mode's first sighting
397416 assert runner .offer () is not None
398417 held = runner .held_bytes
399418 assert adapter .has_releasable_memory ()
@@ -534,28 +553,37 @@ def observed_helper(*args: Any) -> torch.Tensor:
534553 ids = torch .arange (7 ).reshape (1 , 7 ).long ()
535554 mask = torch .ones_like (ids )
536555 outputs : list [torch .Tensor ] = []
556+ baseline_ready , capture_ready = threading .Event (), threading .Event ()
537557
538558 def eager_call () -> None :
539559 with torch .inference_mode ():
540560 outputs .append (runner ._forward (input_ids = ids , attention_mask = mask ))
561+ baseline_ready .set ()
562+ assert capture_ready .wait (timeout = 5 )
563+ outputs .append (runner ._forward (input_ids = ids , attention_mask = mask ))
541564
565+ thread = threading .Thread (target = eager_call )
542566 try :
543567 with torch .inference_mode ():
544568 expected = runner ._forward (input_ids = ids , attention_mask = mask )
569+ thread .start ()
570+ assert baseline_ready .wait (timeout = 5 )
545571 table = runner ._build_relative_pos (torch .empty ((1 , 7 , 1 )))
546572 with runner ._layout_attention_scales (torch .device ("cpu" )), runner ._recording (table ):
547573 calls .clear ()
548574 assert torch .equal (runner ._forward (input_ids = ids , attention_mask = mask ), expected )
549575 assert calls == []
550- thread = threading .Thread (target = eager_call )
551- thread .start ()
576+ capture_ready .set ()
552577 thread .join (timeout = 5 )
553578 assert not thread .is_alive ()
554579 assert calls == [thread .ident ] * len (runner ._layers .layer )
555580 assert deberta_2d .scaled_size_sqrt is observed_helper
556- assert len (outputs ) == 1
557- assert torch .equal (outputs [0 ], expected )
581+ assert len (outputs ) == 2
582+ assert torch .equal (outputs [1 ], outputs [ 0 ] )
558583 finally :
584+ capture_ready .set ()
585+ if thread .ident is not None :
586+ thread .join (timeout = 5 )
559587 runner .close ()
560588
561589
@@ -846,16 +874,14 @@ def _assert_output_close(actual: Any, expected: Any, tolerance: float) -> None:
846874 assert actual == expected
847875
848876
849- @pytest .mark .parametrize ("mode" , ["exact" , "bucketed" ])
850- def test_real_all_task_heads_and_usage_survive_encoder_only_padding (
877+ def test_real_all_task_heads_and_usage_survive_exact_encoder_replay (
851878 tiny_checkpoint : Path ,
852- mode : str ,
853879 room : dict [str , Any ],
854880 monkeypatch : pytest .MonkeyPatch ,
855881) -> None :
856882 """Exercise the actual SIE load/inference/unload path on a generated checkpoint."""
857883 baseline = GLiFormerAdapter (tiny_checkpoint , max_seq_length = 128 , compute_precision = "float32" )
858- candidate = GLiFormerAdapter (tiny_checkpoint , max_seq_length = 128 , compute_precision = "float32" , cuda_graphs = mode )
884+ candidate = GLiFormerAdapter (tiny_checkpoint , max_seq_length = 128 , compute_precision = "float32" , cuda_graphs = "exact" )
859885 monkeypatch .setattr (adapter_module , "unsupported_reason" , lambda * _ : None )
860886 monkeypatch .setattr (adapter_module , "EncoderGraphRunner" , _EagerReplay )
861887 items = [Item (text = "Alice works at Acme." ), Item (text = "Bob likes coffee." )]
@@ -869,19 +895,18 @@ def observe(_module: Any, _inputs: Any, output: Any) -> None:
869895
870896 hook = candidate ._model .model .register_forward_hook (observe )
871897 try :
872- tolerance = 0 if mode == "exact" else 1e-5 # synthetic fp32 padding, not a pretrained/GPU bound
873898 for task in _TASKS :
874899 expected = baseline .extract (items , ** task )
875900 candidate .extract (items , ** task ) # exact mode's first sighting
876901 candidate .extract (items , ** task ) # capture, then test an ordinary replay
877902 actual = candidate .extract (items , ** task )
878903 for field in ("entities" , "classifications" , "relations" , "data" , "errors" , "input_token_counts" ):
879- _assert_output_close (getattr (actual , field ), getattr (expected , field ), tolerance )
904+ _assert_output_close (getattr (actual , field ), getattr (expected , field ), 0 )
880905 eager_dense = baseline .encode (items , ["dense" ]).dense
881906 candidate .encode (items , ["dense" ])
882907 candidate .encode (items , ["dense" ])
883908 dense = candidate .encode (items , ["dense" ]).dense
884- torch .testing .assert_close (torch .from_numpy (dense ), torch .from_numpy (eager_dense ), rtol = 0 , atol = tolerance )
909+ torch .testing .assert_close (torch .from_numpy (dense ), torch .from_numpy (eager_dense ), rtol = 0 , atol = 0 )
885910 assert observed
886911 assert {"ner_logits" , "cat_logits" , "joint_rel_logits" , "structuring_logits" } <= set ().union (
887912 * (set (row ) for row in observed )
@@ -903,14 +928,13 @@ def observe(_module: Any, _inputs: Any, output: Any) -> None:
903928
904929@pytest .mark .skipif (not torch .cuda .is_available (), reason = "actual CUDA graph parity needs a GPU" )
905930@pytest .mark .parametrize ("layout" , [False , True ], ids = ["deberta" , "layout-deberta" ])
906- @pytest .mark .parametrize ("mode" , ["exact" , "bucketed" ])
907- def test_actual_cuda_capture_matches_eager_padded_shape_and_owns_output (layout : bool , mode : str ) -> None :
931+ def test_actual_cuda_capture_matches_eager_shape_and_owns_output (layout : bool ) -> None :
908932 transformer = _tiny_transformer (layout ).to (device = "cuda" , dtype = torch .float16 )
909933 original = transformer .forward
910934 ids = torch .arange (14 , device = "cuda" ).reshape (2 , 7 ).long ()
911935 mask = torch .ones_like (ids )
912936 mask [1 , - 2 :] = 0
913- runner = graphs .EncoderGraphRunner (transformer , mode = mode , window = 64 , pad_token_id = 0 )
937+ runner = graphs .EncoderGraphRunner (transformer , mode = "exact" , window = 64 , pad_token_id = 0 )
914938 try :
915939 with torch .inference_mode ():
916940 key = runner .key (* ids .shape )
0 commit comments