From 4d9bee0ff5c8380f5b02d54b896c77728e5f75f6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=83=91=E6=98=8E=E5=86=8D?= Date: Tue, 9 Jun 2026 21:50:08 +0800 Subject: [PATCH] fix(worker): pop launch_ts in recover_model to avoid strict-constructor crashes launch_ts is an internal timestamp stamped onto the launch snapshot at the entry of launch_builtin_model. recover_model splats the whole snapshot back via launch_builtin_model(**launch_args), which injects launch_ts into the model's self._kwargs. Models that forward the full self._kwargs into a strict constructor (e.g. jina-reranker-v3 -> AutoModelForCausalLM.from_pretrained) then crash with TypeError: unexpected keyword argument 'launch_ts'. The cross-session recovery path (recover_models_on_startup at line 788) already pops launch_ts before relaunch. This commit applies the same cleanup to the in-process recovery path (recover_model), copying the dict first to avoid mutating the cached snapshot in self._model_uid_to_launch_args. Fixes OOM auto-recovery for strict-constructor models. Without this fix, recover_sub_pool detects the crash and triggers recover_model, but the recreated model fails to load, leaving the replica permanently down. Co-Authored-By: Claude Opus 4.7 --- xinference/core/tests/test_worker.py | 56 ++++++++++++++++++++++++++++ xinference/core/worker.py | 14 +++++++ 2 files changed, 70 insertions(+) diff --git a/xinference/core/tests/test_worker.py b/xinference/core/tests/test_worker.py index 95342c93c4..9b259ee9b5 100644 --- a/xinference/core/tests/test_worker.py +++ b/xinference/core/tests/test_worker.py @@ -1432,3 +1432,59 @@ async def update_instance_info(self, model_uid: str, updates: dict): assert "model-x-0" not in supervisor._replica_model_uid_to_worker # Failure gauge marker stays lit (mark_replica_dead must not clear it). assert ("model-x", 0) in supervisor._unexpected_down_replicas + + +@pytest.mark.asyncio +async def test_recover_model_pops_launch_ts_from_kwargs(): + """launch_ts is an internal timestamp stamped onto the launch snapshot at + launch_builtin_model entry. recover_model splats the snapshot back via + launch_builtin_model(**launch_args), which would inject launch_ts into the + model's self._kwargs. Models that forward the full self._kwargs into a + strict constructor (e.g. jina-reranker-v3 -> AutoModelForCausalLM.from_pretrained) + then crash with TypeError. recover_model must pop launch_ts before the splat, + mirroring the existing cleanup in recover_models_on_startup.""" + + # Use a minimal mock object to avoid WorkerActor initialization complexity + class _MockWorker: + async def launch_builtin_model(self, **kwargs): + self._captured_kwargs = kwargs + return "mock-subpool-address" + + async def get_supervisor_ref(self, add_worker=False): + return None + + worker = _MockWorker() + + # Simulate a cached launch snapshot with launch_ts (the internal timestamp) + launch_args = { + "model_uid": "test-model-0", + "model_name": "test-model", + "launch_ts": 1780900592, # Internal timestamp, not a model param + "some_param": "value", + } + original_launch_args = dict(launch_args) + + # Mock parse_replica_model_uid to avoid dependency + def mock_parse(uid): + return ("test-model", 0) + + import xinference.core.worker as worker_module + + original_parse = worker_module.parse_replica_model_uid + worker_module.parse_replica_model_uid = mock_parse + + try: + # Call recover_model directly (it's a standalone async method) + await WorkerActor.recover_model(worker, launch_args) + + # Assert: launch_builtin_model was called without launch_ts + assert hasattr(worker, "_captured_kwargs") + assert "launch_ts" not in worker._captured_kwargs + assert worker._captured_kwargs["model_uid"] == "test-model-0" + assert worker._captured_kwargs["some_param"] == "value" + + # Assert: original launch_args still has launch_ts (copy isolation) + assert launch_args == original_launch_args + assert launch_args["launch_ts"] == 1780900592 + finally: + worker_module.parse_replica_model_uid = original_parse diff --git a/xinference/core/worker.py b/xinference/core/worker.py index 023a5369ab..523aeb36fd 100644 --- a/xinference/core/worker.py +++ b/xinference/core/worker.py @@ -2871,6 +2871,20 @@ async def launch_rank0_model( @no_type_check async def recover_model(self, launch_args: Dict[str, Any]): + # `launch_ts` is an internal timestamp stamped onto the launch snapshot at + # the entry of `launch_builtin_model` (see the `launch_args["launch_ts"]` + # assignment); it is not a model construction parameter. recover_model + # splats the whole snapshot back via `launch_builtin_model(**launch_args)`, + # so `launch_ts` would land in that call's `**kwargs` and sink into the + # model's `self._kwargs`. Models that forward the full `self._kwargs` into a + # strict constructor (e.g. jina-reranker-v3 -> + # `AutoModelForCausalLM.from_pretrained(**model_kwargs)`) then crash with + # `TypeError: ... unexpected keyword argument 'launch_ts'`. The cross-session + # recovery path (`recover_models_on_startup`) already pops it before + # relaunch; do the same here. Copy first so the cached snapshot in + # `self._model_uid_to_launch_args` keeps its `launch_ts` (used as created_ts). + launch_args = dict(launch_args) + launch_args.pop("launch_ts", None) rep_model_uid = launch_args.get("model_uid") origin_uid, _ = parse_replica_model_uid(rep_model_uid) xavier_config: Optional[Dict[str, Any]] = launch_args.get("xavier_config", None)