electron-rare/exo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add more llm stuff

Browse Source
This commit is contained in:
Ryuichi Leo Takashige
2026-01-23 12:55:13 +00:00
parent aa41da8541
commit 00442b3cfd
4 changed files with 46 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
bench/exo_eval.py
+38 -1
View File
@@ -202,7 +202,14 @@ def teardown_instance(client: ExoClient, instance_id: str) -> None:
except ExoHttpError as e:
if e.status != 404:
raise
wait_for_instance_gone(client, instance_id)
except (ConnectionRefusedError, OSError):
logger.warning(f"Could not connect to exo to delete instance {instance_id} (server may be down)")
return
try:
wait_for_instance_gone(client, instance_id)
except (ConnectionRefusedError, OSError, TimeoutError):
logger.warning("Could not verify instance deletion (server may be down)")
return
logger.info(f"Instance {instance_id} deleted")
@@ -516,6 +523,19 @@ def main() -> int:
action="store_true",
help="Skip instance creation (assume instance already running)",
)
ap.add_argument(
"--pipeline",
type=int,
default=None,
metavar="N",
help="Use pipeline sharding with exactly N nodes (overrides config)",
)
ap.add_argument(
"--instance-meta",
choices=["ring", "jaccl", "both"],
default=None,
help="Instance meta preference (overrides config)",
)
ap.add_argument(
"--dry-run",
action="store_true",
@@ -539,6 +559,23 @@ def main() -> int:
logger.info(f"Model: {args.model}")
logger.info(f"API endpoint: http://{args.host}:{args.port}/v1")
# Apply CLI overrides to instance config
if args.pipeline is not None or args.instance_meta is not None:
instance_config = config.setdefault("instance", {})
if args.pipeline is not None:
instance_config["sharding"] = "pipeline"
instance_config["min_nodes"] = args.pipeline
instance_config["max_nodes"] = args.pipeline
logger.info(f"CLI override: pipeline={args.pipeline} nodes")
# Limit concurrency for pipeline to avoid GPU timeouts
if args.pipeline >= 2:
lm_eval_config = config.setdefault("lm_eval", {})
lm_eval_config["num_concurrent"] = 8
logger.info("CLI override: num_concurrent=8 (pipeline>=2)")
if args.instance_meta is not None:
instance_config["instance_meta"] = args.instance_meta
logger.info(f"CLI override: instance_meta={args.instance_meta}")
# Check HuggingFace token if required
if not check_hf_token(config):
return 1
bench/lm_eval_patched.py
+5 -6
View File
@@ -25,12 +25,11 @@ def _patch_amodel_call() -> None:
async def patched_amodel_call(self: Any, *args: Any, **kwargs: Any) -> Any:
try:
return await original(self, *args, **kwargs)
except UnboundLocalError:
# `outputs` referenced before assignment when response.raise_for_status() throws
return []
except Exception:
# After all retries fail, don't crash the entire eval
return []
except (UnboundLocalError, Exception):
# Return one empty-string result per request in the batch so the
# reorderer doesn't assert on missing coverage.
messages = kwargs.get("messages") or (args[2] if len(args) > 2 else [])
return [""] * max(len(messages), 1)
TemplateAPI.amodel_call = patched_amodel_call
src/exo/worker/runner/batched_handler.py
+1 -37
View File
@@ -79,14 +79,12 @@ class BatchedInferenceHandler:
model_id: ModelId,
device_rank: int,
max_batch_size: int = 8,
batch_timeout_ms: int = 50,
):
self.model = model
self.tokenizer = tokenizer
self.model_id = model_id
self.device_rank = device_rank
self.max_batch_size = max_batch_size
self.batch_timeout_ms = batch_timeout_ms
# GPT-OSS model detection
self.is_gpt_oss = isinstance(model, GptOssModel)
@@ -97,7 +95,6 @@ class BatchedInferenceHandler:
# Pending requests waiting to be batched
self.pending: list[PendingRequest] = []
self.pending_start_time: float | None = None
# Active batch generator and request tracking
self.batch_generator: BatchGenerator | None = None
@@ -154,36 +151,11 @@ class BatchedInferenceHandler:
)
self.pending.append(pending_request)
if self.pending_start_time is None:
self.pending_start_time = time.perf_counter()
logger.info(
f"Added request to batch queue (pending={len(self.pending)}, active={self.current_batch_size})"
)
def should_flush(self) -> bool:
"""
Determine if the pending batch should be flushed.
Returns True if:
- We have pending requests AND (batch is full OR timeout reached)
"""
if not self.has_pending:
return False
# Check if batch is full
available_slots = self.max_batch_size - self.current_batch_size
if len(self.pending) >= available_slots:
return True
# Check timeout
if self.pending_start_time is not None:
elapsed_ms = (time.perf_counter() - self.pending_start_time) * 1000
if elapsed_ms >= self.batch_timeout_ms:
return True
return False
def flush(self) -> None:
"""Start processing pending requests by adding them to the BatchGenerator."""
if not self.has_pending:
@@ -194,9 +166,6 @@ class BatchedInferenceHandler:
requests_to_flush = self.pending[:available_slots]
self.pending = self.pending[available_slots:]
if len(self.pending) == 0:
self.pending_start_time = None
# Create batch generator if not exists
if self.batch_generator is None:
logger.info(f"Creating new BatchGenerator for {len(requests_to_flush)} requests")
@@ -205,7 +174,7 @@ class BatchedInferenceHandler:
model=self.model,
max_tokens=MAX_TOKENS,
stop_tokens=self.stop_tokens if self.stop_tokens else None,
prefill_batch_size=min(len(requests_to_flush), 8),
prefill_batch_size=1,
)
else:
logger.info(f"Adding {len(requests_to_flush)} requests to existing BatchGenerator")
@@ -379,10 +348,6 @@ class BatchedInferenceHandler:
for uid in completed_uids:
del self.uid_to_request[uid]
# Close batch generator if no more active requests
if not self.uid_to_request and not self.pending:
self._close_generator()
def emit_error(self, command_id: CommandId, error_message: str) -> Event:
"""Create an error event for a failed request."""
return ChunkGenerated(
@@ -406,4 +371,3 @@ class BatchedInferenceHandler:
"""Close the handler and clean up resources."""
self._close_generator()
self.pending.clear()
self.pending_start_time = None
src/exo/worker/runner/runner.py
+2 -4
View File
@@ -90,7 +90,6 @@ from exo.worker.runner.bootstrap import logger
# Batching configuration
BATCH_ENABLED = True
BATCH_MAX_SIZE = 128
BATCH_TIMEOUT_MS = 20 # Short timeout - flush quickly to avoid request timeouts
def _should_use_serial_processing(
@@ -305,7 +304,6 @@ def main(
model_id=shard_metadata.model_card.model_id,
device_rank=device_rank,
max_batch_size=BATCH_MAX_SIZE,
batch_timeout_ms=BATCH_TIMEOUT_MS,
)
logger.info(
f"Batch handler initialized (max_batch_size={BATCH_MAX_SIZE})"
@@ -735,8 +733,8 @@ def main(
except EndOfStream:
break
# Flush batch if ready
if batch_handler.should_flush():
# Flush pending requests immediately (no timeout delay)
if batch_handler.has_pending:
logger.info(f"Flushing batch (pending={len(batch_handler.pending)}, active={batch_handler.current_batch_size})")
batch_handler.flush()