dynamo/tests/fault_tolerance/deploy/client.py

671 lines
24 KiB
Python

# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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.
"""AI-Perf client implementation for fault tolerance testing."""
import json
import logging
import os
import signal
import subprocess
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
import requests
from kr8s.objects import Pod
from tests.utils.managed_deployment import ManagedDeployment
LOG_FORMAT = "[TEST] %(asctime)s %(levelname)s %(name)s: %(message)s"
DATE_FORMAT = "%Y-%m-%dT%H:%M:%S"
# Configure logging
logging.basicConfig(
level=logging.INFO,
format=LOG_FORMAT,
datefmt=DATE_FORMAT,
)
def get_frontend_port(
managed_deployment: ManagedDeployment,
client_index: int,
deployment_spec: Any,
pod_ports: Dict[str, Any],
logger: logging.Logger,
) -> Tuple[Optional[str], Optional[int], Optional[Pod]]:
"""
Select a frontend pod using round-robin and setup port forwarding.
Args:
managed_deployment: ManagedDeployment instance
client_index: Client index for round-robin selection
deployment_spec: Deployment specification with port info
pod_ports: Dictionary to track existing port forwards
- Key: pod name (str)
- Value: port forward object from managed_deployment.port_forward()
logger: Logger instance
Returns:
Tuple of (pod_name, local_port, pod_instance) or (None, None, None) if failed
"""
pods = managed_deployment.get_pods([managed_deployment.frontend_service_name])
port = 0
pod_name = None
selected_pod = None
# Filter ready pods and cleanup stale port forwards
pods_ready = []
for pod in pods[managed_deployment.frontend_service_name]:
if pod.ready():
pods_ready.append(pod)
else:
# Cleanup port forwards for non-ready pods
if pod.name in pod_ports:
try:
pod_ports[pod.name].stop()
except Exception as e:
logger.debug(f"Error stopping port forward for {pod.name}: {e}")
del pod_ports[pod.name]
if not pods_ready:
logger.error("No ready frontend pods found")
return None, None, None
# Round-robin selection based on client index
selected_pod = pods_ready[client_index % len(pods_ready)]
pod_name = selected_pod.name
# Setup or reuse port forward
if pod_name not in pod_ports:
# Get port from deployment_spec (default: 8000)
port_value = getattr(deployment_spec, "_port", 8000)
port_forward = managed_deployment.port_forward(selected_pod, port_value)
if port_forward:
pod_ports[pod_name] = port_forward
port = port_forward.local_port
else:
logger.error(f"Failed to create port forward for pod {pod_name}")
return None, None, None
else:
# Reuse existing port forward
port = pod_ports[pod_name].local_port
logger.debug(f"Selected pod {pod_name} with local port {port}")
return pod_name, port, selected_pod
def wait_for_model_availability(
url: str,
endpoint: str,
model: str,
logger: logging.Logger,
max_attempts: int = 15,
attempt_timeouts: Optional[List[float]] = None,
) -> bool:
"""
Wait for model to be available before running AI-Perf.
Args:
url: Base URL for the service
endpoint: API endpoint path
model: Model name to test
logger: Logger instance
max_attempts: Maximum number of attempts to check availability
attempt_timeouts: List of timeout values for each attempt
Returns:
True if model is available, False otherwise
"""
if attempt_timeouts is None:
# Default: Start with 60s timeout, then gradually decrease
attempt_timeouts = [60, 60, 45, 30, 30, 20, 20, 15, 15, 15, 10, 10, 10, 10, 10]
test_url = f"{url}{endpoint}"
for attempt in range(max_attempts):
try:
test_payload = {
"model": model,
"messages": [{"role": "user", "content": "test"}],
"max_tokens": 1,
"stream": False,
}
timeout_val = attempt_timeouts[min(attempt, len(attempt_timeouts) - 1)]
logger.info(
f"Testing model availability at {test_url} (attempt {attempt+1}/{max_attempts}, timeout={timeout_val}s)"
)
response = requests.post(test_url, json=test_payload, timeout=timeout_val)
if response.status_code == 200:
logger.info(f"Model '{model}' is available and responding")
# Give a bit more time for stabilization
logger.info("Model ready, waiting 5s for stabilization...")
time.sleep(5)
return True
elif response.status_code == 404:
logger.warning(
f"Model '{model}' not found (404). Response: {response.text[:200]}"
)
elif response.status_code == 400:
logger.warning(f"Bad request (400). Response: {response.text[:200]}")
else:
logger.warning(
f"Unexpected status code {response.status_code}: {response.text[:200]}"
)
except requests.Timeout as e:
logger.warning(
f"Model availability test timed out (attempt {attempt+1}): {e}"
)
except Exception as e:
logger.warning(f"Model availability test failed (attempt {attempt+1}): {e}")
if attempt < max_attempts - 1:
wait_time = 10 if attempt < 5 else 5
logger.info(f"Waiting {wait_time}s before retry...")
time.sleep(wait_time)
logger.warning("Could not confirm model availability after all attempts")
return False
def validate_aiperf_results(
json_path: Path,
requests_per_client: int,
attempt: int,
logger: logging.Logger,
attempt_dir: Path,
pod_name: str,
port: int,
) -> bool:
"""
Validate AI-Perf results from JSON output.
Args:
json_path: Path to the AI-Perf JSON output file
requests_per_client: Expected number of requests
attempt: Current attempt number (0-based)
logger: Logger instance
attempt_dir: Directory containing attempt results
pod_name: Pod name for logging
port: Port number for logging
Returns:
True if the attempt was successful, False if it should be retried
"""
if not json_path.exists():
# No JSON output, but aiperf returned 0 - might be okay
logger.info(f"Attempt {attempt + 1} completed (return code 0, no JSON output)")
log_summary_metrics(attempt_dir, logger, pod_name, port)
return True
try:
with open(json_path, "r") as f:
aiperf_data = json.load(f)
# Check for errors in the output
error_count = 0
if "records" in aiperf_data and "error_request_count" in aiperf_data["records"]:
error_count = int(
aiperf_data["records"]["error_request_count"].get("avg", 0)
)
# Also check error_summary
if "error_summary" in aiperf_data:
error_summary_count = sum(
err.get("count", 0) for err in aiperf_data["error_summary"]
)
error_count = max(error_count, error_summary_count)
# Consider it a failure if most requests failed (> 90%)
failure_threshold = requests_per_client * 0.9
if error_count >= failure_threshold:
logger.warning(
f"Attempt {attempt + 1} had {error_count}/{requests_per_client} failed requests - retrying"
)
return False # Not successful, continue retrying
else:
successful_count = requests_per_client - error_count
logger.info(
f"Attempt {attempt + 1} succeeded with {successful_count}/{requests_per_client} successful requests"
)
log_summary_metrics(attempt_dir, logger, pod_name, port)
return True # Successful
except Exception as e:
logger.warning(f"Could not parse AI-Perf output to check for failures: {e}")
# Assume success if we can't parse the output but aiperf returned 0
logger.info(
f"Attempt {attempt + 1} completed (return code 0, could not verify success)"
)
log_summary_metrics(attempt_dir, logger, pod_name, port)
return True # Assume success
def run_aiperf(
url: str,
endpoint: str,
model: str,
pod_name: str,
port: int,
requests_per_client: int,
input_token_length: int,
output_token_length: int,
output_dir: Path,
logger: logging.Logger,
max_retries: int = 1,
retry_delay: float = 1,
continuous_load: bool = False,
) -> bool:
"""
Execute AI-Perf with specified parameters.
Args:
url: Base URL (http://localhost:port)
endpoint: API endpoint path (e.g., "v1/chat/completions")
model: Model name
pod_name: Selected pod name for logging
port: Local port number
requests_per_client: Number of requests to send (used if continuous load not enabled)
input_token_length: Input token count
output_token_length: Output token count
output_dir: Directory for AI-Perf artifacts
logger: Logger instance
max_retries: Maximum number of retry attempts (default: 1)
retry_delay: Delay in seconds between retries (default: 1)
continuous_load: If True, use continuous load instead of fixed request count
Returns:
True if successful, False otherwise
"""
# Validate required parameters
if not model or not url or not endpoint:
logger.error(
f"Missing required parameter: model={model!r}, url={url!r}, endpoint={endpoint!r}"
)
return False
# Build AI-Perf command
cmd = [
"aiperf",
"profile",
# Model configuration (required)
"--model",
model,
# Endpoint configuration
"--url",
url,
"--endpoint",
endpoint if endpoint.startswith("/") else f"/{endpoint}",
"--endpoint-type",
"chat", # Required: tells AI-Perf the API type
# Enable streaming for TTFT and ITL metrics
"--streaming",
# Request parameters
"--concurrency",
"1", # Optional: we set to 1 for sequential
# Token configuration
"--synthetic-input-tokens-mean",
str(input_token_length),
"--synthetic-input-tokens-stddev",
"0", # Set to 0 for consistent token counts
"--output-tokens-mean",
str(output_token_length),
"--output-tokens-stddev",
"0", # Set to 0 for consistent token counts
# Skip warmup to avoid initial failures
"--warmup-request-count",
"0",
# Output configuration
"--artifact-dir",
str(output_dir),
"--random-seed",
"100", # For reproducible results
]
if continuous_load:
cmd.extend(["--benchmark-duration", "1800"]) # 30 minutes for continuous load
logger.info("Using continuous load with duration: 30 minutes")
timeout = 1860 # 31 minutes default for duration-based tests (30 minutes + 1 minute buffer)
else:
cmd.extend(["--request-count", str(requests_per_client)])
timeout = max(requests_per_client * 2 + 60, 300) # At least 5 minutes
# Log execution
logger.info(f"Starting AI-Perf for Pod {pod_name} Local Port {port}")
logger.info(f"Using model name: {model}")
# Wait for model to be available
model_ready = wait_for_model_availability(url, endpoint, model, logger)
if not model_ready:
logger.warning("Model not ready, but proceeding with AI-Perf test anyway")
# This might result in all requests failing, but the retry logic will handle it
logger.info(f"Command: {' '.join(cmd)}")
# Retry logic for fault tolerance - retry FULL request count until success
# 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
for attempt in range(max_attempts):
if continuous_load:
logger.info(
"AI-Perf continuous load (will run until interrupted by SIGINT)"
)
else:
logger.info(
f"AI-Perf attempt {attempt + 1}/{max_attempts} with {requests_per_client} requests"
)
# Update output directory for this attempt
attempt_dir = output_dir / f"attempt_{attempt}"
attempt_dir.mkdir(parents=True, exist_ok=True)
# Use the original command but update artifact directory
cmd_attempt = cmd.copy()
artifact_dir_idx = cmd_attempt.index("--artifact-dir") + 1
cmd_attempt[artifact_dir_idx] = str(attempt_dir)
try:
result = run_aiperf_with_signal_handling(cmd_attempt, logger, timeout)
# Save logs for this attempt
with open(attempt_dir / "genai_perf.log", "w") as f:
f.write("=== STDOUT ===\n")
f.write(result.stdout)
f.write("\n\n=== STDERR ===\n")
f.write(result.stderr)
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"
success = validate_aiperf_results(
json_path=json_path,
requests_per_client=requests_per_client,
attempt=attempt,
logger=logger,
attempt_dir=attempt_dir,
pod_name=pod_name,
port=port,
)
if success:
break # Success - exit the retry loop
## TODO: bug with aiperf git+https://github.com/ai-dynamo/aiperf.git@54cd6dc820bff8bfebc875da104e59d745e14f75
## where sending a SIGINT on Mac can sometimes have an error code of -9 (SIGABRT) which results in profile_export_aiperf.json not being created
elif result.returncode == -9 and continuous_load:
logger.warning(
f"""
Attempt {attempt + 1} failed with return code {result.returncode}
This is a known bug with aiperf on Mac where sending a SIGINT can sometimes have an error code of -9 (SIGABRT)
which results in profile_export_aiperf.json not being created
"""
)
logger.debug(
f"Stderr: {result.stderr[:500] if result.stderr else 'No stderr'}"
)
else:
logger.warning(
f"Attempt {attempt + 1} failed with return code {result.returncode}"
)
logger.debug(
f"Stderr: {result.stderr[:500] if result.stderr else 'No stderr'}"
)
except Exception as e:
logger.error(f"Error in attempt {attempt + 1}: {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:
time.sleep(retry_delay)
if success and not continuous_load:
logger.info(
f"AI-Perf successfully completed all {requests_per_client} requests for {pod_name}"
)
elif success and continuous_load:
logger.info(
f"AI-Perf sustained continuous load for {pod_name} and existed succesfully"
)
else:
logger.error(f"AI-Perf failed all {max_attempts} attempts for {pod_name}")
return success
# TODO: use file redirection and wait() instead of pipes and communicate
def run_aiperf_with_signal_handling(
cmd_attempt: List[str],
logger: logging.Logger,
timeout: int,
) -> subprocess.CompletedProcess:
"""
Run aiperf with signal handling for graceful shutdown.
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,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
stdin=subprocess.DEVNULL,
)
def signal_handler(signum, frame):
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(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)
returncode = proc.returncode
except subprocess.TimeoutExpired:
logger.warning(f"AI-Perf subprocess timed out after {timeout}s")
proc.kill()
stdout, stderr = proc.communicate()
returncode = proc.returncode
except KeyboardInterrupt:
logger.info("Received KeyboardInterrupt, sending SIGINT to aiperf subprocess")
proc.send_signal(signal.SIGINT)
try:
stdout, stderr = proc.communicate(timeout=30) # Give it time to clean up
returncode = proc.returncode
except subprocess.TimeoutExpired:
logger.warning("Subprocess didn't terminate gracefully, killing it")
proc.kill()
stdout, stderr = proc.communicate()
returncode = proc.returncode
return subprocess.CompletedProcess(cmd_attempt, returncode, stdout, stderr)
def log_summary_metrics(
output_dir: Path, logger: logging.Logger, pod_name: str, port: int
) -> None:
"""
Log summary metrics from AI-Perf results.
Args:
output_dir: Directory containing AI-Perf artifacts
logger: Logger instance
pod_name: Pod name for logging
port: Port number for logging
"""
# Look for AI-Perf output file
profile_json = output_dir / "profile_export_aiperf.json"
if not profile_json.exists():
# Try alternative names
for name in ["profile_export.json", "profile_results.json"]:
alt_path = output_dir / name
if alt_path.exists():
profile_json = alt_path
break
if profile_json.exists():
try:
with open(profile_json) as f:
metrics = json.load(f)
# Request count
request_count = int(metrics.get("request_count", {}).get("avg", 0))
# Check for errors
error_count = len(metrics.get("error_summary", []))
# Latency metrics (in milliseconds)
request_latency = metrics.get("request_latency", {})
avg_latency = request_latency.get("avg", 0) / 1000.0 # Convert to seconds
p99_latency = request_latency.get("p99", 0) / 1000.0 # Convert to seconds
# Throughput metrics
throughput = metrics.get("request_throughput", {}).get("avg", 0)
# Log summary
logger.info(
f"Summary: Pod {pod_name} Port {port} "
f"Requests: {request_count} "
f"Errors: {error_count} "
f"Throughput: {throughput:.1f} req/s "
f"Avg Latency: {avg_latency:.3f}s "
f"P99 Latency: {p99_latency:.3f}s"
)
# Log success rate
if request_count > 0:
success_rate = ((request_count - error_count) / request_count) * 100
logger.info(f"Success rate: {success_rate:.1f}%")
# Also write summary to CSV file for aggregation
csv_path = output_dir / "profile_export_aiperf.csv"
if csv_path.exists():
logger.info(f"AI-Perf results saved to {csv_path}")
except Exception as e:
logger.warning(f"Failed to parse AI-Perf metrics: {e}")
def client(
deployment_spec,
namespace: str,
model: str,
log_dir: str,
index: int,
requests_per_client: int,
input_token_length: int,
output_token_length: int,
max_retries: int,
retry_delay: float = 1,
continuous_load: bool = False,
):
"""
Generate load using AI-Perf for fault tolerance testing.
This function sets up port forwarding to a frontend pod and uses AI-Perf
to generate synthetic requests for performance testing and fault tolerance
evaluation.
Args:
deployment_spec: Deployment specification object
namespace: Kubernetes namespace
model: Model name
log_dir: Directory for output logs and AI-Perf artifacts
index: Client index used for round-robin pod selection
requests_per_client: Number of requests to generate (used if continuous load not enabled)
input_token_length: Number of input tokens per request
output_token_length: Number of output tokens per request
max_retries: Maximum retry attempts for AI-Perf execution
retry_delay: Delay in seconds between retry attempts
continuous_load: If True, use continuous load instead of fixed request count
"""
logger = logging.getLogger(f"CLIENT: {index}")
logging.getLogger("httpx").setLevel(logging.WARNING)
managed_deployment = ManagedDeployment(log_dir, deployment_spec, namespace)
pod_ports: Dict[str, Any] = {}
try:
os.makedirs(log_dir, exist_ok=True)
client_output_dir = Path(log_dir) / f"client_{index}"
client_output_dir.mkdir(parents=True, exist_ok=True)
# Add a startup delay for early clients to give model time to load
time.sleep(15)
# Select frontend pod and setup port forwarding
pod_name, port, selected_pod = get_frontend_port(
managed_deployment=managed_deployment,
client_index=index,
deployment_spec=deployment_spec,
pod_ports=pod_ports,
logger=logger,
)
if not pod_name or not port:
logger.error("Failed to select pod or setup port forwarding")
return
url = f"http://localhost:{port}"
# Get endpoint from deployment_spec (default: /v1/chat/completions)
endpoint = getattr(deployment_spec, "_endpoint", "/v1/chat/completions")
success = run_aiperf(
url=url,
endpoint=endpoint,
model=model,
pod_name=pod_name,
port=port,
requests_per_client=requests_per_client,
input_token_length=input_token_length,
output_token_length=output_token_length,
output_dir=client_output_dir,
logger=logger,
max_retries=max_retries,
retry_delay=retry_delay,
continuous_load=continuous_load,
)
if not success:
logger.error("AI-Perf execution failed")
except Exception as e:
logger.error(f"Client error: {str(e)}")
finally:
for pf_name, port_forward in pod_ports.items():
try:
port_forward.stop()
logger.debug(f"Stopped port forward for {pf_name}")
except Exception as e:
logger.debug(f"Error stopping port forward for {pf_name}: {e}")
logger.info("Exiting")