← back to Exo
feat: mlx memory cache for faster ttft
84c90a6d355a4e80fb9d8eaa3c22e0bd126719ec · 2025-08-26 05:05:42 -0700 · Matt Beton
Co-authored-by: Evan <evanev7@gmail.com>
Co-authored-by: s17 <s17@s17s-Mac-Studio.local>
Files touched
A remote_git.shM run.shM src/exo/engines/mlx/utils_mlx.pyM src/exo/shared/models/model_cards.pyM src/exo/worker/main.pyM src/exo/worker/runner/runner.pyM src/exo/worker/tests/test_integration/integration_utils.pyM src/exo/worker/tests/test_integration/test_inference.pyA src/exo/worker/tests/test_mlx.pyM src/exo/worker/tests/test_multimodel/test_inference_llama70B.py
Diff
commit 84c90a6d355a4e80fb9d8eaa3c22e0bd126719ec
Author: Matt Beton <matthew.beton@gmail.com>
Date: Tue Aug 26 05:05:42 2025 -0700
feat: mlx memory cache for faster ttft
Co-authored-by: Evan <evanev7@gmail.com>
Co-authored-by: s17 <s17@s17s-Mac-Studio.local>
---
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 +
.../tests/test_integration/integration_utils.py | 10 +-
.../tests/test_integration/test_inference.py | 3 +
src/exo/worker/tests/test_mlx.py | 203 +++++++++++++++++
.../test_multimodel/test_inference_llama70B.py | 244 +++++++++++++++++++++
10 files changed, 606 insertions(+), 9 deletions(-)
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 <PASSWORD> <git_command> [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
← 5efe5562 feat: single entrypoint and logging rework
·
back to Exo
·
full mlx caching implementation 1b8b456c →