Address comments and fix mypy errors.

Signed-off-by: [email protected] <[email protected]>

Author: tzulingk

tests/fault_tolerance/deploy/client.py

@@ -19,6 +19,7 @@
 import random
 import time
 from datetime import datetime
+from typing import Any, Dict
 
 import requests
 
@@ -154,12 +155,13 @@ def client(
     logger = logging.getLogger(f"CLIENT: {index}")
     logging.getLogger("httpx").setLevel(logging.WARNING)
 
-    managed_deployment = ManagedDeployment(None, deployment_spec, namespace)
-    pod_ports = {}
+    managed_deployment = ManagedDeployment(log_dir, deployment_spec, namespace)
+    pod_ports: Dict[str, Any] = {}
 
     min_elapsed_time = 1 / max_request_rate
 
     try:
+        os.makedirs(log_dir, exist_ok=True)
         log_path = os.path.join(log_dir, f"client_{index}.log.txt")
         with open(log_path, "w") as log:
             for i in range(requests_per_client):

tests/fault_tolerance/deploy/parse_results.py

@@ -18,7 +18,7 @@
 import os
 import re
 from datetime import datetime
-from typing import Any
+from typing import Any, Dict, List, Tuple
 
 import pandas as pd
 from tabulate import tabulate
@@ -38,7 +38,7 @@ def parse_test_log(file_path):
                 start_time = datetime.fromisoformat(
                     line.split(" ")[1].replace("T", " ")
                 )
-                start_cmd = []
+                start_cmd: List[str] = []
             elif "Deployment fault-tolerance-test is ready" in line:
                 ready_time = datetime.fromisoformat(
                     line.split(" ")[1].replace("T", " ")
@@ -170,7 +170,7 @@ def parse_process_log(log_dir, process_name):
     }
     if not os.path.isdir(log_dir):
         return {}
-    ready_times = {}
+    ready_times: Dict[str, List[Tuple[datetime, str, float]]] = {}
 
     for entry in os.listdir(log_dir):
         if entry.endswith(".log") and "metrics" not in entry:
@@ -317,7 +317,7 @@ def process_test_directory(test_dir, sla):
     }
 
 
-def main(logs_dir, tablefmt, log_paths=[], sla=None):
+def main(logs_dir, tablefmt, log_paths=None, sla=None):
     results = []
     if log_paths:
         for log_path in log_paths:

tests/fault_tolerance/deploy/scenarios.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 from dataclasses import dataclass
+from typing import Optional
 
 from tests.utils.managed_deployment import DeploymentSpec
 
@@ -26,7 +27,7 @@ class Load:
     output_token_length: int = 100
     max_retries: int = 1
     max_request_rate: float = 1
-    sla: float = None
+    sla: Optional[float] = None
 
 
 @dataclass
@@ -43,7 +44,7 @@ class Scenario:
     deployment: DeploymentSpec
     load: Load
     failures: list[Failure]
-    model: str = None
+    model: Optional[str] = None
 
 
 # Each Deployment Spec contains

tests/fault_tolerance/deploy/test_deployment.py

@@ -5,7 +5,6 @@
 import multiprocessing
 import time
 from contextlib import contextmanager
-from multiprocessing import Process
 
 import pytest
 
@@ -14,8 +13,6 @@
 from tests.fault_tolerance.deploy.scenarios import scenarios
 from tests.utils.managed_deployment import ManagedDeployment
 
-multiprocessing.set_start_method("spawn")
-
 
 @pytest.fixture(params=scenarios.keys())
 def scenario(request):
@@ -37,9 +34,10 @@ def _clients(
     max_request_rate,
 ):
     procs = []
+    ctx = multiprocessing.get_context("spawn")
     for i in range(num_clients):
         procs.append(
-            Process(
+            ctx.Process(
                 target=client,
                 args=(
                     deployment_spec,

tests/utils/managed_deployment.py

@@ -8,7 +8,7 @@
 import shlex
 import time
 from dataclasses import dataclass
-from typing import Optional
+from typing import Any, Optional
 
 import kr8s
 import kubernetes
@@ -179,6 +179,10 @@ def name(self) -> str:
         """Deployment name"""
         return self._deployment_spec["metadata"]["name"]
 
+    @name.setter
+    def name(self, value: str):
+        self._deployment_spec["metadata"]["name"] = value
+
     @property
     def port(self) -> int:
         """Deployment port"""
@@ -193,10 +197,6 @@ def system_port(self) -> int:
     def endpoint(self) -> str:
         return self._endpoint
 
-    @name.setter
-    def name(self, value: str):
-        self._deployment_spec["metadata"]["name"] = value
-
     @property
     def namespace(self) -> str:
         """Deployment namespace"""
@@ -353,12 +353,13 @@ class ManagedDeployment:
     namespace: str
     frontend_service_name: Optional[str] = "Frontend"
 
-    _custom_api = None
-    _core_api = None
-    _in_cluster = False
-    _logger = logging.getLogger()
-    _port_forward = None
-    _deployment_name = None
+    _custom_api: Optional[Any] = None
+    _core_api: Optional[Any] = None
+    _in_cluster: bool = False
+    _logger: logging.Logger = logging.getLogger()
+    _port_forward: Optional[Any] = None
+    _deployment_name: Optional[str] = None
+    _apps_v1: Optional[Any] = None
 
     def __post_init__(self):
         self._deployment_name = self.deployment_spec.name
@@ -379,6 +380,7 @@ async def _init_kubernetes(self):
 
     async def _wait_for_pods(self, label, expected, timeout=300):
         for _ in range(timeout):
+            assert self._core_api is not None, "Kubernetes API not initialized"
             pods = await self._core_api.list_namespaced_pod(
                 self.namespace, label_selector=label
             )
@@ -397,6 +399,7 @@ async def _wait_for_pods(self, label, expected, timeout=300):
 
     async def _scale_statfulset(self, name, label, replicas):
         body = {"spec": {"replicas": replicas}}
+        assert self._apps_v1 is not None, "Kubernetes API not initialized"
         await self._apps_v1.patch_namespaced_stateful_set_scale(
             name, self.namespace, body
         )
@@ -406,6 +409,7 @@ async def _restart_stateful(self, name, label):
         self._logger.info(f"Restarting {name} {label}")
 
         await self._scale_statfulset(name, label, 0)
+        assert self._core_api is not None, "Kubernetes API not initialized"
         nats_pvc = await self._core_api.list_namespaced_persistent_volume_claim(
             self.namespace, label_selector=label
         )
@@ -434,6 +438,7 @@ async def _wait_for_ready(self, timeout: int = 1800, sleep=1, log_interval=60):
         while (time.time() - start_time) < timeout:
             try:
                 attempt += 1
+                assert self._custom_api is not None, "Kubernetes API not initialized"
                 status = await self._custom_api.get_namespaced_custom_object(
                     group="nvidia.com",
                     version="v1alpha1",
@@ -520,6 +525,7 @@ async def _create_deployment(self):
         )
 
         try:
+            assert self._custom_api is not None, "Kubernetes API not initialized"
             await self._custom_api.create_namespaced_custom_object(
                 group="nvidia.com",
                 version="v1alpha1",
@@ -652,7 +658,7 @@ async def _delete_deployment(self):
         Delete the DynamoGraphDeployment CR.
         """
         try:
-            if self._deployment_name:
+            if self._deployment_name and self._custom_api is not None:
                 await self._custom_api.delete_namespaced_custom_object(
                     group="nvidia.com",
                     version="v1alpha1",