[metrics] Report logical prefill token counts (#38566)

test/registered/unit/managers/test_prefill_adder.py

@@ -443,6 +443,13 @@ class TestPrefillAdder(CustomTestCase):
         req1.full_untruncated_fill_ids = list(range(56))
         req1.last_node = MagicMock()
         req1.sampling_params.ignore_eos = False
+        # add_one_req reads req.extend_range.length after set_extend_range;
+        # emulate the real Req writer (a spec=Req mock lacks the attribute).
+        req1.set_extend_range = MagicMock(
+            side_effect=lambda start, end: setattr(
+                req1, "extend_range", Range(start, end)
+            )
+        )
 
         result1 = adder.add_one_req(
             req1, has_chunked_req=False, truncation_align_size=None
@@ -476,6 +483,11 @@ class TestPrefillAdder(CustomTestCase):
         req2.full_untruncated_fill_ids = list(range(56))
         req2.last_node = MagicMock()
         req2.sampling_params.ignore_eos = False
+        req2.set_extend_range = MagicMock(
+            side_effect=lambda start, end: setattr(
+                req2, "extend_range", Range(start, end)
+            )
+        )
 
         result2 = adder2.add_one_req(
             req2, has_chunked_req=False, truncation_align_size=None
@@ -492,6 +504,11 @@ class TestPrefillAdder(CustomTestCase):
         req3.full_untruncated_fill_ids = list(range(3))
         req3.last_node = MagicMock()
         req3.sampling_params.ignore_eos = False
+        req3.set_extend_range = MagicMock(
+            side_effect=lambda start, end: setattr(
+                req3, "extend_range", Range(start, end)
+            )
+        )
 
         result3 = adder2.add_one_req(
             req3, has_chunked_req=False, truncation_align_size=None