fix: performance tracker (#503)

Author: rebel-eunji

vllm_rbln/v1/worker/metrics.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import atexit
 from collections import defaultdict
 from dataclasses import dataclass, field
 
@@ -193,13 +192,6 @@ def __init__(self, name: str | None = None):
         self.decode_metrics = StepMetrics()
         self.prefill_metrics_by_request_id = PrefillMetricsByRequestID()
         self.padded_decode_metrics = StepMetrics()
-        self._registered_cleanup = False
-
-    def register_cleanup(self):
-        """Register cleanup function to print stats on exit."""
-        if not self._registered_cleanup:
-            atexit.register(self.print_final_stats)
-            self._registered_cleanup = True
 
     def check_dummy_request(self, request_ids: list[str] | None) -> bool:
         if request_ids:
@@ -227,7 +219,7 @@ def record_prefill(
                 f"got {len(request_ids)}: {request_ids}"
             )
             request_id = request_ids[0]
-        self.prefill_metrics.add_measurement(latency, token_count)
+        self.prefill_metrics.add_measurement(latency, token_count, host_time, device_time, ccl_time)
         if request_id:
             self.prefill_metrics_by_request_id.add_measurement(
                 request_id, latency, token_count, host_time, device_time, ccl_time

vllm_rbln/v1/worker/optimum_model_runner.py

@@ -11,11 +11,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import contextlib
 import logging
 import time
 from typing import TYPE_CHECKING, NamedTuple, Union, cast
 
 import numpy as np
+import rebel
 import torch
 import torch.distributed
 import torch.nn as nn
@@ -240,8 +242,8 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
         )
 
         if envs.VLLM_RBLN_METRICS:
-            self.performance_tracker = PerformanceTracker()
-            self.performance_tracker.register_cleanup()
+            self.model_performance_tracker = PerformanceTracker("MODEL")
+            self.sampler_performance_tracker = PerformanceTracker("SAMPLER")
 
         # Ephemeral state transferred
         # between execute_model() and sample_tokens().
@@ -309,22 +311,26 @@ def execute_model(
             )
 
         with record_function_or_nullcontext("rbln_model_runner: forward"):
-            start_time = time.perf_counter()
-            # FIXME model_input must be modified to be padded
-            hidden_states = self.model(model_input)
+            if hasattr(rebel, "capture_reports"):
+                capture_ctx = rebel.capture_reports()
+            else:
+                # use a dummy context manager that does nothing
+                capture_ctx = contextlib.nullcontext()
+            model_start_time = time.perf_counter()
+            with capture_ctx as model_reports:
+                # FIXME model_input must be modified to be padded
+                hidden_states = self.model(model_input)
+            if envs.VLLM_RBLN_METRICS and self.model_performance_tracker is not None:
+                self.collect_metrics(
+                    self.model_performance_tracker,
+                    model_input.is_prompt,
+                    start_time=model_start_time,
+                    end_time=time.perf_counter(),
+                    reports=model_reports,
+                    token_count=0,
+                    # the performance of sampler doesn't depend on token count
+                )
             sample_hidden_states = hidden_states.clone()
-            end_time = time.perf_counter()
-            if envs.VLLM_RBLN_METRICS:
-                # Record performance metrics
-                execution_time = end_time - start_time
-                if model_input.is_prompt:
-                    self.performance_tracker.record_prefill(
-                        execution_time, num_scheduled_tokens
-                    )
-                else:
-                    self.performance_tracker.record_decode(
-                        execution_time, num_scheduled_tokens
-                    )
 
         with record_function_or_nullcontext("rbln_model_runner: postprocess"):
             if self.is_pooling_model:
@@ -450,7 +456,7 @@ def _prepare_inputs(
             finished_requests_ids=list(finished_requests_ids),
             cached_block_tables=cached_block_tables,
             cached_lengths=cached_lengths,
-            is_prompt=is_prefill,
+            is_prompt=is_prefill, # FIXME unify the variable name is_prefill and is_prompt
             dummy_block=scheduler_output.dummy_block,
         )
         return model_input, num_scheduled_tokens
@@ -1308,7 +1314,24 @@ def sample_tokens(
                 padded_logits = logits.reshape(1, -1)
             else:
                 padded_logits = logits
-            sampler_output = self._sample(padded_logits, spec_decode_metadata=None)
+            sampler_start_time = time.perf_counter()
+            if hasattr(rebel, "capture_reports"):
+                capture_ctx = rebel.capture_reports()
+            else:
+                # use a dummy context manager that does nothing
+                capture_ctx = contextlib.nullcontext()
+            with capture_ctx as sampler_reports:
+                sampler_output = self._sample(padded_logits, spec_decode_metadata=None)
+            if envs.VLLM_RBLN_METRICS and self.sampler_performance_tracker is not None:
+                self.collect_metrics(
+                    self.sampler_performance_tracker,
+                    is_prompt,
+                    start_time=sampler_start_time,
+                    end_time=time.perf_counter(),
+                    reports=sampler_reports,
+                    token_count=0,
+                    # the performance of sampler doesn't depend on token count
+                )
         self.input_batch.prev_sampled_token_ids = None
 
         with record_function_or_nullcontext("rbln_model_runner: bookkeep"):
@@ -1481,3 +1504,37 @@ def postprocess_sampler_output(
             logprobs_tensors = LogprobsTensors(**dict)
 
         return num_sampled_tokens, sampled_token_ids, logprobs_tensors
+
+    def collect_metrics(
+        self,
+        performance_tracker: PerformanceTracker,
+        is_prefill: bool,
+        start_time: float,
+        end_time: float,
+        reports: list[dict],
+        token_count: int,
+    ) -> None:
+        execution_time = end_time - start_time
+        host_time = None
+        device_time = None
+        ccl_time = None
+        if reports is not None and len(reports) > 0:
+            host_time = reports[0].get("total_host", None)
+            device_time = reports[0].get("total_device", None)
+            ccl_time = reports[0].get("total_ccl", None)
+        if is_prefill:
+            performance_tracker.record_prefill(
+                execution_time,
+                token_count,
+                host_time=host_time,
+                device_time=device_time,
+                ccl_time=ccl_time,
+            )
+        else:
+            performance_tracker.record_decode(
+                execution_time,
+                token_count,
+                host_time=host_time,
+                device_time=device_time,
+                ccl_time=ccl_time,
+            )

vllm_rbln/v1/worker/optimum_worker.py

@@ -243,9 +243,11 @@ def pin_lora(self, lora_id: int) -> bool:
 
     def shutdown(self) -> None:
         logger.info("v1 optimum_worker shutdown called")
-        if envs.VLLM_RBLN_METRICS and self.model_runner.performance_tracker:
-            # FIXME - performance tracker atexit is not called
-            self.model_runner.performance_tracker.print_final_stats()
+        if envs.VLLM_RBLN_METRICS:
+            if self.model_runner.model_performance_tracker:
+                self.model_runner.model_performance_tracker.print_final_stats()
+            if self.model_runner.sampler_performance_tracker:
+                self.model_runner.sampler_performance_tracker.print_final_stats()
 
 
 def init_worker_distributed_environment(

vllm_rbln/v1/worker/rbln_model_runner.py

@@ -544,11 +544,8 @@ def __init__(
     def _enable_performance_tracker(self):
         if envs.VLLM_RBLN_METRICS:
             self.performance_tracker = PerformanceTracker("MODEL")
-            self.performance_tracker.register_cleanup()
             self.sampler_performance_tracker = PerformanceTracker("SAMPLER")
-            self.sampler_performance_tracker.register_cleanup()
             self.e2e_performance_tracker = PerformanceTracker("E2E")
-            self.e2e_performance_tracker.register_cleanup()
 
     def _get_positions(self, num_tokens: Any):
         if isinstance(num_tokens, int):