small nits

Author: tmonty12

tests/fault_tolerance/deploy/client.py

@@ -363,7 +363,6 @@ def run_aiperf(
     # Note: For continuous load, we only run once and expect SIGINT to stop it
     max_attempts = 1 if continuous_load else (max_retries if max_retries > 0 else 1)
     success = False
-    all_results = []
 
     for attempt in range(max_attempts):
         if continuous_load:
@@ -394,16 +393,6 @@ def run_aiperf(
                 f.write("\n\n=== STDERR ===\n")
                 f.write(result.stderr)
 
-            all_results.append(
-                {
-                    "attempt": attempt + 1,
-                    "returncode": result.returncode,
-                    "stdout": result.stdout,
-                    "stderr": result.stderr,
-                }
-            )
-
-            # Even with continuous load, with SIGINT, aiperf should return 0 and create the profile_export_aiperf.json file
             if result.returncode == 0:
                 # AI-Perf returns 0 even if all requests failed, so we need to check the output
                 json_path = attempt_dir / "profile_export_aiperf.json"
@@ -440,7 +429,6 @@ def run_aiperf(
                 )
         except Exception as e:
             logger.error(f"Error in attempt {attempt + 1}: {str(e)}")
-            all_results.append({"attempt": attempt + 1, "error": str(e)})
 
         # Sleep before next attempt (if not the last attempt and not continuous load)
         if not success and attempt < max_attempts - 1 and not continuous_load:
@@ -468,9 +456,9 @@ def run_aiperf_with_signal_handling(
     """
     Run aiperf with signal handling for graceful shutdown.
 
-    Handles SIGINT forwarding and timeout when running with subprocess.Popen.
-    This ensures that Ctrl-C and SIGINT are properly forwarded to the subprocess
-    so it can clean up gracefully and write results files.
+    Handles SIGINT and SIGTERM forwarding and timeout when running with subprocess.Popen.
+    This ensures that Ctrl-C (SIGINT) and graceful termination signals (SIGTERM)
+    are properly forwarded to the subprocess so it can clean up gracefully and write results files.
     """
     proc = subprocess.Popen(
         cmd_attempt,
@@ -480,15 +468,20 @@ def run_aiperf_with_signal_handling(
         stdin=subprocess.DEVNULL,
     )
 
-    # Set up signal handler to forward SIGINT to subprocess
     def signal_handler(signum, frame):
-        logger.info(f"Received signal {signum}, forwarding to aiperf subprocess")
+        signal_names = {
+            signal.SIGINT: "SIGINT",
+            signal.SIGTERM: "SIGTERM",
+        }
+        signal_name = signal_names.get(signum, f"signal {signum}")
+        logger.info(f"Received {signal_name}, forwarding to aiperf subprocess")
         try:
-            proc.send_signal(signal.SIGINT)
+            proc.send_signal(signum)
         except ProcessLookupError:
             pass  # Process already terminated
 
     signal.signal(signal.SIGINT, signal_handler)
+    signal.signal(signal.SIGTERM, signal_handler)
 
     try:
         stdout, stderr = proc.communicate(timeout=timeout)

tests/fault_tolerance/deploy/scenarios.py

@@ -13,9 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 import logging
 import re
-import time
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from enum import Enum, auto
@@ -188,7 +188,7 @@ async def execute(
 
         await deployment._wait_for_ready(timeout=1800)  # 30 minute timeout
 
-        time.sleep(
+        await asyncio.sleep(
             self.time
         )  # have some requests processed after the rolling upgrade has completed
 

tests/fault_tolerance/deploy/test_deployment.py

@@ -1,12 +1,12 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import asyncio
 import logging
 import multiprocessing
 import os
 import re
 import signal
-import time
 from contextlib import contextmanager
 from multiprocessing.context import SpawnProcess
 
@@ -224,7 +224,7 @@ async def _inject_failures(
     deployment: ManagedDeployment,
 ):  # noqa: F811
     for failure in failures:
-        time.sleep(failure.time)
+        await asyncio.sleep(failure.time)
 
         logger.info(f"Injecting failure for: {failure}")
 

tests/utils/managed_deployment.py

@@ -57,6 +57,15 @@ def image(self) -> Optional[str]:
         except KeyError:
             return None
 
+    @property
+    def envs(self) -> list[dict[str, str]]:
+        """Environment variables for the service"""
+        return self._spec.get("envs", [])
+
+    @envs.setter
+    def envs(self, value: list[dict[str, str]]):
+        self._spec["envs"] = value
+
     @image.setter
     def image(self, value: str):
         if "extraPodSpec" not in self._spec:
@@ -318,13 +327,9 @@ def set_service_env_var(self, service_name: str, name: str, value: str):
         """
         Set an environment variable for a specific service
         """
-        # Check service exists
-        if service_name not in self._deployment_spec["spec"]["services"]:
-            raise ValueError(f"Service '{service_name}' not found in deployment spec")
-
-        service = self._deployment_spec["spec"]["services"][service_name]
-        if "envs" not in service:
-            service["envs"] = []
+        service = self.get_service(service_name)
+        if service.envs is None:
+            service.envs = []
 
         # if env var already exists, update it
         for env in service["envs"]:
@@ -342,11 +347,8 @@ def get_service_env_vars(self, service_name: str) -> list[dict]:
         Returns:
             List of environment variable dicts (e.g., [{"name": "VAR", "value": "val"}])
         """
-        # Check service exists
-        if service_name not in self._deployment_spec["spec"]["services"]:
-            raise ValueError(f"Service '{service_name}' not found in deployment spec")
-
-        return self._deployment_spec["spec"]["services"][service_name].get("envs", [])
+        service = self.get_service(service_name)
+        return service.envs
 
     @property
     def services(self) -> list[ServiceSpec]:
@@ -374,11 +376,7 @@ def add_arg_to_service(self, service_name: str, arg_name: str, arg_value: str):
             arg_name: Argument name (e.g., "--max-model-len", "--max-seq-len")
             arg_value: Argument value (e.g., "1024")
         """
-        # Get the service
-        if service_name not in self._deployment_spec["spec"]["services"]:
-            raise ValueError(f"Service '{service_name}' not found in deployment spec")
-
-        service = self._deployment_spec["spec"]["services"][service_name]
+        service = self.get_service(service_name)
 
         # Ensure args list exists
         if "extraPodSpec" not in service:
@@ -418,15 +416,23 @@ def add_arg_to_service(self, service_name: str, arg_name: str, arg_value: str):
             # Add new argument
             args_list.extend([arg_name, arg_value])
 
-    def set_service_replicas(self, service_name: str, replicas: int):
+    def get_service(self, service_name: str) -> ServiceSpec:
         """
-        Set the number of replicas for a specific service
+        Get a specific service from the deployment spec
         """
-        # Check service exists
         if service_name not in self._deployment_spec["spec"]["services"]:
             raise ValueError(f"Service '{service_name}' not found in deployment spec")
 
-        self._deployment_spec["spec"]["services"][service_name]["replicas"] = replicas
+        return ServiceSpec(
+            service_name, self._deployment_spec["spec"]["services"][service_name]
+        )
+
+    def set_service_replicas(self, service_name: str, replicas: int):
+        """
+        Set the number of replicas for a specific service
+        """
+        service = self.get_service(service_name)
+        service.replicas = replicas
 
     def save(self, out_file: str):
         """Save updated deployment to file"""