feat: mlx memory cache for faster ttft

Co-authored-by: Evan <[email protected]>
Co-authored-by: s17 <[email protected]>

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

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

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)

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",

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,

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.")

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")

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

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()

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