import argparse import datetime import json import os import traceback from typing import Dict, Tuple, List import GPUtil import docker from docker.models.containers import Container from loguru import logger import pandas as pd class InferenceEngineRunner: def __init__(self, model: str): self.model = model def run(self, parameters: list[tuple], gpus: int = 0): NotImplementedError("This method should be implemented by the subclass") def stop(self): NotImplementedError("This method should be implemented by the subclass") class TGIDockerRunner(InferenceEngineRunner): def __init__(self, model: str, image: str = "ghcr.io/huggingface/text-generation-inference:latest", volumes=None): super().__init__(model) if volumes is None: volumes = [] self.container = None self.image = image self.volumes = volumes def run(self, parameters: list[tuple], gpus: int = 0): params = f"--model-id {self.model} --port 8080" for p in parameters: params += f" --{p[0]} {str(p[1])}" logger.info(f"Running TGI with parameters: {params}") volumes = {} for v in self.volumes: volumes[v[0]] = {"bind": v[1], "mode": "rw"} self.container = run_docker(self.image, params, "Connected", "ERROR", volumes=volumes, gpus=gpus, ports={"8080/tcp": 8080} ) def stop(self): if self.container: self.container.stop() class BenchmarkRunner: def __init__(self, image: str = "ghcr.io/huggingface/text-generation-inference-benchmark:latest", volumes: List[Tuple[str, str]] = None): if volumes is None: volumes = [] self.container = None self.image = image self.volumes = volumes def run(self, parameters: list[tuple], network_mode): params = "text-generation-inference-benchmark" for p in parameters: params += f" --{p[0]} {str(p[1])}" if p[1] is not None else f" --{p[0]}" logger.info(f"Running text-generation-inference-benchmarks with parameters: {params}") volumes = {} for v in self.volumes: volumes[v[0]] = {"bind": v[1], "mode": "rw"} self.container = run_docker(self.image, params, "Benchmark finished", "Fatal:", volumes=volumes, extra_env={"RUST_LOG": "text_generation_inference_benchmark=info", "RUST_BACKTRACE": "full"}, network_mode=network_mode) def stop(self): if self.container: self.container.stop() def run_docker(image: str, args: str, success_sentinel: str, error_sentinel: str, ports: Dict[str, int] = None, volumes=None, network_mode: str = "bridge", gpus: int = 0, extra_env: Dict[str, str] = None) -> Container: if ports is None: ports = {} if volumes is None: volumes = {} if extra_env is None: extra_env = {} client = docker.from_env(timeout=300) # retrieve the GPU devices from CUDA_VISIBLE_DEVICES devices = [f"{i}" for i in range(get_num_gpus())][:gpus] environment = {"HF_TOKEN": os.environ.get("HF_TOKEN")} environment.update(extra_env) container = client.containers.run(image, args, detach=True, device_requests=[ docker.types.DeviceRequest(device_ids=devices, capabilities=[['gpu']]) ] if gpus > 0 else None, volumes=volumes, shm_size="1g", ports=ports, network_mode=network_mode, environment=environment, ) for line in container.logs(stream=True): print(line.decode("utf-8"), end="") if success_sentinel.encode("utf-8") in line: break if error_sentinel.encode("utf-8") in line: container.stop() raise Exception(f"Error starting container: {line}") return container def get_gpu_names() -> str: gpus = GPUtil.getGPUs() if len(gpus) == 0: return '' return f'{len(gpus)}x{gpus[0].name if gpus else "No GPU available"}' def get_gpu_name() -> str: gpus = GPUtil.getGPUs() if len(gpus) == 0: return '' return gpus[0].name def get_num_gpus() -> int: return len(GPUtil.getGPUs()) def build_df(model: str, data_files: dict[str, str]) -> pd.DataFrame: df = pd.DataFrame() now = datetime.datetime.now(datetime.timezone.utc) created_at = now.isoformat() # '2024-10-02T11:53:17.026215+00:00' # Load the results for key, filename in data_files.items(): with open(filename, 'r') as f: data = json.load(f) for result in data['results']: entry = result [config] = pd.json_normalize(result['config']).to_dict(orient='records') entry.update(config) entry['engine'] = data['config']['meta']['engine'] entry['tp'] = data['config']['meta']['tp'] entry['version'] = data['config']['meta']['version'] entry['model'] = model entry['created_at'] = created_at del entry['config'] df = pd.concat([df, pd.DataFrame(entry, index=[0])]) return df def main(sha, results_file): results_dir = 'results' # get absolute path results_dir = os.path.join(os.path.dirname(__file__), results_dir) logger.info('Starting benchmark') models = [ ('meta-llama/Llama-3.1-8B-Instruct', 1), # ('meta-llama/Llama-3.1-70B-Instruct', 4), # ('mistralai/Mixtral-8x7B-Instruct-v0.1', 2), ] success = True for model in models: tgi_runner = TGIDockerRunner(model[0]) # create results directory model_dir = os.path.join(results_dir, f'{model[0].replace("/", "_").replace(".", "_")}') os.makedirs(model_dir, exist_ok=True) runner = BenchmarkRunner( volumes=[(model_dir, '/opt/text-generation-inference-benchmark/results')] ) try: tgi_runner.run([('max-concurrent-requests', 512)], gpus=model[1]) logger.info(f'TGI started for model {model[0]}') parameters = [ ('tokenizer-name', model[0]), ('max-vus', 800), ('url', 'http://localhost:8080'), ('duration', '120s'), ('warmup', '30s'), ('benchmark-kind', 'rate'), ('prompt-options', 'num_tokens=200,max_tokens=220,min_tokens=180,variance=10'), ('decode-options', 'num_tokens=200,max_tokens=220,min_tokens=180,variance=10'), ('extra-meta', f'"engine=TGI,tp={model[1]},version={sha},gpu={get_gpu_name()}"'), ('no-console', None) ] rates = [('rates', f'{r / 10.}') for r in list(range(8, 248, 8))] parameters.extend(rates) runner.run(parameters, f'container:{tgi_runner.container.id}') except Exception as e: logger.error(f'Error running benchmark for model {model[0]}: {e}') # print the stack trace print(traceback.format_exc()) success = False finally: tgi_runner.stop() runner.stop() if not success: logger.error('Some benchmarks failed') exit(1) df = pd.DataFrame() # list recursively directories directories = [f'{results_dir}/{d}' for d in os.listdir(results_dir) if os.path.isdir(f'{results_dir}/{d}')] logger.info(f'Found result directories: {directories}') for directory in directories: data_files = {} for filename in os.listdir(directory): if filename.endswith('.json'): data_files[filename.split('.')[-2]] = f'{directory}/{filename}' logger.info(f'Processing directory {directory}') df = pd.concat([df, build_df(directory.split('/')[-1], data_files)]) df['device'] = get_gpu_name() df['error_rate'] = df['failed_requests'] / (df['failed_requests'] + df['successful_requests']) * 100.0 df.to_parquet(results_file) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--sha", help="SHA of the commit to add to the results", required=True) parser.add_argument("--results-file", help="The file where to store the results, can be a local file or a s3 path") args = parser.parse_args() if args.results_file is None: results_file = f'{args.sha}.parquet' else: results_file = args.results_file main(args.sha, results_file)