Compare commits

..

1 Commits

Author SHA1 Message Date
ciaranbor eb6ae9fd3c Prevent failed instance retries (#1763)
## Motivation

Currently, when a runner fails, the master retries the instance. Most of
the time, this causes a loop over failure. Retries need backoff and a
cap.

## Changes

- src/exo/worker/main.py: Before creating a runner, check an exponential
backoff timer per instance. After EXO_MAX_INSTANCE_RETRIES failures,
send DeleteInstance to permanently remove the instance. Record attempts
on Shutdown; reset on InstanceDeleted.
- src/exo/utils/keyed_backoff.py: Add attempts() method to query retry
count
- src/exo/shared/constants.py: Add EXO_MAX_INSTANCE_RETRIES = 3.

## Why It Works

The worker gates CreateRunner tasks behind a KeyedBackoff, adding
exponential delay (2s base, 30s cap) between retries. After 3 failures
the worker sends DeleteInstance, stopping retries entirely. The backoff
resets when the instance is deleted, so a fresh placement starts clean.

---------

Co-authored-by: Evan <evanev7@gmail.com>
2026-04-01 21:03:34 +01:00
12 changed files with 254 additions and 561 deletions
+22 -68
View File
@@ -7,7 +7,7 @@
# name, patterns, reasoning
#
# Optional per-model overrides (CLI flags take priority over these):
# temperature, top_p, max_tokens, reasoning_effort, enable_thinking
# temperature, top_p, max_tokens, reasoning_effort
#
# Fallback defaults (when no per-model config):
# reasoning: temperature=1.0, max_tokens=131072, reasoning_effort="high"
@@ -18,9 +18,10 @@
# ─── Qwen3.5 (Feb 2026) ─────────────────────────────────────────────
# Source: HuggingFace model cards (Qwen/Qwen3.5-*)
# Model card recommends: temp=0.6, top_p=0.95, top_k=20
# We omit top_k to match vllm eval (which doesn't set it).
# max_tokens=121072 to match vllm eval (131072 context - 10000 safety margin).
# 35B-A3B thinking general: temp=1.0, top_p=0.95, top_k=20
# 397B thinking: temp=0.6, top_p=0.95, top_k=20
# Non-thinking: temp=0.7, top_p=0.8, top_k=20
# max_tokens: 32768 general, 81920 for complex math/code
[[model]]
name = "Qwen3.5 2B"
@@ -28,8 +29,7 @@ patterns = ["Qwen3.5-2B"]
reasoning = true
temperature = 0.6
top_p = 0.95
enable_thinking = true
max_tokens = 121072
max_tokens = 81920
[[model]]
name = "Qwen3.5 9B"
@@ -37,8 +37,7 @@ patterns = ["Qwen3.5-9B"]
reasoning = true
temperature = 0.6
top_p = 0.95
enable_thinking = true
max_tokens = 121072
max_tokens = 81920
[[model]]
name = "Qwen3.5 27B"
@@ -46,17 +45,15 @@ patterns = ["Qwen3.5-27B"]
reasoning = true
temperature = 0.6
top_p = 0.95
enable_thinking = true
max_tokens = 121072
max_tokens = 81920
[[model]]
name = "Qwen3.5 35B A3B"
patterns = ["Qwen3.5-35B-A3B"]
reasoning = true
temperature = 0.6
temperature = 1.0
top_p = 0.95
enable_thinking = true
max_tokens = 121072
max_tokens = 81920
[[model]]
name = "Qwen3.5 122B A10B"
@@ -64,8 +61,7 @@ patterns = ["Qwen3.5-122B-A10B"]
reasoning = true
temperature = 0.6
top_p = 0.95
enable_thinking = true
max_tokens = 121072
max_tokens = 81920
[[model]]
name = "Qwen3.5 397B A17B"
@@ -73,14 +69,12 @@ patterns = ["Qwen3.5-397B-A17B"]
reasoning = true
temperature = 0.6
top_p = 0.95
enable_thinking = true
max_tokens = 121072
max_tokens = 81920
# ─── Qwen3 (Apr 2025) ───────────────────────────────────────────────
# Source: HuggingFace model cards (Qwen/Qwen3-*)
# Model card recommends: temp=0.6, top_p=0.95, top_k=20
# We omit top_k to match vllm eval (which doesn't set it).
# Non-thinking: temp=0.7, top_p=0.8
# Thinking: temp=0.6, top_p=0.95, top_k=20
# Non-thinking: temp=0.7, top_p=0.8, top_k=20
# max_tokens: 32768 general, 38912 for complex math/code
[[model]]
@@ -89,7 +83,6 @@ patterns = ["Qwen3-0.6B"]
reasoning = true
temperature = 0.6
top_p = 0.95
enable_thinking = true
max_tokens = 38912
[[model]]
@@ -98,7 +91,6 @@ patterns = ["Qwen3-30B-A3B"]
reasoning = true
temperature = 0.6
top_p = 0.95
enable_thinking = true
max_tokens = 38912
[[model]]
@@ -107,7 +99,6 @@ patterns = ["Qwen3-235B-A22B"]
reasoning = true
temperature = 0.6
top_p = 0.95
enable_thinking = true
max_tokens = 38912
[[model]]
@@ -116,7 +107,6 @@ patterns = ["Qwen3-Next-80B-A3B-Thinking"]
reasoning = true
temperature = 0.6
top_p = 0.95
enable_thinking = true
max_tokens = 38912
[[model]]
@@ -139,9 +129,9 @@ max_tokens = 16384
name = "Qwen3 Coder Next"
patterns = ["Qwen3-Coder-Next"]
reasoning = false
temperature = 1.0
top_p = 0.95
max_tokens = 121072
temperature = 0.7
top_p = 0.8
max_tokens = 16384
# ─── GPT-OSS (OpenAI) ───────────────────────────────────────────────
# Source: OpenAI GitHub README + HuggingFace discussion #21
@@ -175,38 +165,10 @@ patterns = ["DeepSeek-V3.1"]
reasoning = true
temperature = 0.0
[[model]]
name = "DeepSeek V3.2"
patterns = ["DeepSeek-V3.2"]
reasoning = true
temperature = 1.0
top_p = 0.95
enable_thinking = true
# ─── NVIDIA Nemotron ───────────────────────────────────────────────────
# Source: HuggingFace model cards
# All variants: temp=1.0, top_p=0.95, enable_thinking=true
[[model]]
name = "Nemotron Cascade 2 30B A3B"
patterns = ["Nemotron-Cascade-2-30B-A3B"]
reasoning = true
temperature = 1.0
top_p = 0.95
enable_thinking = true
[[model]]
name = "Nemotron 3 Super 120B A12B"
patterns = ["Nemotron-3-Super-120B-A12B", "NVIDIA-Nemotron-3-Super-120B-A12B"]
reasoning = true
temperature = 1.0
top_p = 0.95
enable_thinking = true
# ─── GLM (ZhipuAI / THUDM) ──────────────────────────────────────────
# Source: HuggingFace model cards + generation_config.json + docs.z.ai
# GLM 4.5+: temp=1.0, top_p=0.95
# max_tokens=121072 to match vllm eval (131072 context - 10000 safety margin)
# Reasoning tasks: 131072 max_tokens; coding/SWE tasks: temp=0.7
[[model]]
name = "GLM-5"
@@ -214,8 +176,7 @@ patterns = ["GLM-5"]
reasoning = true
temperature = 1.0
top_p = 0.95
enable_thinking = true
max_tokens = 121072
max_tokens = 131072
[[model]]
name = "GLM 4.5 Air"
@@ -230,8 +191,7 @@ patterns = ["GLM-4.7-"]
reasoning = true
temperature = 1.0
top_p = 0.95
enable_thinking = true
max_tokens = 121072
max_tokens = 131072
# Note: matches both GLM-4.7 and GLM-4.7-Flash
# ─── Kimi (Moonshot AI) ─────────────────────────────────────────────
@@ -253,8 +213,7 @@ patterns = ["Kimi-K2.5"]
reasoning = true
temperature = 1.0
top_p = 0.95
enable_thinking = true
max_tokens = 121072
max_tokens = 131072
[[model]]
name = "Kimi K2 Instruct"
@@ -264,8 +223,7 @@ temperature = 0.6
# ─── MiniMax ─────────────────────────────────────────────────────────
# Source: HuggingFace model cards + generation_config.json
# All models: temp=1.0, top_p=0.95
# max_tokens=90000 to match vllm eval (100000 context - 10000 safety margin)
# All models: temp=1.0, top_p=0.95, top_k=40
[[model]]
name = "MiniMax M2.5"
@@ -273,8 +231,6 @@ patterns = ["MiniMax-M2.5"]
reasoning = true
temperature = 1.0
top_p = 0.95
enable_thinking = true
max_tokens = 90000
[[model]]
name = "MiniMax M2.1"
@@ -295,8 +251,6 @@ patterns = ["Step-3.5-Flash"]
reasoning = true
temperature = 1.0
top_p = 0.95
enable_thinking = true
max_tokens = 121072
# ─── Llama (Meta) ───────────────────────────────────────────────────
# Source: generation_config.json + meta-llama/llama-models generation.py
+3 -2
View File
@@ -79,7 +79,7 @@ def load_tokenizer_for_bench(model_id: str) -> Any:
model_path = Path(
snapshot_download(
model_id,
allow_patterns=["*.json", "*.py", "*.tiktoken", "*.model", "*.jinja"],
allow_patterns=["*.json", "*.py", "*.tiktoken", "*.model"],
)
)
@@ -430,7 +430,7 @@ def main() -> int:
key=lambda p: (
str(p.get("instance_meta", "")),
str(p.get("sharding", "")),
nodes_used_in_instance(p["instance"]),
-nodes_used_in_instance(p["instance"]),
),
reverse=True,
)
@@ -637,6 +637,7 @@ def main() -> int:
peak = mean(
x["stats"]["peak_memory_usage"]["inBytes"] for x in runs
)
summary = (
f"prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f} "
f"prompt_tokens={ptok} gen_tokens={gtok} "
+89 -465
View File
@@ -62,12 +62,6 @@ from loguru import logger
# ---------------------------------------------------------------------------
MAX_RETRIES = 30
INSTANCE_HEALTH_CHECK_AFTER = 3 # Check instance health after this many consecutive failures
class InstanceFailedError(RuntimeError):
"""Raised when the exo instance is detected as failed/gone."""
pass
DEFAULT_MAX_TOKENS = 16_384
REASONING_MAX_TOKENS = 131_072
TEMPERATURE_NON_REASONING = 0.0
@@ -277,7 +271,7 @@ def run_humaneval_test(
@dataclass
class QuestionResult:
question_id: int | str
question_id: int
prompt: str
response: str
extracted_answer: str | None
@@ -287,11 +281,7 @@ class QuestionResult:
prompt_tokens: int = 0
completion_tokens: int = 0
reasoning_tokens: int = 0
reasoning_content: str = ""
finish_reason: str = ""
elapsed_s: float = 0.0
power_watts: float = 0.0
energy_joules: float = 0.0
@dataclass
@@ -527,10 +517,6 @@ class ApiResult:
prompt_tokens: int
completion_tokens: int
reasoning_tokens: int
reasoning_content: str = ""
finish_reason: str = ""
power_watts: float = 0.0
energy_joules: float = 0.0
async def _call_api(
@@ -544,9 +530,6 @@ async def _call_api(
system_message: str | None = None,
reasoning_effort: str | None = None,
top_p: float | None = None,
top_k: int | None = None,
min_p: float | None = None,
enable_thinking: bool | None = None,
) -> ApiResult:
messages = []
if system_message:
@@ -563,12 +546,6 @@ async def _call_api(
body["reasoning_effort"] = reasoning_effort
if top_p is not None:
body["top_p"] = top_p
if top_k is not None:
body["top_k"] = top_k
if min_p is not None:
body["min_p"] = min_p
if enable_thinking is not None:
body["enable_thinking"] = enable_thinking
resp = await client.post(
f"{base_url}/v1/chat/completions",
@@ -577,40 +554,19 @@ async def _call_api(
)
resp.raise_for_status()
data = resp.json()
choice = data["choices"][0]
message = choice["message"]
content = message.get("content") or ""
reasoning_content = message.get("reasoning_content") or ""
finish_reason = choice.get("finish_reason") or ""
# For thinking models, empty content is expected when finish_reason is "length"
if not content.strip() and finish_reason != "length" and not reasoning_content:
content = data["choices"][0]["message"]["content"]
if not content or not content.strip():
raise ValueError("Empty response from model")
usage = data.get("usage", {})
details = usage.get("completion_tokens_details", {})
power = data.get("power_usage") or {}
return ApiResult(
content=content,
prompt_tokens=usage.get("prompt_tokens", 0),
completion_tokens=usage.get("completion_tokens", 0),
reasoning_tokens=details.get("reasoning_tokens", 0) if details else 0,
reasoning_content=reasoning_content,
finish_reason=finish_reason,
power_watts=power.get("total_avg_sys_power_watts", 0.0),
energy_joules=power.get("total_energy_joules", 0.0),
)
async def _check_instance_health(base_url: str) -> bool:
"""Return True if the exo instance is still reachable."""
try:
async with httpx.AsyncClient() as c:
resp = await c.get(f"{base_url}/models", timeout=5.0)
return resp.status_code == 200
except Exception:
return False
async def call_with_retries(
client: httpx.AsyncClient,
base_url: str,
@@ -622,14 +578,8 @@ async def call_with_retries(
system_message: str | None = None,
reasoning_effort: str | None = None,
top_p: float | None = None,
top_k: int | None = None,
min_p: float | None = None,
enable_thinking: bool | None = None,
instance_failed: asyncio.Event | None = None,
) -> ApiResult | None:
for attempt in range(MAX_RETRIES):
if instance_failed and instance_failed.is_set():
raise InstanceFailedError("Instance already marked as failed")
try:
return await _call_api(
client,
@@ -642,17 +592,8 @@ async def call_with_retries(
system_message,
reasoning_effort,
top_p,
top_k,
min_p,
enable_thinking,
)
except Exception as e:
is_conn_error = isinstance(e, (httpx.ConnectError, httpx.RemoteProtocolError, ConnectionRefusedError, OSError))
if is_conn_error and attempt >= INSTANCE_HEALTH_CHECK_AFTER:
if not await _check_instance_health(base_url):
if instance_failed:
instance_failed.set()
raise InstanceFailedError(f"Instance is down after {attempt + 1} failures: {e}")
if attempt < MAX_RETRIES - 1:
wait = min(2**attempt, 60)
logger.warning(
@@ -677,16 +618,10 @@ async def evaluate_benchmark(
max_tokens: int,
concurrency: int = 1,
limit: int | None = None,
offset: int = 0,
timeout: float | None = None,
reasoning_effort: str | None = None,
top_p: float | None = None,
top_k: int | None = None,
min_p: float | None = None,
enable_thinking: bool | None = None,
difficulty: str | None = None,
checkpoint_path: Path | None = None,
release_version: str | None = None,
) -> list[QuestionResult]:
"""Run a benchmark. Returns per-question results."""
import datasets
@@ -717,21 +652,7 @@ async def evaluate_benchmark(
ds = ds.filter(lambda x: x["difficulty"] == difficulty)
logger.info(f"Filtered to {len(ds)} {difficulty} problems")
if release_version and "release_version" in ds.column_names:
ds = ds.filter(lambda x: x["release_version"] == release_version)
logger.info(
f"Filtered to {len(ds)} problems with release_version={release_version}"
)
# Sort by question_id to match LCB runner ordering (scenario_router.py:60).
# This ensures [offset:offset+limit] slices select the same problems as vllm.
if "question_id" in ds.column_names:
ds = ds.sort("question_id")
total = len(ds)
if offset > 0:
ds = ds.select(range(min(offset, total), total))
total = len(ds)
if limit and limit < total:
ds = ds.select(range(limit))
total = limit
@@ -739,13 +660,6 @@ async def evaluate_benchmark(
logger.info(
f"Evaluating {benchmark_name}: {total} questions, concurrency={concurrency}, "
f"temperature={temperature}, max_tokens={max_tokens}"
+ (f", top_k={top_k}" if top_k is not None else "")
+ (f", min_p={min_p}" if min_p is not None else "")
+ (
f", enable_thinking={enable_thinking}"
if enable_thinking is not None
else ""
)
)
if config.kind == "code":
@@ -753,64 +667,16 @@ async def evaluate_benchmark(
"Code benchmarks execute model-generated code. Use a sandboxed environment."
)
# Load checkpoint for resume
checkpoint_data: dict[str | int, dict[str, Any]] = {}
if checkpoint_path and checkpoint_path.exists():
with open(checkpoint_path) as f:
for line in f:
entry = json.loads(line)
checkpoint_data[entry["question_id"]] = entry
logger.info(f"Loaded {len(checkpoint_data)} checkpointed results")
semaphore = asyncio.Semaphore(concurrency)
instance_failed = asyncio.Event()
results: list[QuestionResult | None] = [None] * total
completed = 0
lock = asyncio.Lock()
def _get_question_id(idx: int, doc: dict) -> str | int:
"""Get a stable question ID for checkpointing."""
if benchmark_name == "livecodebench":
return doc.get("question_id", idx)
elif benchmark_name == "humaneval":
return doc.get("task_id", idx)
return idx
async def process_question(
idx: int, doc: dict, http_client: httpx.AsyncClient
) -> None:
nonlocal completed
system_msg = None
question_id = _get_question_id(idx, doc)
# Bail out early if instance is already dead
if instance_failed.is_set():
return
# Check checkpoint
if question_id in checkpoint_data:
cached = checkpoint_data[question_id]
results[idx] = QuestionResult(
question_id=question_id,
prompt=cached.get("prompt", ""),
response=cached.get("response", ""),
extracted_answer=cached.get("extracted_answer"),
gold_answer=cached.get("gold_answer", ""),
correct=cached.get("correct", False),
error=cached.get("error"),
prompt_tokens=cached.get("prompt_tokens", 0),
completion_tokens=cached.get("completion_tokens", 0),
reasoning_tokens=cached.get("reasoning_tokens", 0),
reasoning_content=cached.get("reasoning_content", ""),
finish_reason=cached.get("finish_reason", ""),
elapsed_s=cached.get("elapsed_s", 0.0),
power_watts=cached.get("power_watts", 0.0),
energy_joules=cached.get("energy_joules", 0.0),
)
async with lock:
completed += 1
logger.info(f" [{completed}/{total}] {question_id} (cached)")
return
if benchmark_name == "gpqa_diamond":
prompt, gold = format_gpqa_question(doc, idx)
@@ -831,48 +697,24 @@ async def evaluate_benchmark(
raise ValueError(f"Unknown benchmark: {benchmark_name}")
async with semaphore:
if instance_failed.is_set():
return
t0 = time.monotonic()
try:
# Race the API call against the instance_failed event
api_task = asyncio.create_task(call_with_retries(
http_client,
base_url,
model,
prompt,
temperature,
max_tokens,
timeout,
system_message=system_msg,
reasoning_effort=reasoning_effort,
top_p=top_p,
top_k=top_k,
min_p=min_p,
enable_thinking=enable_thinking,
instance_failed=instance_failed,
))
failed_waiter = asyncio.create_task(instance_failed.wait())
done, pending = await asyncio.wait(
[api_task, failed_waiter],
return_when=asyncio.FIRST_COMPLETED,
)
for p in pending:
p.cancel()
with contextlib.suppress(asyncio.CancelledError):
await p
if instance_failed.is_set() and api_task not in done:
logger.error(f"Instance failed, aborting {question_id}")
return
api_result = api_task.result()
except InstanceFailedError:
logger.error(f"Instance failed, skipping {question_id}")
return
api_result = await call_with_retries(
http_client,
base_url,
model,
prompt,
temperature,
max_tokens,
timeout,
system_message=system_msg,
reasoning_effort=reasoning_effort,
top_p=top_p,
)
elapsed = time.monotonic() - t0
if api_result is None:
result = QuestionResult(
question_id=question_id,
question_id=idx,
prompt=prompt,
response="",
extracted_answer=None,
@@ -887,17 +729,13 @@ async def evaluate_benchmark(
"prompt_tokens": api_result.prompt_tokens,
"completion_tokens": api_result.completion_tokens,
"reasoning_tokens": api_result.reasoning_tokens,
"reasoning_content": api_result.reasoning_content,
"finish_reason": api_result.finish_reason,
"elapsed_s": elapsed,
"power_watts": api_result.power_watts,
"energy_joules": api_result.energy_joules,
}
if config.kind == "mc":
extracted = extract_mc_answer(response, valid_letters)
result = QuestionResult(
question_id=question_id,
question_id=idx,
prompt=prompt,
response=response,
extracted_answer=extracted,
@@ -911,7 +749,7 @@ async def evaluate_benchmark(
check_aime_answer(extracted, int(gold)) if extracted else False
)
result = QuestionResult(
question_id=question_id,
question_id=idx,
prompt=prompt,
response=response,
extracted_answer=extracted,
@@ -925,7 +763,7 @@ async def evaluate_benchmark(
code = extract_code_block(response, preserve_indent=keep_indent)
if code is None:
result = QuestionResult(
question_id=question_id,
question_id=idx,
prompt=prompt,
response=response,
extracted_answer=None,
@@ -940,7 +778,7 @@ async def evaluate_benchmark(
code,
)
result = QuestionResult(
question_id=question_id,
question_id=idx,
prompt=prompt,
response=response,
extracted_answer="pass" if passed else "fail",
@@ -955,7 +793,7 @@ async def evaluate_benchmark(
exec_meta["sample"],
)
result = QuestionResult(
question_id=question_id,
question_id=idx,
prompt=prompt,
response=response,
extracted_answer="pass" if passed else "fail",
@@ -966,7 +804,7 @@ async def evaluate_benchmark(
)
else:
result = QuestionResult(
question_id=question_id,
question_id=idx,
prompt=prompt,
response=response,
extracted_answer=None,
@@ -977,7 +815,7 @@ async def evaluate_benchmark(
)
else:
result = QuestionResult(
question_id=question_id,
question_id=idx,
prompt=prompt,
response=response,
extracted_answer=None,
@@ -989,80 +827,24 @@ async def evaluate_benchmark(
results[idx] = result
# Write checkpoint (skip failures so they get retried on resume)
if checkpoint_path is not None and not result.error and result.response:
_write_checkpoint(checkpoint_path, result)
async with lock:
completed += 1
n = completed
# Log progress
thinking_info = ""
if result.reasoning_content:
thinking_info = f", {len(result.reasoning_content)} chars thinking"
logger.info(
f" [{n}/{total}] {question_id}: {len(result.response)} chars{thinking_info}, "
f"tokens: {result.prompt_tokens}+{result.completion_tokens} "
f"[{result.finish_reason}]"
+ (f" {result.extracted_answer}" if result.extracted_answer else "")
)
async def _health_monitor() -> None:
"""Periodically check if the instance is still alive."""
# Wait a bit before first check to let things start
await asyncio.sleep(10)
while not instance_failed.is_set():
if not await _check_instance_health(base_url):
# Double-check to avoid false positives
await asyncio.sleep(2)
if not await _check_instance_health(base_url):
logger.error("Health monitor: instance is down!")
instance_failed.set()
return
await asyncio.sleep(5)
if n % max(1, total // 20) == 0 or n == total:
correct_so_far = sum(1 for r in results if r is not None and r.correct)
answered = sum(1 for r in results if r is not None)
logger.info(
f" [{n}/{total}] {correct_so_far}/{answered} correct "
f"({correct_so_far / max(answered, 1):.1%})"
)
async with httpx.AsyncClient() as http_client:
monitor = asyncio.create_task(_health_monitor())
tasks = [process_question(i, doc, http_client) for i, doc in enumerate(ds)]
await asyncio.gather(*tasks)
monitor.cancel()
with contextlib.suppress(asyncio.CancelledError):
await monitor
if instance_failed.is_set():
completed_count = sum(1 for r in results if r is not None)
logger.error(
f"Instance failed! Completed {completed_count}/{total} problems. "
f"Checkpoint saved — restart to resume remaining problems."
)
raise InstanceFailedError("Instance failed during evaluation")
return [r for r in results if r is not None]
def _write_checkpoint(path: Path, result: QuestionResult) -> None:
"""Append a single result to the JSONL checkpoint file."""
entry = {
"question_id": result.question_id,
"response": result.response,
"extracted_answer": result.extracted_answer,
"gold_answer": result.gold_answer,
"correct": result.correct,
"error": result.error,
"prompt_tokens": result.prompt_tokens,
"completion_tokens": result.completion_tokens,
"reasoning_tokens": result.reasoning_tokens,
"reasoning_content": result.reasoning_content,
"finish_reason": result.finish_reason,
"elapsed_s": round(result.elapsed_s, 2),
"power_watts": round(result.power_watts, 2),
"energy_joules": round(result.energy_joules, 2),
}
with open(path, "a") as f:
f.write(json.dumps(entry) + "\n")
# ---------------------------------------------------------------------------
# Results display
# ---------------------------------------------------------------------------
@@ -1085,8 +867,6 @@ def print_results(
total_elapsed = sum(r.elapsed_s for r in results)
wall_clock = max(r.elapsed_s for r in results) if results else 0.0
avg_gen_tps = total_completion_tokens / total_elapsed if total_elapsed > 0 else 0.0
total_energy = sum(r.energy_joules for r in results)
avg_power = sum(r.power_watts for r in results) / max(total, 1)
label = f"[c={concurrency}] " if concurrency is not None else ""
print(f"\n{label}{benchmark_name}: {correct}/{total} ({accuracy:.1%})")
@@ -1098,10 +878,6 @@ def print_results(
f" | total time: {total_elapsed:.1f}s wall clock: {wall_clock:.1f}s"
)
print(tok_line)
if total_energy > 0:
print(
f" power: avg {avg_power:.1f}W | total energy: {total_energy:.1f}J ({total_energy / 3600:.2f}Wh)"
)
if errors:
print(f" API errors: {errors}")
if no_extract:
@@ -1120,8 +896,6 @@ def print_results(
"total_elapsed_s": total_elapsed,
"wall_clock_s": wall_clock,
"avg_gen_tps": avg_gen_tps,
"avg_power_watts": avg_power,
"total_energy_joules": total_energy,
}
@@ -1279,11 +1053,7 @@ def save_results(
"prompt_tokens": r.prompt_tokens,
"completion_tokens": r.completion_tokens,
"reasoning_tokens": r.reasoning_tokens,
"reasoning_content": r.reasoning_content,
"finish_reason": r.finish_reason,
"elapsed_s": round(r.elapsed_s, 2),
"power_watts": round(r.power_watts, 2),
"energy_joules": round(r.energy_joules, 2),
}
for r in results
],
@@ -1299,32 +1069,6 @@ def save_results(
# ---------------------------------------------------------------------------
def _find_existing_instance(client: ExoClient, model_id: str) -> str | None:
"""Find an existing instance for the given model."""
try:
state = client.request_json("GET", "/state")
except Exception:
return None
for inst_id, inst in state.get("instances", {}).items():
# Instance structure is nested: {"MlxJacclInstance": {"shardAssignments": {"modelId": ...}}}
for _inst_type, inner in inst.items():
if not isinstance(inner, dict):
continue
sa = inner.get("shardAssignments", {})
if sa.get("modelId") == model_id:
return inst_id
return None
def _checkpoint_path(
results_dir: str, benchmark: str, model: str, concurrency: int
) -> Path:
"""Return the JSONL checkpoint path for a benchmark run."""
out_dir = Path(results_dir) / model.replace("/", "_") / benchmark
out_dir.mkdir(parents=True, exist_ok=True)
return out_dir / f"c{concurrency}.checkpoint.jsonl"
def parse_int_list(values: list[str]) -> list[int]:
items: list[int] = []
for v in values:
@@ -1352,12 +1096,6 @@ def main() -> int:
default=None,
help="Max questions per benchmark (for fast iteration).",
)
ap.add_argument(
"--offset",
type=int,
default=0,
help="Skip first N questions (0-based).",
)
reasoning_group = ap.add_mutually_exclusive_group()
reasoning_group.add_argument(
@@ -1377,8 +1115,6 @@ def main() -> int:
"--temperature", type=float, default=None, help="Override temperature."
)
ap.add_argument("--top-p", type=float, default=None, help="Override top_p.")
ap.add_argument("--top-k", type=int, default=None, help="Override top_k.")
ap.add_argument("--min-p", type=float, default=None, help="Override min_p.")
ap.add_argument(
"--max-tokens", type=int, default=None, help="Override max output tokens."
)
@@ -1412,42 +1148,16 @@ def main() -> int:
choices=["easy", "medium", "hard"],
help="Filter by difficulty (livecodebench only). E.g. --difficulty hard",
)
ap.add_argument(
"--release-version",
default=None,
help="LCB dataset release version (livecodebench only). E.g. release_v5",
)
ap.add_argument(
"--results-dir",
default="eval_results",
help="Directory for result JSON files (default: eval_results).",
)
ap.add_argument(
"--enable-thinking",
type=lambda v: v.lower() in ("true", "1", "yes"),
default=None,
help="Enable thinking mode for models that support it.",
)
ap.add_argument(
"--skip-instance-setup",
action="store_true",
help="Skip exo instance management (assumes model is already running).",
)
ap.add_argument(
"--force",
action="store_true",
help="Discard any existing checkpoint and run from scratch.",
)
ap.add_argument(
"--keep-instance",
action="store_true",
help="Skip deleting the instance after eval (for chaining runs).",
)
ap.add_argument(
"--reuse-instance",
action="store_true",
help="Reuse an existing ready instance instead of creating a new one.",
)
args, _ = ap.parse_known_args()
@@ -1469,84 +1179,61 @@ def main() -> int:
instance_id: str | None = None
if not args.skip_instance_setup:
_short_id, full_model_id = resolve_model_short_id(
short_id, full_model_id = resolve_model_short_id(
client,
args.model,
force_download=args.force_download,
)
selected = settle_and_fetch_placements(
client,
full_model_id,
args,
settle_timeout=args.settle_timeout,
)
if not selected:
logger.error("No valid placements matched your filters.")
return 1
# Try to reuse an existing instance if --reuse-instance is set
if args.reuse_instance:
existing = _find_existing_instance(client, full_model_id)
if existing:
instance_id = existing
logger.info(f"Reusing existing instance {instance_id}")
selected.sort(
key=lambda p: (
str(p.get("instance_meta", "")),
str(p.get("sharding", "")),
-nodes_used_in_instance(p["instance"]),
),
reverse=True,
)
preview = selected[0]
instance = preview["instance"]
instance_id = instance_id_from_instance(instance)
if instance_id is None:
selected = settle_and_fetch_placements(
client,
full_model_id,
args,
settle_timeout=args.settle_timeout,
)
if not selected:
logger.error("No valid placements matched your filters.")
return 1
logger.info(
f"PLACEMENT: {preview['sharding']} / {preview['instance_meta']} / "
f"nodes={nodes_used_in_instance(instance)}"
)
selected.sort(
key=lambda p: (
str(p.get("instance_meta", "")),
str(p.get("sharding", "")),
nodes_used_in_instance(p["instance"]),
),
reverse=True,
)
preview = selected[0]
instance = preview["instance"]
instance_id = instance_id_from_instance(instance)
settle_deadline = (
time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
)
download_duration = run_planning_phase(
client,
full_model_id,
preview,
args.danger_delete_downloads,
args.timeout,
settle_deadline,
)
if download_duration is not None:
logger.info(f"Download: {download_duration:.1f}s")
logger.info(
f"PLACEMENT: {preview['sharding']} / {preview['instance_meta']} / "
f"nodes={nodes_used_in_instance(instance)}"
)
settle_deadline = (
time.monotonic() + args.settle_timeout
if args.settle_timeout > 0
else None
)
download_duration = run_planning_phase(
client,
full_model_id,
preview,
args.danger_delete_downloads,
args.timeout,
settle_deadline,
)
if download_duration is not None:
logger.info(f"Download: {download_duration:.1f}s")
# Delete any existing instances to free resources before placing
try:
state = client.request_json("GET", "/state")
for old_id in list(state.get("instances", {}).keys()):
logger.info(f"Deleting stale instance {old_id}")
with contextlib.suppress(ExoHttpError):
client.request_json("DELETE", f"/instance/{old_id}")
if state.get("instances"):
time.sleep(2)
except Exception:
pass
client.request_json("POST", "/instance", body={"instance": instance})
try:
wait_for_instance_ready(client, instance_id)
except (RuntimeError, TimeoutError) as e:
logger.error(f"Failed to initialize: {e}")
with contextlib.suppress(ExoHttpError):
client.request_json("DELETE", f"/instance/{instance_id}")
return 1
time.sleep(1)
client.request_json("POST", "/instance", body={"instance": instance})
try:
wait_for_instance_ready(client, instance_id)
except (RuntimeError, TimeoutError) as e:
logger.error(f"Failed to initialize: {e}")
with contextlib.suppress(ExoHttpError):
client.request_json("DELETE", f"/instance/{instance_id}")
return 1
time.sleep(1)
cluster_snapshot = capture_cluster_snapshot(client)
else:
full_model_id = args.model
@@ -1604,57 +1291,16 @@ def main() -> int:
reasoning_effort = str(cfg["reasoning_effort"])
else:
reasoning_effort = "high" if is_reasoning else None
if args.top_k is not None:
top_k: int | None = args.top_k
elif "top_k" in cfg:
top_k = int(cfg["top_k"])
else:
top_k = None
if args.min_p is not None:
min_p: float | None = args.min_p
elif "min_p" in cfg:
min_p = float(cfg["min_p"])
else:
min_p = None
if args.enable_thinking is not None:
enable_thinking: bool | None = args.enable_thinking
elif "enable_thinking" in cfg:
enable_thinking = bool(cfg["enable_thinking"])
else:
enable_thinking = None
base_url = f"http://{args.host}:{args.port}"
logger.info(f"Model: {full_model_id}")
logger.info(
f"Settings: temperature={temperature}, max_tokens={max_tokens}, "
+ (f"top_p={top_p}, " if top_p is not None else "")
+ (f"top_k={top_k}, " if top_k is not None else "")
+ (f"min_p={min_p}, " if min_p is not None else "")
+ f"reasoning={'yes' if is_reasoning else 'no'}"
+ (f", reasoning_effort={reasoning_effort}" if reasoning_effort else "")
+ (
f", enable_thinking={enable_thinking}"
if enable_thinking is not None
else ""
)
)
# Common kwargs for evaluate_benchmark
eval_kwargs: dict[str, Any] = {
"reasoning_effort": reasoning_effort,
"top_p": top_p,
"top_k": top_k,
"min_p": min_p,
"enable_thinking": enable_thinking,
"difficulty": args.difficulty,
"offset": args.offset,
"release_version": args.release_version,
}
try:
if args.compare_concurrency:
concurrency_levels = parse_int_list(args.compare_concurrency)
@@ -1663,11 +1309,6 @@ def main() -> int:
for c in concurrency_levels:
logger.info(f"\n{'=' * 50}")
logger.info(f"Running {task_name} at concurrency={c}")
checkpoint_path = _checkpoint_path(
args.results_dir, task_name, full_model_id, c
)
if args.force and checkpoint_path.exists():
checkpoint_path.unlink()
results = asyncio.run(
evaluate_benchmark(
task_name,
@@ -1678,8 +1319,9 @@ def main() -> int:
concurrency=c,
limit=args.limit,
timeout=args.request_timeout,
checkpoint_path=checkpoint_path,
**eval_kwargs,
reasoning_effort=reasoning_effort,
top_p=top_p,
difficulty=args.difficulty,
)
)
if results:
@@ -1694,18 +1336,10 @@ def main() -> int:
cluster=cluster_snapshot,
)
results_by_c[c] = results
# Clean up checkpoint on success
if checkpoint_path.exists():
checkpoint_path.unlink()
if len(results_by_c) >= 2:
print_comparison(task_name, results_by_c)
else:
for task_name in task_names:
checkpoint_path = _checkpoint_path(
args.results_dir, task_name, full_model_id, args.num_concurrent
)
if args.force and checkpoint_path.exists():
checkpoint_path.unlink()
results = asyncio.run(
evaluate_benchmark(
task_name,
@@ -1716,8 +1350,9 @@ def main() -> int:
concurrency=args.num_concurrent,
limit=args.limit,
timeout=args.request_timeout,
checkpoint_path=checkpoint_path,
**eval_kwargs,
reasoning_effort=reasoning_effort,
top_p=top_p,
difficulty=args.difficulty,
)
)
if results:
@@ -1731,25 +1366,14 @@ def main() -> int:
scores,
cluster=cluster_snapshot,
)
# Clean up checkpoint on success
if checkpoint_path.exists():
checkpoint_path.unlink()
finally:
if instance_id is not None:
if args.keep_instance:
logger.info(f"Keeping instance {instance_id} (--keep-instance)")
else:
try:
client.request_json("DELETE", f"/instance/{instance_id}")
except ExoHttpError as e:
if e.status != 404:
raise
try:
wait_for_instance_gone(client, instance_id)
except TimeoutError:
logger.warning(
f"Timed out waiting for instance {instance_id} to be deleted"
)
try:
client.request_json("DELETE", f"/instance/{instance_id}")
except ExoHttpError as e:
if e.status != 404:
raise
wait_for_instance_gone(client, instance_id)
return 0
+8 -5
View File
@@ -462,8 +462,9 @@ def run_planning_phase(
)
logger.info(f"Started download on {node_id}")
# Wait for downloads (no timeout — poll until complete or failed)
while True:
# Wait for downloads
start = time.time()
while time.time() - start < timeout:
all_done = True
for node_id in node_ids:
node_downloads = client.get_node_downloads(node_id) or []
@@ -513,7 +514,9 @@ def run_planning_phase(
if download_t0 is not None:
return time.perf_counter() - download_t0
return None
time.sleep(10)
time.sleep(1)
raise TimeoutError("Downloads did not complete in time")
def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
@@ -561,8 +564,8 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
ap.add_argument(
"--settle-timeout",
type=float,
default=7200,
help="Max seconds to wait for the cluster to produce valid placements (default: 7200).",
default=0,
help="Max seconds to wait for the cluster to produce valid placements (0 = try once).",
)
ap.add_argument(
"--danger-delete-downloads",
+2
View File
@@ -97,3 +97,5 @@ EXO_OFFLINE = os.getenv("EXO_OFFLINE", "false").lower() == "true"
EXO_TRACING_ENABLED = os.getenv("EXO_TRACING_ENABLED", "false").lower() == "true"
EXO_MAX_CONCURRENT_REQUESTS = int(os.getenv("EXO_MAX_CONCURRENT_REQUESTS", "8"))
EXO_MAX_INSTANCE_RETRIES = 3
+7 -4
View File
@@ -1,10 +1,9 @@
import time
from typing import Generic, TypeVar
K = TypeVar("K")
from typing import final
class KeyedBackoff(Generic[K]):
@final
class KeyedBackoff[K]:
"""Tracks exponential backoff state per key."""
def __init__(self, base: float = 0.5, cap: float = 10.0):
@@ -26,6 +25,10 @@ class KeyedBackoff(Generic[K]):
self._last_time[key] = time.monotonic()
self._attempts[key] = self._attempts.get(key, 0) + 1
def attempts(self, key: K) -> int:
"""Return the number of recorded attempts for a key."""
return self._attempts.get(key, 0)
def reset(self, key: K) -> None:
"""Reset backoff state for a key (e.g., on success)."""
self._attempts.pop(key, None)
+25 -5
View File
@@ -9,9 +9,11 @@ from loguru import logger
from exo.api.types import ImageEditsTaskParams
from exo.download.download_utils import is_read_only_model_dir, resolve_existing_model
from exo.shared.apply import apply
from exo.shared.constants import EXO_MAX_INSTANCE_RETRIES
from exo.shared.models.model_cards import ModelId, add_to_card_cache, delete_custom_card
from exo.shared.types.chunks import InputImageChunk
from exo.shared.types.commands import (
DeleteInstance,
ForwarderCommand,
ForwarderDownloadCommand,
StartDownload,
@@ -23,6 +25,7 @@ from exo.shared.types.events import (
Event,
IndexedEvent,
InputChunkReceived,
InstanceDeleted,
NodeDownloadProgress,
NodeGatheredInfo,
TaskCreated,
@@ -44,6 +47,7 @@ from exo.shared.types.tasks import (
)
from exo.shared.types.topology import Connection, SocketConnection
from exo.shared.types.worker.downloads import DownloadCompleted
from exo.shared.types.worker.instances import InstanceId
from exo.shared.types.worker.runners import RunnerId
from exo.utils.channels import Receiver, Sender, channel
from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer
@@ -84,6 +88,9 @@ class Worker:
self.image_cache: dict[str, str] = {}
self._download_backoff: KeyedBackoff[ModelId] = KeyedBackoff(base=0.5, cap=10.0)
self._instance_backoff: KeyedBackoff[InstanceId] = KeyedBackoff(
base=0.5, cap=10.0
)
self._stopped: anyio.Event = anyio.Event()
async def run(self):
@@ -127,6 +134,9 @@ class Worker:
self.state = apply(self.state, event=event)
event = event.event
if isinstance(event, InstanceDeleted):
self._instance_backoff.reset(event.instance_id)
# Buffer input image chunks for image editing
if isinstance(event, InputChunkReceived):
cmd_id = event.command_id
@@ -156,15 +166,24 @@ class Worker:
self.state.runners,
self.state.tasks,
self.input_chunk_buffer,
self._instance_backoff,
self._download_backoff,
)
if task is None:
continue
# Gate DownloadModel on backoff BEFORE emitting TaskCreated
# to prevent flooding the event log with useless events
if isinstance(task, DownloadModel):
model_id = task.shard_metadata.model_card.model_id
if not self._download_backoff.should_proceed(model_id):
if isinstance(task, CreateRunner):
iid = task.instance_id
if self._instance_backoff.attempts(iid) >= EXO_MAX_INSTANCE_RETRIES:
logger.warning(
f"Instance {iid} exceeded {EXO_MAX_INSTANCE_RETRIES} retries, requesting deletion"
)
await self.command_sender.send(
ForwarderCommand(
origin=self._system_id,
command=DeleteInstance(instance_id=iid),
)
)
continue
logger.info(f"Worker plan: {task.__class__.__name__}")
@@ -175,6 +194,7 @@ class Worker:
match task:
case CreateRunner():
self._create_supervisor(task)
self._instance_backoff.record_attempt(task.instance_id)
await self.event_sender.send(
TaskStatusUpdated(
task_id=task.task_id, task_status=TaskStatus.Complete
+25 -12
View File
@@ -3,7 +3,7 @@
from collections.abc import Mapping, Sequence
from exo.shared.types.chunks import InputImageChunk
from exo.shared.types.common import CommandId, NodeId
from exo.shared.types.common import CommandId, ModelId, NodeId
from exo.shared.types.tasks import (
CancelTask,
ConnectToGroup,
@@ -39,6 +39,7 @@ from exo.shared.types.worker.runners import (
RunnerStatus,
RunnerWarmingUp,
)
from exo.utils.keyed_backoff import KeyedBackoff
from exo.worker.runner.runner_supervisor import RunnerSupervisor
@@ -50,18 +51,22 @@ def plan(
instances: Mapping[InstanceId, Instance],
all_runners: Mapping[RunnerId, RunnerStatus], # all global
tasks: Mapping[TaskId, Task],
input_chunk_buffer: Mapping[CommandId, Mapping[int, InputImageChunk]] | None = None,
input_chunk_buffer: Mapping[CommandId, Mapping[int, InputImageChunk]],
instance_backoff: KeyedBackoff[InstanceId],
download_backoff: KeyedBackoff[ModelId],
) -> Task | None:
# Python short circuiting OR logic should evaluate these sequentially.
return (
_cancel_tasks(runners, tasks)
or _kill_runner(runners, all_runners, instances)
or _create_runner(node_id, runners, instances)
or _model_needs_download(node_id, runners, global_download_status)
or _create_runner(node_id, runners, instances, instance_backoff)
or _model_needs_download(
node_id, runners, global_download_status, download_backoff
)
or _init_distributed_backend(runners, all_runners)
or _load_model(runners, all_runners, global_download_status)
or _ready_to_warmup(runners, all_runners)
or _pending_tasks(runners, tasks, all_runners, input_chunk_buffer or {})
or _pending_tasks(runners, tasks, all_runners, input_chunk_buffer)
)
@@ -92,8 +97,12 @@ def _create_runner(
node_id: NodeId,
runners: Mapping[RunnerId, RunnerSupervisor],
instances: Mapping[InstanceId, Instance],
instance_backoff: KeyedBackoff[InstanceId],
) -> CreateRunner | None:
for instance in instances.values():
if not instance_backoff.should_proceed(instance.instance_id):
continue
runner_id = instance.shard_assignments.node_to_runner.get(node_id, None)
if runner_id is None:
continue
@@ -116,6 +125,7 @@ def _model_needs_download(
node_id: NodeId,
runners: Mapping[RunnerId, RunnerSupervisor],
global_download_status: Mapping[NodeId, Sequence[DownloadProgress]],
download_backoff: KeyedBackoff[ModelId],
) -> DownloadModel | None:
local_downloads = global_download_status.get(node_id, [])
download_status = {
@@ -124,12 +134,16 @@ def _model_needs_download(
for runner in runners.values():
model_id = runner.bound_instance.bound_shard.model_card.model_id
if isinstance(runner.status, RunnerIdle) and (
model_id not in download_status
or not isinstance(
download_status[model_id],
(DownloadOngoing, DownloadCompleted, DownloadFailed),
if (
isinstance(runner.status, RunnerIdle)
and (
model_id not in download_status
or not isinstance(
download_status[model_id],
(DownloadOngoing, DownloadCompleted, DownloadFailed),
)
)
and download_backoff.should_proceed(model_id)
):
# We don't invalidate download_status randomly in case a file gets deleted on disk
return DownloadModel(
@@ -272,7 +286,7 @@ def _pending_tasks(
runners: Mapping[RunnerId, RunnerSupervisor],
tasks: Mapping[TaskId, Task],
all_runners: Mapping[RunnerId, RunnerStatus],
input_chunk_buffer: Mapping[CommandId, Mapping[int, InputImageChunk]] | None,
input_chunk_buffer: Mapping[CommandId, Mapping[int, InputImageChunk]],
) -> Task | None:
for task in tasks.values():
# for now, just forward chat completions
@@ -287,7 +301,6 @@ def _pending_tasks(
if isinstance(task, (ImageEdits, TextGeneration)):
expected_image_chunks = task.task_params.total_input_chunks
if expected_image_chunks > 0:
assert input_chunk_buffer is not None
cmd_id = task.command_id
received = len(input_chunk_buffer.get(cmd_id, {}))
if received < expected_image_chunks:
@@ -8,6 +8,7 @@ from exo.shared.types.worker.runners import (
RunnerConnected,
RunnerIdle,
)
from exo.utils.keyed_backoff import KeyedBackoff
from exo.worker.tests.constants import (
INSTANCE_1_ID,
MODEL_A_ID,
@@ -52,6 +53,9 @@ def test_plan_requests_download_when_waiting_and_shard_not_downloaded():
instances=instances,
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
assert isinstance(result, plan_mod.DownloadModel)
@@ -104,6 +108,9 @@ def test_plan_loads_model_when_all_shards_downloaded_and_waiting():
instances=instances,
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
assert isinstance(result, LoadModel)
@@ -146,6 +153,9 @@ def test_plan_does_not_request_download_when_shard_already_downloaded():
instances=instances,
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
assert not isinstance(result, plan_mod.DownloadModel)
@@ -193,6 +203,9 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
instances=instances,
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
assert result is None
@@ -213,6 +226,9 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
instances=instances,
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
assert result is not None
@@ -9,6 +9,7 @@ from exo.shared.types.worker.runners import (
RunnerReady,
RunnerStatus,
)
from exo.utils.keyed_backoff import KeyedBackoff
from exo.worker.tests.constants import (
INSTANCE_1_ID,
MODEL_A_ID,
@@ -52,6 +53,9 @@ def test_plan_kills_runner_when_instance_missing():
instances=instances,
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
assert isinstance(result, Shutdown)
@@ -91,6 +95,9 @@ def test_plan_kills_runner_when_sibling_failed():
instances=instances,
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
assert isinstance(result, Shutdown)
@@ -122,6 +129,9 @@ def test_plan_creates_runner_when_missing_for_node():
instances=instances,
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
# We patched plan_mod.CreateRunner → CreateRunner
@@ -160,6 +170,9 @@ def test_plan_does_not_create_runner_when_supervisor_already_present():
instances=instances,
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
assert result is None
@@ -189,6 +202,9 @@ def test_plan_does_not_create_runner_for_unassigned_node():
instances=instances,
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
assert result is None
@@ -9,6 +9,7 @@ from exo.shared.types.worker.runners import (
RunnerReady,
RunnerRunning,
)
from exo.utils.keyed_backoff import KeyedBackoff
from exo.worker.tests.constants import (
COMMAND_1_ID,
INSTANCE_1_ID,
@@ -71,6 +72,9 @@ def test_plan_forwards_pending_chat_completion_when_runner_ready():
instances=instances,
all_runners=all_runners,
tasks={TASK_1_ID: task},
input_chunk_buffer={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
assert result is task
@@ -120,6 +124,9 @@ def test_plan_does_not_forward_chat_completion_if_any_runner_not_ready():
instances=instances,
all_runners=all_runners,
tasks={TASK_1_ID: task},
input_chunk_buffer={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
assert result is None
@@ -166,6 +173,9 @@ def test_plan_does_not_forward_tasks_for_other_instances():
instances=instances,
all_runners=all_runners,
tasks={foreign_task.task_id: foreign_task},
input_chunk_buffer={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
assert result is None
@@ -230,6 +240,9 @@ def test_plan_ignores_non_pending_or_non_chat_tasks():
instances=instances,
all_runners=all_runners,
tasks={TASK_1_ID: completed_task, other_task_id: other_task},
input_chunk_buffer={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
assert result is None
@@ -269,6 +282,9 @@ def test_plan_returns_none_when_nothing_to_do():
instances=instances,
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
assert result is None
@@ -7,6 +7,7 @@ from exo.shared.types.worker.runners import (
RunnerLoading,
RunnerWarmingUp,
)
from exo.utils.keyed_backoff import KeyedBackoff
from exo.worker.tests.constants import (
INSTANCE_1_ID,
MODEL_A_ID,
@@ -61,6 +62,9 @@ def test_plan_starts_warmup_for_accepting_rank_when_all_loaded_or_warming():
instances=instances,
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
assert isinstance(result, StartWarmup)
@@ -102,6 +106,9 @@ def test_plan_starts_warmup_for_rank_zero_after_others_warming():
instances=instances,
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
assert isinstance(result, StartWarmup)
@@ -142,6 +149,9 @@ def test_plan_does_not_start_warmup_for_non_zero_rank_until_all_loaded_or_warmin
instances=instances,
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
assert result is None
@@ -186,6 +196,9 @@ def test_plan_does_not_start_warmup_for_rank_zero_until_others_warming():
instances=instances,
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
assert result is None
@@ -202,6 +215,9 @@ def test_plan_does_not_start_warmup_for_rank_zero_until_others_warming():
instances=instances,
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
assert isinstance(result, StartWarmup)
@@ -245,6 +261,9 @@ def test_plan_starts_warmup_for_connecting_rank_after_others_warming():
instances=instances,
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
assert isinstance(result, StartWarmup)
@@ -287,6 +306,9 @@ def test_plan_does_not_start_warmup_for_accepting_rank_until_all_loaded_or_warmi
instances=instances,
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
assert result is None
@@ -328,6 +350,9 @@ def test_plan_does_not_start_warmup_for_connecting_rank_until_others_warming():
instances=instances,
all_runners=all_runners,
tasks={},
input_chunk_buffer={},
instance_backoff=KeyedBackoff(),
download_backoff=KeyedBackoff(),
)
assert result is None