From 84c90a6d355a4e80fb9d8eaa3c22e0bd126719ec Mon Sep 17 00:00:00 2001 From: Matt Beton Date: Tue, 26 Aug 2025 05:05:42 -0700 Subject: [PATCH] feat: mlx memory cache for faster ttft Co-authored-by: Evan Co-authored-by: s17 --- remote_git.sh | 91 +++++++ run.sh | 7 +- src/exo/engines/mlx/utils_mlx.py | 38 +++ src/exo/shared/models/model_cards.py | 13 + src/exo/worker/main.py | 2 +- src/exo/worker/runner/runner.py | 4 + .../test_integration/integration_utils.py | 10 +- .../tests/test_integration/test_inference.py | 3 + src/exo/worker/tests/test_mlx.py | 203 +++++++++++++++ .../test_inference_llama70B.py | 244 ++++++++++++++++++ 10 files changed, 606 insertions(+), 9 deletions(-) create mode 100755 remote_git.sh create mode 100644 src/exo/worker/tests/test_mlx.py diff --git a/remote_git.sh b/remote_git.sh new file mode 100755 index 00000000..c224fe0e --- /dev/null +++ b/remote_git.sh @@ -0,0 +1,91 @@ +#!/usr/bin/env bash +set -euo pipefail + +############################################################################### +# Args & prerequisites +############################################################################### +if [[ $# -lt 2 ]]; then + echo "Usage: $0 [git_args...]" >&2 + echo "Examples:" >&2 + echo " $0 mypassword pull" >&2 + echo " $0 mypassword checkout main" >&2 + echo " $0 mypassword status" >&2 + echo " $0 mypassword fetch --all" >&2 + exit 1 +fi + +PASSWORD=$1 +shift # Remove password from args +GIT_CMD="$*" # Remaining args form the git command +HOSTS_FILE=${HOSTS_FILE:-hosts.json} + +for prog in jq sshpass; do + command -v "$prog" >/dev/null || + { echo "Error: $prog not installed."; exit 1; } +done + +############################################################################### +# Load hosts.json (works on macOS Bash 3.2 and Bash 4+) +############################################################################### +if builtin command -v mapfile >/dev/null 2>&1; then + mapfile -t HOSTS < <(jq -r '.[]' "$HOSTS_FILE") +else + HOSTS=() + while IFS= read -r h; do HOSTS+=("$h"); done < <(jq -r '.[]' "$HOSTS_FILE") +fi +[[ ${#HOSTS[@]} -gt 0 ]] || { echo "No hosts found in $HOSTS_FILE"; exit 1; } + +############################################################################### +# Helper – run a remote command and capture rc/stderr/stdout +############################################################################### +ssh_opts=(-o StrictHostKeyChecking=no + -o NumberOfPasswordPrompts=1 # allow sshpass to answer exactly once + -o LogLevel=ERROR) + +run_remote () { # $1 host $2 command + local host=$1 cmd=$2 rc + if sshpass -p "$PASSWORD" ssh "${ssh_opts[@]}" "$host" "$cmd"; then + rc=0 + else + rc=$? + fi + return $rc +} + +############################################################################### +# Run git command locally +############################################################################### +echo "=== Running 'git $GIT_CMD' locally ===" +if (cd ~/exo && git $GIT_CMD); then + echo "✓ Local git command succeeded" +else + echo "❌ Local git command failed" + exit 1 +fi + +############################################################################### +# Run git command on remote hosts (parallel) +############################################################################### +echo "" +echo "=== Running 'git $GIT_CMD' on ${#HOSTS[@]} remote host(s) ===" +fail=0 +for h in "${HOSTS[@]}"; do + ( + echo "→ Running on $h..." + if run_remote "$h" "cd ~/exo && git $GIT_CMD"; then + echo " ✓ $h: success" + else + echo " ❌ $h: failed" + exit 1 + fi + ) || fail=1 & +done +wait + +echo "" +if (( fail == 0 )); then + echo "🎉 Git command executed successfully on all hosts!" +else + echo "⚠️ Some hosts failed—see above." + exit 1 +fi \ No newline at end of file diff --git a/run.sh b/run.sh index 4daf8186..82e29cf1 100755 --- a/run.sh +++ b/run.sh @@ -28,7 +28,6 @@ done if [ "$CLEAN" = true ]; then echo "Cleaning databases..." rm -f ~/.exo/*db* - rm -f ~/.exo_replica/*db* fi # Configure MLX @@ -36,14 +35,14 @@ fi # First command (worker) - changes based on replica flag if [ "$REPLICA" = true ]; then - osascript -e "tell app \"Terminal\" to do script \"cd '$DIR'; nix develop -c bash -c 'export EXO_HOME=.exo_replica; uv run exo-worker'\"" + osascript -e "tell app \"Terminal\" to do script \"cd '$DIR'; nix develop -c bash -c 'export EXO_HOME=.exo; uv run exo-worker'\"" else osascript -e "tell app \"Terminal\" to do script \"cd '$DIR'; nix develop -c uv run exo-worker\"" fi # Second command (master) - changes based on replica flag if [ "$REPLICA" = true ]; then - osascript -e "tell app \"Terminal\" to do script \"cd '$DIR'; nix develop -c bash -c 'export RUST_LOG=true EXO_RUN_AS_REPLICA=1 EXO_HOME=.exo_replica API_PORT=8001; uv run exo-master'\"" + osascript -e "tell app \"Terminal\" to do script \"cd '$DIR'; nix develop -c bash -c 'export RUST_LOG=true EXO_RUN_AS_REPLICA=1 EXO_HOME=.exo API_PORT=8001; uv run exo-master'\"" else osascript -e "tell app \"Terminal\" to do script \"cd '$DIR'; nix develop -c bash -c 'export RUST_LOG=true; uv run exo-master'\"" -fi +fi \ No newline at end of file diff --git a/src/exo/engines/mlx/utils_mlx.py b/src/exo/engines/mlx/utils_mlx.py index 955cbb88..6ac3dc6e 100644 --- a/src/exo/engines/mlx/utils_mlx.py +++ b/src/exo/engines/mlx/utils_mlx.py @@ -1,5 +1,6 @@ import asyncio import concurrent.futures +import contextlib import os import resource from asyncio import AbstractEventLoop @@ -39,6 +40,43 @@ class HostList(RootModel[list[str]]): return cls(root=[str(host) for host in hosts]) +def mlx_setup( + model_size_mb: int, + cache_frac_of_mrwss: float = 0.65, # main workhorse + wired_frac_of_mrwss: float = 0.00, # start with no wiring +) -> None: + info = mx.metal.device_info() + mrwss = int(info["max_recommended_working_set_size"]) # bytes + memsize = int(info["memory_size"]) # bytes + + runner_print(f"model size mb {model_size_mb}") + runner_print(f"{mrwss=}") + runner_print(f"{memsize=}") + + model_bytes = int(model_size_mb * 1024**2) + kv_bytes = int(0.02 * model_bytes) + + # Cache: keep most of weights+KV “on ice”, but don’t starve the OS. + target_cache = int(1.10 * (model_bytes + kv_bytes)) # +10% slack + target_cache = min(target_cache, int(cache_frac_of_mrwss * mrwss)) + target_cache = min(target_cache, memsize) + runner_print(f"{target_cache=}") + + mx.set_cache_limit(max(target_cache, 0)) + return + + # Optional hard cap (keeps total MLX usage under control) + with contextlib.suppress(Exception): + mx.set_memory_limit(int(0.85 * mrwss)) + + # Wiring: off by default; if you re‑enable, wire at most a small fraction. + if wired_frac_of_mrwss > 0.0: + target_wired = min(int(wired_frac_of_mrwss * mrwss), int(0.5 * model_bytes)) + target_wired = min(target_wired, target_cache) # don’t wire more than cache + with contextlib.suppress(Exception): # older macOS won’t have this + mx.set_wired_limit(max(target_wired, 0)) + + def mlx_distributed_init(rank: int, hosts: list[Host]) -> mx.distributed.Group: # type: ignore """ Initialize the MLX distributed (runs in thread pool) diff --git a/src/exo/shared/models/model_cards.py b/src/exo/shared/models/model_cards.py index 1bf4822e..ff0669ec 100644 --- a/src/exo/shared/models/model_cards.py +++ b/src/exo/shared/models/model_cards.py @@ -42,6 +42,19 @@ MODEL_CARDS: dict[str, ModelCard] = { n_layers=61, ), ), + "deepseek-v3.1": ModelCard( + short_id="deepseek-v3.1", + model_id="mlx-community/DeepSeek-V3.1-8bit", + name="DeepSeek V3.1 (8-bit)", + description="""DeepSeek V3.1 is a large language model trained on the DeepSeek V3.1 dataset.""", + tags=[], + metadata=ModelMetadata( + model_id="mlx-community/DeepSeek-V3.1-8bit", + pretty_name="DeepSeek V3.1 (8-bit)", + storage_size_kilobytes=754706307, + n_layers=61, + ), + ), # deepseek r1 "deepseek-r1-0528:4bit": ModelCard( short_id="deepseek-r1-0528:4bit", diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py index 2db7eedb..a44280a1 100644 --- a/src/exo/worker/main.py +++ b/src/exo/worker/main.py @@ -7,7 +7,7 @@ from exo.shared.apply import apply from exo.shared.constants import EXO_WORKER_LOG from exo.shared.db.sqlite.event_log_manager import EventLogConfig, EventLogManager from exo.shared.keypair import Keypair, get_node_id_keypair -from exo.shared.logging import logger_setup, logger_cleanup +from exo.shared.logging import logger_cleanup, logger_setup from exo.shared.types.common import NodeId from exo.shared.types.events import ( NodePerformanceMeasured, diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py index 440dcdef..307de378 100644 --- a/src/exo/worker/runner/runner.py +++ b/src/exo/worker/runner/runner.py @@ -14,6 +14,7 @@ from exo.engines.mlx.utils_mlx import ( apply_chat_template, initialize_mlx, mlx_force_oom, + mlx_setup, warmup_inference, ) from exo.shared.openai_compat import FinishReason @@ -34,6 +35,7 @@ from exo.worker.runner.communication import ( runner_write_error, runner_write_response, ) +from exo.worker.runner.utils import get_weights_size_kb async def _mlx_generate( @@ -110,6 +112,8 @@ async def main(): model_shard_meta = setup_message.model_shard_meta hosts = setup_message.hosts + mlx_setup(int(get_weights_size_kb(model_shard_meta) // 2**10)) + # For testing - these are fake break conditions if model_shard_meta.immediate_exception: raise Exception("Fake exception - runner failed to spin up.") diff --git a/src/exo/worker/tests/test_integration/integration_utils.py b/src/exo/worker/tests/test_integration/integration_utils.py index 1112dbd2..0654ad77 100644 --- a/src/exo/worker/tests/test_integration/integration_utils.py +++ b/src/exo/worker/tests/test_integration/integration_utils.py @@ -9,11 +9,12 @@ from exo.shared.types.tasks import TaskId, TaskStatus async def read_streaming_response( global_events: AsyncSQLiteEventStorage, filter_task: Optional[TaskId] = None -) -> Tuple[bool, bool, str]: +) -> Tuple[bool, bool, str, int]: # Read off all events - these should be our GenerationChunk events seen_task_started, seen_task_finished = 0, 0 response_string = "" finish_reason: str | None = None + token_count = 0 if not filter_task: idx = await global_events.get_last_idx() @@ -50,8 +51,9 @@ async def read_streaming_response( if event.task_status == TaskStatus.COMPLETE: seen_task_finished += 1 - if isinstance(event, ChunkGenerated): - assert isinstance(event.chunk, TokenChunk) + if isinstance(event, ChunkGenerated) and isinstance( + event.chunk, TokenChunk + ): response_string += event.chunk.text if event.chunk.finish_reason: finish_reason = event.chunk.finish_reason @@ -60,7 +62,7 @@ async def read_streaming_response( print(f"event log: {await global_events.get_events_since(0)}") - return seen_task_started == 1, seen_task_finished == 1, response_string + return seen_task_started == 1, seen_task_finished == 1, response_string, token_count T = TypeVar("T") diff --git a/src/exo/worker/tests/test_integration/test_inference.py b/src/exo/worker/tests/test_integration/test_inference.py index 53e40abe..3d430f41 100644 --- a/src/exo/worker/tests/test_integration/test_inference.py +++ b/src/exo/worker/tests/test_integration/test_inference.py @@ -72,6 +72,7 @@ async def test_runner_inference( seen_task_started, seen_task_finished, response_string, + _, ) = await read_streaming_response(global_events) assert seen_task_started @@ -152,6 +153,7 @@ async def test_2_runner_inference( seen_task_started, seen_task_finished, response_string, + _, ) = await read_streaming_response(global_events) assert seen_task_started @@ -264,6 +266,7 @@ async def test_2_runner_multi_message( seen_task_started, seen_task_finished, response_string, + _, ) = await read_streaming_response(global_events) assert seen_task_started diff --git a/src/exo/worker/tests/test_mlx.py b/src/exo/worker/tests/test_mlx.py new file mode 100644 index 00000000..a9f50b2a --- /dev/null +++ b/src/exo/worker/tests/test_mlx.py @@ -0,0 +1,203 @@ +# type: ignore + +import contextlib +import os +import time +from pathlib import Path + +import mlx.core as mx +import pytest +from mlx_lm.generate import stream_generate +from mlx_lm.sample_utils import make_sampler +from mlx_lm.tokenizer_utils import load_tokenizer +from mlx_lm.utils import load_model + +MODEL_ID = "mlx-community/Llama-3.3-70B-Instruct-4bit" +MODEL_PATH = Path( + os.path.expanduser("~/.exo/models/mlx-community--Llama-3.3-70B-Instruct-4bit/") +) + + +def _get_model_size_gb(path: str) -> float: + """Calculate total size of directory recursively in GB.""" + total_size = 0 + for dirpath, _, filenames in os.walk(path): + for filename in filenames: + filepath = os.path.join(dirpath, filename) + if os.path.isfile(filepath): + total_size += os.path.getsize(filepath) + return total_size / (1024**3) # Convert bytes to GB + + +@pytest.mark.skipif( + not (os.path.exists(MODEL_PATH) and _get_model_size_gb(MODEL_PATH) > 30), + reason=f"This test only runs when model {MODEL_ID} is downloaded", +) +def test_mlx_profiling(): + """ + Test MLX generation directly to profile: + - Time to first token (TTFT) + - Prefill tokens per second (TPS) + - Generation tokens per second (TPS) + For two consecutive prompts using the 70B Llama model. + """ + + # How much memory to keep "wired" (resident) and how much freed memory MLX should keep cached + info = mx.metal.device_info() # returns limits & sizes + # Start conservatively: e.g., 70–90% of recommended working set + target_bytes = int(0.8 * info["max_recommended_working_set_size"]) + + # Keep more freed buffers around for instant reuse + mx.set_cache_limit(target_bytes) + + # On macOS 15+ you can wire resident memory to avoid OS paging/compression + with contextlib.suppress(Exception): + mx.set_wired_limit(target_bytes) + + print(f"\n=== Loading Model {MODEL_ID} ===") + load_start = time.time() + + # Load model and tokenizer + model, _ = load_model(MODEL_PATH, lazy=True, strict=False) + tokenizer = load_tokenizer(MODEL_PATH) + + # Evaluate model parameters to load them into memory + mx.eval(model.parameters()) + + # Create sampler with temperature 0.7 + sampler = make_sampler(temp=0.7) + + load_time = time.time() - load_start + print(f"Model loaded in {load_time:.2f}s") + + # Define test prompts + prompts = [ + "Write me a haiku about a robot.", + "Please write a haiku about a flower.", + "Please write a haiku about headlights.", + ] + + # Prepare messages in chat format + test_messages = [[{"role": "user", "content": prompt}] for prompt in prompts] + + results = [] + + for i, (messages, prompt_text) in enumerate( + zip(test_messages, prompts, strict=False), 1 + ): + print(f"\n=== Prompt {i}: '{prompt_text}' ===") + + # Apply chat template + formatted_prompt = tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + + # Tokenize to count prompt tokens + prompt_tokens = tokenizer.encode(formatted_prompt) + num_prompt_tokens = len(prompt_tokens) + + print(f"Prompt tokens: {num_prompt_tokens}") + + # Start timing + start_time = time.time() + first_token_time = None + tokens_generated = 0 + generated_text = "" + + # Stream generate tokens + for generation in stream_generate( + model=model, + tokenizer=tokenizer, + prompt=formatted_prompt, + max_tokens=100, + sampler=sampler, + ): + if first_token_time is None: + first_token_time = time.time() + ttft = first_token_time - start_time + print(f"Time to first token: {ttft:.3f}s") + + tokens_generated += 1 + generated_text += generation.text + + # Stop if we hit the finish reason + if generation.finish_reason: + break + + total_time = time.time() - start_time + generation_time = total_time - ttft if first_token_time else total_time + + # Calculate metrics + prefill_tps = num_prompt_tokens / ttft if ttft > 0 else 0 + generation_tps = ( + tokens_generated / generation_time if generation_time > 0 else 0 + ) + + # Store results + result = { + "prompt": prompt_text, + "ttft": ttft, + "total_time": total_time, + "generation_time": generation_time, + "prompt_tokens": num_prompt_tokens, + "tokens_generated": tokens_generated, + "prefill_tps": prefill_tps, + "generation_tps": generation_tps, + "generated_text": generated_text, + } + results.append(result) + + # Print results for this prompt + print(f"Total completion time: {total_time:.3f}s") + print(f"Tokens generated: {tokens_generated}") + print(f"Response length: {len(generated_text)} chars") + print( + f"Prefill TPS: {prefill_tps:.1f} tokens/sec ({num_prompt_tokens} prompt tokens / {ttft:.3f}s)" + ) + print( + f"Generation TPS: {generation_tps:.1f} tokens/sec ({tokens_generated} tokens / {generation_time:.3f}s)" + ) + print(f"Generated text preview: {generated_text[:100]}...") + + # Small delay between prompts + if i < len(prompts): + time.sleep(3.0) + + # Compare results + print("\n=== Comparison ===") + if len(results) == 2: + r1, r2 = results[0], results[1] + + print(f"Second prompt TTFT: {r2['ttft'] / r1['ttft']:.2f}x the first") + print( + f"Second prompt prefill TPS: {r2['prefill_tps'] / r1['prefill_tps']:.2f}x the first" + ) + print( + f"Second prompt generation TPS: {r2['generation_tps'] / r1['generation_tps']:.2f}x the first" + ) + + # Performance expectations + print("\n=== Performance Summary ===") + print("First prompt:") + print(f" TTFT: {r1['ttft']:.3f}s") + print(f" Prefill: {r1['prefill_tps']:.1f} tok/s") + print(f" Generation: {r1['generation_tps']:.1f} tok/s") + + print("Second prompt (warmed up):") + print(f" TTFT: {r2['ttft']:.3f}s") + print(f" Prefill: {r2['prefill_tps']:.1f} tok/s") + print(f" Generation: {r2['generation_tps']:.1f} tok/s") + + # Basic assertions + for result in results: + assert result["ttft"] > 0, "TTFT must be positive" + assert result["tokens_generated"] > 0, "Must generate at least one token" + assert len(result["generated_text"]) > 0, "Must generate some text" + assert result["prefill_tps"] > 0, "Prefill TPS must be positive" + assert result["generation_tps"] > 0, "Generation TPS must be positive" + + print("\n✅ All tests passed!") + + +if __name__ == "__main__": + test_mlx_profiling() diff --git a/src/exo/worker/tests/test_multimodel/test_inference_llama70B.py b/src/exo/worker/tests/test_multimodel/test_inference_llama70B.py index b98d9c5f..6e9ace7f 100644 --- a/src/exo/worker/tests/test_multimodel/test_inference_llama70B.py +++ b/src/exo/worker/tests/test_multimodel/test_inference_llama70B.py @@ -1,5 +1,6 @@ import asyncio import os +import time from logging import Logger from typing import Callable @@ -11,8 +12,10 @@ from exo.shared.models.model_meta import get_model_meta from exo.shared.types.api import ChatCompletionMessage, ChatCompletionTaskParams from exo.shared.types.common import Host from exo.shared.types.events import ( + ChunkGenerated, InstanceCreated, InstanceDeleted, + RunnerStatusUpdated, TaskCreated, ) from exo.shared.types.models import ModelId, ModelMetadata @@ -29,6 +32,7 @@ from exo.shared.types.worker.instances import ( InstanceStatus, ShardAssignments, ) +from exo.shared.types.worker.runners import LoadedRunnerStatus from exo.shared.types.worker.shards import PipelineShardMetadata from exo.worker.download.shard_downloader import NoopShardDownloader from exo.worker.main import run @@ -46,6 +50,7 @@ from exo.worker.tests.constants import ( ) from exo.worker.tests.test_integration.integration_utils import ( read_streaming_response, + until_event_with_timeout, ) from exo.worker.worker import Worker @@ -68,6 +73,242 @@ def _get_model_size_gb(path: str) -> float: return total_size / (1024**3) # Convert bytes to GB +@pytest.mark.skipif( + not ( + os.path.exists( + os.path.expanduser( + "~/.exo/models/mlx-community--Llama-3.3-70B-Instruct-4bit/" + ) + ) + and _get_model_size_gb( + os.path.expanduser( + "~/.exo/models/mlx-community--Llama-3.3-70B-Instruct-4bit/" + ) + ) + > 30 + ), + reason="This test only runs when model mlx-community/Llama-3.3-70B-Instruct-4bit is downloaded", +) +async def test_ttft( + logger: Logger, + pipeline_shard_meta: Callable[[int, int], PipelineShardMetadata], + hosts: Callable[[int], list[Host]], +): + logger_test_install(logger) + event_log_manager = EventLogManager(EventLogConfig()) + await event_log_manager.initialize() + shard_downloader = NoopShardDownloader() + + global_events = event_log_manager.global_events + await global_events.delete_all_events() + + worker1 = Worker( + NODE_A, + shard_downloader=shard_downloader, + worker_events=global_events, + global_events=global_events, + ) + asyncio.create_task(run(worker1)) + + ## Instance + model_id = ModelId(MODEL_ID) + + shard_assignments = ShardAssignments( + model_id=model_id, + runner_to_shard={RUNNER_1_ID: pipeline_shard_meta(1, 0)}, + node_to_runner={NODE_A: RUNNER_1_ID}, + ) + + instance = Instance( + instance_id=INSTANCE_1_ID, + instance_type=InstanceStatus.ACTIVE, + shard_assignments=shard_assignments, + hosts=hosts(1), + ) + + # Create instance first + await global_events.append_events( + [InstanceCreated(instance=instance)], origin=MASTER_NODE_ID + ) + + await until_event_with_timeout( + global_events, + event_type=RunnerStatusUpdated, + condition=lambda x: isinstance(x.runner_status, LoadedRunnerStatus), + ) + logger.info("model loaded.") + + # First inference + task1_params = ChatCompletionTaskParams( + model="gpt-4", + messages=[ + ChatCompletionMessage( + role="user", content="Please write a haiku about a flower." + ) + ], + stream=True, + max_tokens=100, + ) + task1 = ChatCompletionTask( + task_id=TASK_1_ID, + command_id=COMMAND_1_ID, + instance_id=INSTANCE_1_ID, + task_type=TaskType.CHAT_COMPLETION, + task_status=TaskStatus.PENDING, + task_params=task1_params, + ) + + print("Starting first inference...") + # Record the current event index before creating the task + idx_before_task1 = await global_events.get_last_idx() + + task_created_time_1 = time.time() + await global_events.append_events( + [TaskCreated(task_id=task1.task_id, task=task1)], origin=MASTER_NODE_ID + ) + + # Wait for first chunk to measure time to first token + first_chunk_seen_1 = False + time_to_first_token_1: None | float = None + while not first_chunk_seen_1: + events = await global_events.get_events_since(idx_before_task1) + for wrapped_event in events: + if isinstance(wrapped_event.event, ChunkGenerated) and hasattr( + wrapped_event.event, "chunk" + ): + first_chunk_time_1 = time.time() + time_to_first_token_1 = first_chunk_time_1 - task_created_time_1 + first_chunk_seen_1 = True + break + if not first_chunk_seen_1: + await asyncio.sleep(0.01) + + _, seen_task_finished_1, response_string_1, _ = await read_streaming_response( + global_events + ) + # # total_time_1 = time.time() - task_created_time_1 + + assert seen_task_finished_1 + + # Wait for first task to complete + await asyncio.sleep(3.0) + + # Second inference + task2_params = ChatCompletionTaskParams( + model="gpt-4", + messages=[ + ChatCompletionMessage( + role="user", content="Write me a haiku about a robot." + ) + ], + stream=True, + max_tokens=150, + ) + task2 = ChatCompletionTask( + task_id=TASK_2_ID, + command_id=COMMAND_2_ID, + instance_id=INSTANCE_1_ID, + task_type=TaskType.CHAT_COMPLETION, + task_status=TaskStatus.PENDING, + task_params=task2_params, + ) + + print("Starting second inference...") + # Record the current event index before creating the second task + idx_before_task2 = await global_events.get_last_idx() + + task_created_time_2 = time.time() + await global_events.append_events( + [TaskCreated(task_id=task2.task_id, task=task2)], origin=MASTER_NODE_ID + ) + + # Wait for first chunk of second task to measure time to first token + first_chunk_seen_2 = False + time_to_first_token_2: float | None = None + while not first_chunk_seen_2: + events = await global_events.get_events_since(idx_before_task2) + for wrapped_event in events: + if isinstance(wrapped_event.event, ChunkGenerated) and hasattr( + wrapped_event.event, "chunk" + ): + first_chunk_time_2 = time.time() + time_to_first_token_2 = first_chunk_time_2 - task_created_time_2 + first_chunk_seen_2 = True + break + if not first_chunk_seen_2: + await asyncio.sleep(0.01) + + _, seen_task_finished_2, response_string_2, _ = await read_streaming_response( + global_events, filter_task=TASK_2_ID + ) + # # total_time_2 = time.time() - task_created_time_2 + + assert seen_task_finished_2 + assert time_to_first_token_1 + assert time_to_first_token_2 + + # Calculate TPS metrics + # Prompt is approximately 45 tokens according to user + # prompt_tokens = 45 + + # # Prefill TPS = prompt tokens / time to first token + # prefill_tps_1 = prompt_tokens / time_to_first_token_1 if time_to_first_token_1 > 0 else 0 + # prefill_tps_2 = prompt_tokens / time_to_first_token_2 if time_to_first_token_2 > 0 else 0 + + # # Generation TPS = generated tokens / generation time + # # Generation time = total time - time to first token + # generation_time_1 = total_time_1 - time_to_first_token_1 + # generation_time_2 = total_time_2 - time_to_first_token_2 + # generation_tps_1 = token_count_1 / generation_time_1 if generation_time_1 > 0 else 0 + # generation_tps_2 = token_count_2 / generation_time_2 if generation_time_2 > 0 else 0 + + # # Display time to first token profiling results + # print("\n=== Time to First Token Profiling ===") + # print(f"First inference ('{task1.task_params.messages[0].content}'):") + # print(f" Time to first token: {time_to_first_token_1:.3f}s") + # print(f" Total completion time: {total_time_1:.3f}s") + # print(f" Tokens generated: {token_count_1}") + # print(f" Response length: {len(response_string_1)} chars") + # print(f" Prefill TPS: {prefill_tps_1:.1f} tokens/sec ({prompt_tokens} prompt tokens / {time_to_first_token_1:.3f}s)") + # print(f" Generation TPS: {generation_tps_1:.1f} tokens/sec ({token_count_1} tokens / {generation_time_1:.3f}s)") + + # print(f"\nSecond inference ('{task2.task_params.messages[0].content}'):") + # print(f" Time to first token: {time_to_first_token_2:.3f}s") + # print(f" Total completion time: {total_time_2:.3f}s") + # print(f" Tokens generated: {token_count_2}") + # print(f" Response length: {len(response_string_2)} chars") + # print(f" Prefill TPS: {prefill_tps_2:.1f} tokens/sec ({prompt_tokens} prompt tokens / {time_to_first_token_2:.3f}s)") + # print(f" Generation TPS: {generation_tps_2:.1f} tokens/sec ({token_count_2} tokens / {generation_time_2:.3f}s)") + + # print("\nComparison:") + # print(f" Second inference time to first token: {time_to_first_token_2/time_to_first_token_1:.2f}x the first") + # print(f" Second inference prefill TPS: {prefill_tps_2/prefill_tps_1:.2f}x the first") + # print(f" Second inference generation TPS: {generation_tps_2/generation_tps_1:.2f}x the first") + + # Basic assertions to ensure responses make sense + assert len(response_string_1) > 0 + assert len(response_string_2) > 0 + assert time_to_first_token_1 and time_to_first_token_1 > 0 + assert time_to_first_token_2 and time_to_first_token_2 > 0 + + # Cleanup + idx = await global_events.get_last_idx() + await asyncio.sleep(1.0) + events = await global_events.get_events_since(idx) + assert len(events) == 0 + + await global_events.append_events( + [ + InstanceDeleted( + instance_id=instance.instance_id, + ), + ], + origin=MASTER_NODE_ID, + ) + + await asyncio.sleep(2.0) + + @pytest.mark.skipif( not ( os.path.exists( @@ -153,6 +394,7 @@ async def test_2_runner_inference( seen_task_started, seen_task_finished, response_string, + _, ) = await read_streaming_response(global_events) assert seen_task_started @@ -292,6 +534,7 @@ async def test_parallel_inference( seen_task_started_1, seen_task_finished_1, response_string_1, + _, ) = await read_streaming_response(global_events) incomplete_task = ( @@ -303,6 +546,7 @@ async def test_parallel_inference( seen_task_started_2, seen_task_finished_2, response_string_2, + _, ) = await read_streaming_response(global_events, filter_task=incomplete_task) assert seen_task_started_1