diff --git a/moe_infinity/common/constants.py b/moe_infinity/common/constants.py index fd6420af..41d72bd8 100644 --- a/moe_infinity/common/constants.py +++ b/moe_infinity/common/constants.py @@ -39,7 +39,7 @@ "gptoss": 4, "qwen3": 5, "dbrx": 4, - "olmoe": 4, + "olmoe": 5, "jamba": 4, } diff --git a/moe_infinity/entrypoints/openai/api_server_v2.py b/moe_infinity/entrypoints/openai/api_server_v2.py index bc04631d..6a846043 100644 --- a/moe_infinity/entrypoints/openai/api_server_v2.py +++ b/moe_infinity/entrypoints/openai/api_server_v2.py @@ -142,6 +142,7 @@ def _decorator(func: Any) -> Any: _contextpilot_fallback_count: int = 0 _contextpilot_last_fallback_count: int = 0 +logger = logging.getLogger(__name__) _cp_logger = logging.getLogger("moe_infinity.contextpilot") _cp_middleware: Optional[Any] = None _eviction_sync: Optional[Any] = None @@ -1093,12 +1094,19 @@ async def _initialize_model() -> None: ) _ensure_engine_loop_running() + _health_state.set_healthy() + except Exception as exc: + # _model_init_task is fire-and-forget (asyncio.create_task, never + # awaited) -- without this, an exception here is silently dropped: + # the process keeps running, /health stays stuck on "starting" + # forever, and every request gets a 503 indistinguishable from a + # genuine hang. Log it and surface it through /health instead. + logger.exception("Model initialization failed") + _health_state.set_unhealthy(f"{type(exc).__name__}: {exc}") finally: if _startup_watchdog is not None: _startup_watchdog.cancel() - _health_state.set_healthy() - @app.on_event("startup") async def startup_event() -> None: diff --git a/moe_infinity/utils/hf_config.py b/moe_infinity/utils/hf_config.py index bed2e19c..b1a3b7b0 100644 --- a/moe_infinity/utils/hf_config.py +++ b/moe_infinity/utils/hf_config.py @@ -78,7 +78,7 @@ def parse_moe_param(config: PretrainedConfig) -> Tuple[int, int, int]: num_decoder_layers = config.num_hidden_layers num_layers = config.num_hidden_layers num_experts = config.num_local_experts - elif "qwen3" in arch: + elif "qwen3" in arch or "olmoe" in arch: num_encoder_layers = 0 num_decoder_layers = config.num_hidden_layers num_layers = config.num_hidden_layers @@ -147,7 +147,7 @@ def parse_expert_id( layer_id, expert_id = result[0] layer_id = int(layer_id) expert_id = int(expert_id) - elif "deepseek" in arch or "qwen3" in arch: + elif "deepseek" in arch or "qwen3" in arch or "olmoe" in arch: decoder_sparse_step = 1 layer_type = "decoder" diff --git a/pyproject.toml b/pyproject.toml index b11a21c0..c31eb89c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["setuptools==75.3.2", "wheel", "torch"] +requires = ["setuptools>=75.3.2", "wheel", "torch"] build-backend = "setuptools.build_meta"