fix: harden ModelWeightsManager background task reliability

- Hold a strong set reference to each asyncio.Task to prevent GC cancellation
- Deduplicate concurrent sync requests for the same hf_repo via _in_progress dict
- Surface task exceptions via logger.error in _on_task_done callback
- Store ModelWeightsManager as app.state singleton so state persists across requests
- Add recover_hf_syncs startup handler to re-trigger syncs after server restart

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>

Author: lilyz-ai

model-engine/model_engine_server/api/app.py

@@ -312,6 +312,56 @@ def load_redis():
     get_or_create_aioredis_pool()
 
 
+@app.on_event("startup")
+def init_model_weights_manager():
+    from model_engine_server.core.config import infra_config
+    from model_engine_server.domain.use_cases.model_weights_manager import ModelWeightsManager
+    from model_engine_server.infra.gateways import (
+        ABSLLMArtifactGateway,
+        GCSLLMArtifactGateway,
+        S3LLMArtifactGateway,
+    )
+
+    provider = infra_config().cloud_provider
+    if provider == "azure":
+        gateway = ABSLLMArtifactGateway()
+    elif provider == "gcp":
+        gateway = GCSLLMArtifactGateway()
+    else:
+        gateway = S3LLMArtifactGateway()
+    app.state.model_weights_manager = ModelWeightsManager(llm_artifact_gateway=gateway)
+
+
+@app.on_event("startup")
+async def recover_hf_syncs():
+    """Re-trigger weight syncs for endpoints that were syncing when server last stopped."""
+    from model_engine_server.db.base import get_session_async
+    from model_engine_server.infra.repositories.live_tokenizer_repository import (
+        SUPPORTED_MODELS_INFO,
+    )
+    from sqlalchemy import text
+
+    session_factory = get_session_async()
+    try:
+        async with session_factory() as session:
+            result = await session.execute(
+                text(
+                    "SELECT DISTINCT endpoint_metadata->'_llm'->>'model_name' AS model_name "
+                    "FROM endpoints "
+                    "WHERE (endpoint_metadata->'_llm'->>'hf_weights_syncing')::boolean = true"
+                )
+            )
+            model_names = [row.model_name for row in result if row.model_name]
+    except Exception:
+        logger.warning("Could not query pending HF sync endpoints at startup", exc_info=True)
+        return
+    for model_name in model_names:
+        info = SUPPORTED_MODELS_INFO.get(model_name)
+        if info and info.hf_repo:
+            app.state.model_weights_manager.ensure_model_weights_available(info.hf_repo)
+            logger.info(f"Startup: re-triggered HF weight sync for {model_name}")
+
+
 def healthcheck() -> Response:
     """Returns 200 if the app is healthy."""
     return Response(status_code=200)

model-engine/model_engine_server/api/llms_v1.py

@@ -86,7 +86,6 @@
     UpdateLLMModelEndpointV1UseCase,
 )
 from model_engine_server.domain.use_cases.model_bundle_use_cases import CreateModelBundleV2UseCase
-from model_engine_server.domain.use_cases.model_weights_manager import ModelWeightsManager
 from pydantic import RootModel
 from sse_starlette.sse import EventSourceResponse
 
@@ -149,14 +148,15 @@ def handle_streaming_exception(
 @llm_router_v1.post("/model-endpoints", response_model=CreateLLMModelEndpointV1Response)
 async def create_model_endpoint(
     wrapped_request: RootModel[CreateLLMModelEndpointV1Request],
+    request: Request,
     auth: User = Depends(verify_authentication),
     external_interfaces: ExternalInterfaces = Depends(get_external_interfaces),
 ) -> CreateLLMModelEndpointV1Response:
-    request = wrapped_request.root
+    llm_request = wrapped_request.root
     """
     Creates an LLM endpoint for the current user.
     """
-    logger.info(f"POST /llm/model-endpoints with {request} for {auth}")
+    logger.info(f"POST /llm/model-endpoints with {llm_request} for {auth}")
     try:
         create_model_bundle_use_case = CreateModelBundleV2UseCase(
             model_bundle_repository=external_interfaces.model_bundle_repository,
@@ -169,17 +169,15 @@ async def create_model_endpoint(
             llm_artifact_gateway=external_interfaces.llm_artifact_gateway,
             docker_repository=external_interfaces.docker_repository,
         )
-        model_weights_manager = ModelWeightsManager(
-            llm_artifact_gateway=external_interfaces.llm_artifact_gateway,
-        )
+        model_weights_manager = request.app.state.model_weights_manager
         use_case = CreateLLMModelEndpointV1UseCase(
             create_llm_model_bundle_use_case=create_llm_model_bundle_use_case,
             model_endpoint_service=external_interfaces.model_endpoint_service,
             docker_repository=external_interfaces.docker_repository,
             llm_artifact_gateway=external_interfaces.llm_artifact_gateway,
             model_weights_manager=model_weights_manager,
         )
-        return await use_case.execute(user=auth, request=request)
+        return await use_case.execute(user=auth, request=llm_request)
     except ObjectAlreadyExistsException as exc:
         raise HTTPException(
             status_code=400,

model-engine/model_engine_server/domain/use_cases/model_weights_manager.py

@@ -1,7 +1,7 @@
 import asyncio
 import functools
 import tempfile
-from typing import List
+from typing import Dict, List, Set
 
 from huggingface_hub import snapshot_download
 from model_engine_server.common.config import hmi_config
@@ -24,6 +24,8 @@
 class ModelWeightsManager:
     def __init__(self, llm_artifact_gateway: LLMArtifactGateway):
         self.llm_artifact_gateway = llm_artifact_gateway
+        self._background_tasks: Set[asyncio.Task] = set()
+        self._in_progress: Dict[str, asyncio.Task] = {}
 
     def get_remote_path(self, hf_repo: str) -> str:
         prefix = hmi_config.hf_user_fine_tuned_weights_prefix.rstrip("/")
@@ -38,16 +40,35 @@ def ensure_model_weights_available(self, hf_repo: str) -> str:
         Callers receive the checkpoint path right away and can proceed with
         any following actions (e.g. endpoint creation) without blocking.
 
+        A second call for the same ``hf_repo`` while a sync is already in
+        progress is a no-op: the existing task is reused and the same remote
+        path is returned.
+
         Args:
             hf_repo: HuggingFace repository ID, e.g. ``"meta-llama/Meta-Llama-3-8B"``.
 
         Returns:
             The remote path (s3://, gs://, or https://) where the weights will be stored.
         """
         remote_path = self.get_remote_path(hf_repo)
-        asyncio.create_task(self._sync_weights(hf_repo, remote_path))
+        if hf_repo not in self._in_progress:
+            task = asyncio.create_task(self._sync_weights(hf_repo, remote_path))
+            self._background_tasks.add(task)
+            self._in_progress[hf_repo] = task
+            task.add_done_callback(lambda t: self._on_task_done(t, hf_repo))
         return remote_path
 
+    def _on_task_done(self, task: asyncio.Task, hf_repo: str) -> None:
+        self._background_tasks.discard(task)
+        self._in_progress.pop(hf_repo, None)
+        if not task.cancelled():
+            exc = task.exception()
+            if exc:
+                logger.error(
+                    f"Background weight sync failed for {hf_repo}: {exc}",
+                    exc_info=exc,
+                )
+
     async def _sync_weights(self, hf_repo: str, remote_path: str) -> None:
         """Downloads weights from HuggingFace Hub and uploads to remote storage if not cached."""
         files = self.llm_artifact_gateway.list_files(remote_path)

model-engine/tests/unit/domain/test_model_weights_manager.py

@@ -116,6 +116,61 @@ def test_s3_path_construction(monkeypatch):
     assert path == "s3://bucket/prefix/myorg/mymodel"
 
 
+def test_deduplication_same_hf_repo():
+    """Second call for same hf_repo while a sync is in progress should not create a new task."""
+    gateway = FakeArtifactGateway(existing_files=[])
+    manager = ModelWeightsManager(llm_artifact_gateway=gateway)
+
+    mwm_base = "model_engine_server.domain.use_cases.model_weights_manager"
+    with patch(f"{mwm_base}.asyncio.create_task") as mock_create_task:
+        result1 = manager.ensure_model_weights_available("org/model")
+        result2 = manager.ensure_model_weights_available("org/model")
+
+    assert mock_create_task.call_count == 1
+    assert result1 == result2
+
+
+def test_task_reference_held_until_done():
+    """_background_tasks should hold a reference to the task until _on_task_done fires."""
+    gateway = FakeArtifactGateway(existing_files=[])
+    manager = ModelWeightsManager(llm_artifact_gateway=gateway)
+
+    mwm_base = "model_engine_server.domain.use_cases.model_weights_manager"
+    mock_task = MagicMock()
+    with patch(f"{mwm_base}.asyncio.create_task", return_value=mock_task):
+        manager.ensure_model_weights_available("org/model")
+
+    assert mock_task in manager._background_tasks
+    assert "org/model" in manager._in_progress
+
+    # Simulate successful task completion via the done callback
+    mock_task.cancelled.return_value = False
+    mock_task.exception.return_value = None
+    manager._on_task_done(mock_task, "org/model")
+
+    assert mock_task not in manager._background_tasks
+    assert "org/model" not in manager._in_progress
+
+
+def test_error_surfaced_on_task_failure():
+    """When the background task raises, _on_task_done should log the error."""
+    gateway = FakeArtifactGateway(existing_files=[])
+    manager = ModelWeightsManager(llm_artifact_gateway=gateway)
+
+    mock_task = MagicMock()
+    mock_task.cancelled.return_value = False
+    exc = RuntimeError("Download failed")
+    mock_task.exception.return_value = exc
+
+    mwm_base = "model_engine_server.domain.use_cases.model_weights_manager"
+    with patch(f"{mwm_base}.logger") as mock_logger:
+        manager._on_task_done(mock_task, "org/model")
+        mock_logger.error.assert_called_once()
+        call_args = mock_logger.error.call_args
+        assert "org/model" in call_args[0][0]
+        assert call_args[1]["exc_info"] == exc
+
+
 @pytest.mark.asyncio
 async def test_create_llm_model_endpoint_calls_weights_manager_on_hf_source():
     """CreateLLMModelEndpointV1UseCase should call ensure_model_weights_available (sync),