tests/integration/defs/perf/utils.py (4)

387-389: Docstring formatting inconsistency.

The one-line docstring should be on a single line as per Python conventions (PEP 257).

-    def run_disagg_server(self,
-                          full_test_name: str,
-                          venv: Optional[PythonVenvRunnerImpl],
-                          gpu_clock_lock: GPUClockLock,
-                          session_data_writer: SessionDataWriter,
-                          output_dir: str,
-                          outputs: Dict[int, str] = {},
-                          original_test_name: str = None,
-                          cmd_idx: int = 0,
-                          **kwargs) -> List[str]:
-        """
-        Run the disaggregated server and write the results to the output csv and/or yaml files.
-        """
+    def run_disagg_server(self,
+                          full_test_name: str,
+                          venv: Optional[PythonVenvRunnerImpl],
+                          gpu_clock_lock: GPUClockLock,
+                          session_data_writer: SessionDataWriter,
+                          output_dir: str,
+                          outputs: Dict[int, str] = {},
+                          original_test_name: str = None,
+                          cmd_idx: int = 0,
+                          **kwargs) -> List[str]:
+        """Run the disaggregated server and write the results to the output csv and/or yaml files."""

---

409-409: Unused variable tmpDir can be removed.

The variable tmpDir is assigned but never used in the function. The temp_wd context manager is likely managing the working directory change automatically.

-        tmpDir = temp_wd(self.get_working_dir())

---

474-479: Consider more graceful process termination.

The processes are terminated unconditionally in the finally block. Consider checking process status first and attempting a graceful shutdown before termination.

             finally:
-                server_proc.terminate()
-                ctx_workers_proc.terminate()
-                gen_workers_proc.terminate()
-                server_proc.wait()
-                ctx_workers_proc.wait()
-                gen_workers_proc.wait()
+                # Gracefully shutdown processes
+                for proc, name in [(server_proc, "server"), 
+                                  (ctx_workers_proc, "ctx_workers"),
+                                  (gen_workers_proc, "gen_workers")]:
+                    if proc.poll() is None:  # Process is still running
+                        print_info(f"Terminating {name} process...")
+                        proc.terminate()
+                        try:
+                            proc.wait(timeout=5)
+                        except subprocess.TimeoutExpired:
+                            print_warning(f"{name} process didn't terminate gracefully, killing...")
+                            proc.kill()
+                            proc.wait()

---

637-638: Consider making the placeholder string more descriptive.

The serialized command is replaced with a generic placeholder. Consider making it more informative about what's being hidden.

-        # serialized_cmd = self.get_commands().get_cmd_str(cmd_idx)
-        serialized_cmd = "placeholder for disaggregated server"
+        # serialized_cmd = self.get_commands().get_cmd_str(cmd_idx)  
+        serialized_cmd = "[REDACTED: disaggregated server multi-process command sequence]"

tests/integration/defs/perf/test_perf.py (3)

597-702: Large block of commented-out code should be removed or documented.

There's a substantial amount of commented-out code in the load_from_str method. If this code is no longer needed, it should be removed. If it's being kept for reference or future use, add a clear comment explaining why.

Would you like me to help clean up this commented code or convert it to proper documentation if it's meant to serve as a reference?

---

1823-2108: Consider breaking down the DisaggregatedServerPerfTest class.

The DisaggregatedServerPerfTest class is quite large (285 lines) with multiple responsibilities. Consider extracting configuration generation and command building into separate helper classes for better maintainability.

Consider using a builder pattern or factory pattern for creating the various configurations (ctx_config, gen_config, server_config) and commands. This would improve testability and make the code more modular.

---

1849-1901: Many commented configuration options in worker config generation.

The gen_worker_config method contains many commented-out configuration options. These should either be removed if not needed, or properly documented if they represent future configuration options.

     def gen_worker_config(self):
+        """Generate worker configurations for context and generation servers.
+        
+        Note: Additional configuration options are available but not currently used:
+        - max_batch_size, max_num_tokens, max_seq_len
+        - tensor_parallel_size, pipeline_parallel_size
+        - enable_attention_dp, kv_cache_config
+        See documentation for full configuration options.
+        """
         ctx_config = {
-            # 'max_batch_size': ctx_batch_size,
-            # 'max_num_tokens': ctx_max_num_tokens,
             # ... (remove other commented lines)
             'disable_overlap_scheduler': True,

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

tests/integration/defs/perf/utils.py (1)

312-356: Add timeout handling for disaggregated server processes.

The disaggregated server orchestration launches multiple processes but doesn't have explicit timeout handling for each process. Consider adding timeout parameters to prevent indefinite hangs.

Apply this diff to add timeout handling:

         finally:
-            server_proc.terminate()
-            ctx_workers_proc.terminate()
-            gen_workers_proc.terminate()
-            server_proc.wait()
-            ctx_workers_proc.wait()
-            gen_workers_proc.wait()
+            # Terminate processes gracefully with timeout
+            for proc, name in [(server_proc, 'server'), 
+                                (ctx_workers_proc, 'ctx_workers'),
+                                (gen_workers_proc, 'gen_workers')]:
+                if proc is not None:
+                    proc.terminate()
+                    try:
+                        proc.wait(timeout=60)
+                    except subprocess.TimeoutExpired:
+                        print_error(f"Process {name} did not terminate gracefully, killing...")
+                        proc.kill()
+                        proc.wait()

tests/integration/defs/perf/test_perf.py (2)

311-312: Review threshold values for disaggregated server metrics.

The thresholds for DISAGG_SERVER_E2EL and DISAGG_SERVER_TTFT are set to (0.1, 50) with a comment "Ignore < 50ms?". Please confirm these are the intended values for your performance regression detection.

---

774-776: Add validation for disaggregated server configurations.

The validate method returns early for disaggregated server configurations, bypassing all validation logic. Consider adding at least minimal validation for disaggregated server-specific parameters.

tests/integration/defs/perf/test_perf.py (3)

1570-1570: Remove or guard debug print statement.

The print_info call outputs the entire command output which could be verbose and unnecessary for production use.

-        print_info(outputs[cmd_idx].split("\n"))
+        # Debug output - uncomment if needed for troubleshooting
+        # print_info(outputs[cmd_idx].split("\n"))

---

1899-1945: Extract magic numbers into constants for configuration generation.

The disaggregated worker configuration contains several hardcoded values that should be configurable or at least defined as constants.

+    # Configuration constants
+    _DEFAULT_MAX_BATCH_SIZE = 32
+    _DEFAULT_MAX_NUM_TOKENS = 4096
+    _DEFAULT_MAX_SEQ_LEN = 4096
+    _DEFAULT_FREE_GPU_MEM_FRACTION = 0.5
+    _DEFAULT_CUDA_GRAPH_BATCH_SIZES = [1, 2, 4, 8, 16, 32]
+
     def _gen_disagg_worker_config(self):
         ctx_config = {
-            'max_batch_size': 32,
-            'max_num_tokens': 4096,
-            'max_seq_len': 4096,
+            'max_batch_size': self._DEFAULT_MAX_BATCH_SIZE,
+            'max_num_tokens': self._DEFAULT_MAX_NUM_TOKENS,
+            'max_seq_len': self._DEFAULT_MAX_SEQ_LEN,
             'tensor_parallel_size': self._config.ctx_tp_size,
             'enable_attention_dp': self._config.ctx_dp_size > 1,
             'print_iter_log': True,
             'disable_overlap_scheduler': True,
             'kv_cache_config': {
                 'enable_block_reuse': False,
-                'free_gpu_memory_fraction': 0.5,
+                'free_gpu_memory_fraction': self._DEFAULT_FREE_GPU_MEM_FRACTION,
                 'dtype': 'fp8',
             },

---

2008-2046: Consider making benchmark parameters configurable.

The benchmark command uses hardcoded values for input/output lengths and number of prompts. Consider making these configurable based on the test configuration.

     def _get_disagg_benchmark_command(self):
         benchmark_script = os.path.join(self._llm_root, "tensorrt_llm", "serve",
                                         "scripts", "benchmark_serving.py")
         model_path = MODEL_PATH_DICT[self._config.model_name]
         model_dir = os.path.join(llm_models_root(), model_path)
         shared_gpt_path = os.path.join(
             llm_models_root(), "datasets",
             "ShareGPT_V3_unfiltered_cleaned_split.json")
+        
+        # Use test configuration values if available, otherwise use defaults
+        input_len = self._config.input_lens[0] if self._config.input_lens else 1024
+        output_len = self._config.output_lens[0] if self._config.output_lens else 1024
+        num_prompts = self._config.num_reqs if hasattr(self._config, 'num_reqs') else 320
+        
         benchmark_cmd = [
             'python3',
             benchmark_script,
             '--model',
             model_dir,
             '--tokenizer',
             model_dir,
             '--dataset-name',
             'random',
             '--dataset-path',
             shared_gpt_path,
             '--random-input-len',
-            '1024',
+            str(input_len),
             '--random-output-len',
-            '1024',
+            str(output_len),
             '--random-prefix-len',
             '0',
             '--num-prompts',
-            '320',
+            str(num_prompts),

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

tests/integration/defs/perf/utils.py (2)

319-351: Harden teardown: guard unbound locals and add wait timeouts with kill fallback.

Unconditional .terminate()/.wait() can raise UnboundLocalError if setup fails and can hang without timeouts.

Apply:

     def run_cmd(self, cmd_idx: int, venv) -> str:
         output = ""
+        server_proc = ctx_workers_proc = gen_workers_proc = None
+        TIMEOUT_SEC = 60
         try:
             with (  # Start ctx workers
@@
                 check_output(self.client_cmd, env=venv._new_env)
                 output += check_output(self.benchmark_cmd, env=venv._new_env)
         finally:
-            server_proc.terminate()
-            ctx_workers_proc.terminate()
-            gen_workers_proc.terminate()
-            server_proc.wait()
-            ctx_workers_proc.wait()
-            gen_workers_proc.wait()
+            for p in (server_proc, ctx_workers_proc, gen_workers_proc):
+                if p is None:
+                    continue
+                try:
+                    p.terminate()
+                    p.wait(timeout=TIMEOUT_SEC)
+                except Exception:
+                    try:
+                        p.kill()
+                        p.wait(timeout=10)
+                    except Exception:
+                        pass
         return output

---

354-356: Return a string and document placeholder rationale.

get_cmd_str returns a list here but other paths return a string; this can break result serialization. Also, add a short comment with the TODO/tracking ID for the placeholder.

Apply:

-    def get_cmd_str(self, cmd_idx) -> List[str]:
-        return ["disaggregated server tests, please check config files"]
+    def get_cmd_str(self, cmd_idx) -> List[str]:
+        # TODO(TRACK-ID): disagg command serialization not yet implemented; emit placeholder for now.
+        return "disaggregated server tests, please check config files"

tests/integration/defs/perf/test_perf.py (2)

766-779: Add minimal validation for disagg configs.

Early return skips all validation. Assert sane worker counts to fail fast.

Apply:

-        if self.runtime == "disagg_server":
-            # TODO: validate disaggregated server config
-            return
+        if self.runtime == "disagg_server":
+            assert self.ctx_server_workers >= 0 and self.gen_server_workers >= 0, "Worker counts must be non-negative"
+            assert (self.ctx_server_workers + self.gen_server_workers) > 0, "At least one disagg worker must be configured"
+            return

---

1901-1947: Duplicate key in ctx_config ('disable_overlap_scheduler').

Key appears twice; Python keeps the last, hiding the first. Remove the duplicate.

Apply:

         ctx_config = {
@@
-            'disable_overlap_scheduler': True,
             'kv_cache_config': {
                 'enable_block_reuse': False,
                 # 'free_gpu_memory_fraction': ctx_free_gpu_memory_fraction,
                 'free_gpu_memory_fraction': 0.5,
                 'dtype': 'fp8',
             },
-            'disable_overlap_scheduler': True,
+            'disable_overlap_scheduler': True,

tests/integration/defs/perf/test_perf.py (2)

1568-1574: Avoid dumping entire output logs on every metric parse.

This is noisy and bloats logs. Consider gating behind a debug flag.

Apply:

-        print_info(outputs[cmd_idx].split("\n"))
+        if os.environ.get("TRTLLM_DEBUG") == "1":
+            print_info(outputs[cmd_idx].split("\n"))

---

278-284: Remove unused DISAGG_SERVER_METRICS_LOG_QUERIES.

get_metric_regex reads from PERF_METRIC_LOG_QUERIES; this map is redundant.

Apply:

- DISAGG_SERVER_METRICS_LOG_QUERIES = {
-     PerfMetricType.DISAGG_SERVER_E2EL:
-     re.compile(r"Median E2EL \(ms\):\s*(\d+\.?\d*)"),
-     PerfMetricType.DISAGG_SERVER_TTFT:
-     re.compile(r"Median TTFT \(ms\):\s*(\d+\.?\d*)"),
- }

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.