diff --git a/evals/benchmark/stresscli/commands/load_test.py b/evals/benchmark/stresscli/commands/load_test.py index 42c86ee6..133e2b46 100644 --- a/evals/benchmark/stresscli/commands/load_test.py +++ b/evals/benchmark/stresscli/commands/load_test.py @@ -82,7 +82,6 @@ def collect_metrics(collector, services, output_dir, namespace=None): services=services, output_dir=output_dir, restart_pods_flag=False, - ) else: # If namespace is not provided, call without namespace @@ -175,10 +174,12 @@ def run_locust_test(kubeconfig, global_settings, run_settings, output_folder, in services = global_settings.get("service-list") or [] if runspec["deployment-type"] == "k8s": from .metrics import MetricsCollector + collector = MetricsCollector() collect_metrics(collector, services, start_output_folder, namespace) elif runspec["deployment-type"] == "docker": from .metrics_docker import DockerMetricsCollector + collector = DockerMetricsCollector() collect_metrics(collector, services, start_output_folder) @@ -192,6 +193,7 @@ def run_locust_test(kubeconfig, global_settings, run_settings, output_folder, in if service_metric: from .metrics_util import export_metric + services = global_settings.get("service-list") or [] if runspec["deployment-type"] == "k8s": collect_metrics(collector, services, end_output_folder, namespace) diff --git a/evals/benchmark/stresscli/commands/metrics_docker.py b/evals/benchmark/stresscli/commands/metrics_docker.py index 052b964d..48ddfcdc 100644 --- a/evals/benchmark/stresscli/commands/metrics_docker.py +++ b/evals/benchmark/stresscli/commands/metrics_docker.py @@ -1,20 +1,26 @@ +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + import json -import docker -import requests import logging import os import time +import requests + +import docker + # Setup logs log_level = os.getenv("LOG_LEVEL", "ERROR").upper() logging.basicConfig(level=getattr(logging, log_level)) + class DockerMetricsCollector: def __init__(self): self.docker_client = docker.from_env() def get_docker_container(self, container_name): - """Retrieve Docker container information""" + """Retrieve Docker container information.""" try: container = self.docker_client.containers.get(container_name) logging.info(f"Found Docker container {container_name}") @@ -24,24 +30,28 @@ def get_docker_container(self, container_name): return None def get_exposed_port(self, container): - """Get the port exposed to the external environment by the Docker container""" + """Get the port exposed to the external environment by the Docker container.""" try: # Retrieve ports in JSON format - ports_json = container.attrs['NetworkSettings']['Ports'] + ports_json = container.attrs["NetworkSettings"]["Ports"] logging.debug(f"Container ports: {ports_json}") # Parse the ports to find the host port for container_port, host_infos in ports_json.items(): for host_info in host_infos: - host_ip = host_info['HostIp'] - host_port = host_info['HostPort'] + host_ip = host_info["HostIp"] + host_port = host_info["HostPort"] # Use localhost if the port is mapped to 0.0.0.0 or empty - if host_ip in ['0.0.0.0', '']: - logging.debug(f"Found host port {host_port} for container port {container_port} (mapped to all interfaces)") - return ('localhost', host_port) + if host_ip in ["0.0.0.0", ""]: + logging.debug( + f"Found host port {host_port} for container port {container_port} (mapped to all interfaces)" + ) + return ("localhost", host_port) else: - logging.debug(f"Found host port {host_port} for container port {container_port} (mapped to {host_ip})") + logging.debug( + f"Found host port {host_port} for container port {container_port} (mapped to {host_ip})" + ) return (host_ip, host_port) logging.error("No valid host port found.") @@ -51,7 +61,7 @@ def get_exposed_port(self, container): return (None, None) def collect_metrics(self, container_name, metrics_path="/metrics"): - """Collect metrics from the Docker container""" + """Collect metrics from the Docker container.""" container = self.get_docker_container(container_name) if container: try: @@ -71,7 +81,7 @@ def collect_metrics(self, container_name, metrics_path="/metrics"): return None def start_collecting_data(self, services, output_dir="/data"): - """Start collecting metrics from services""" + """Start collecting metrics from services.""" timestamp = int(time.time()) for container_name in services: metrics = self.collect_metrics(container_name) @@ -84,10 +94,18 @@ def start_collecting_data(self, services, output_dir="/data"): logging.error(f"No metrics collected for container {container_name}") return {"status": "success"} + if __name__ == "__main__": docker_collector = DockerMetricsCollector() result = docker_collector.start_collecting_data( - services=["llm-tgi-server", "retriever-redis-server", "embedding-tei-server", "tei-embedding-server", "tgi-service", "tei-reranking-server"], + services=[ + "llm-tgi-server", + "retriever-redis-server", + "embedding-tei-server", + "tei-embedding-server", + "tgi-service", + "tei-reranking-server", + ], output_dir="/path/to/data", ) print(result) diff --git a/evals/benchmark/stresscli/commands/metrics_util.py b/evals/benchmark/stresscli/commands/metrics_util.py index 73b17a95..270f153e 100644 --- a/evals/benchmark/stresscli/commands/metrics_util.py +++ b/evals/benchmark/stresscli/commands/metrics_util.py @@ -117,7 +117,7 @@ def calculate_diff(start_dir, end_dir, output_dir, services=None): for service_name in services: # Create a regex pattern to match files starting with the service_name followed by symbol @ - pattern = rf'^{re.escape(service_name)}@.*\.txt$' + pattern = rf"^{re.escape(service_name)}@.*\.txt$" start_service_files = [f for f in start_files if re.match(pattern, f)] end_service_files = [f for f in end_files if re.match(pattern, f)]