From 00442b3cfd5147c2dbb73f9bb8779dc0e4cdea43 Mon Sep 17 00:00:00 2001 From: Ryuichi Leo Takashige Date: Fri, 23 Jan 2026 12:55:13 +0000 Subject: [PATCH] Add more llm stuff --- bench/exo_eval.py | 39 +++++++++++++++++++++++- bench/lm_eval_patched.py | 11 +++---- src/exo/worker/runner/batched_handler.py | 38 +---------------------- src/exo/worker/runner/runner.py | 6 ++-- 4 files changed, 46 insertions(+), 48 deletions(-) diff --git a/bench/exo_eval.py b/bench/exo_eval.py index eaac95e1..9133239f 100644 --- a/bench/exo_eval.py +++ b/bench/exo_eval.py @@ -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 diff --git a/bench/lm_eval_patched.py b/bench/lm_eval_patched.py index 6096ddfd..4f3faf2f 100644 --- a/bench/lm_eval_patched.py +++ b/bench/lm_eval_patched.py @@ -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 diff --git a/src/exo/worker/runner/batched_handler.py b/src/exo/worker/runner/batched_handler.py index 43076192..2cc00e41 100644 --- a/src/exo/worker/runner/batched_handler.py +++ b/src/exo/worker/runner/batched_handler.py @@ -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 diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py index e583e6ef..9021d131 100644 --- a/src/exo/worker/runner/runner.py +++ b/src/exo/worker/runner/runner.py @@ -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()