Improve model status alerts

Author: PaulHax

align_app/adm/decider/__init__.py

@@ -1,12 +1,12 @@
 from align_utils.models import ADMResult, Decision, ChoiceInfo
 from .decider import MultiprocessDecider
-from .client import get_decision, is_model_cached
+from .client import get_decision, get_model_cache_status
 from .types import DeciderParams
 
 __all__ = [
     "MultiprocessDecider",
     "get_decision",
-    "is_model_cached",
+    "get_model_cache_status",
     "DeciderParams",
     "ADMResult",
     "Decision",

align_app/adm/decider/client.py

@@ -7,6 +7,7 @@
 from typing import Dict, Any
 from align_utils.models import ADMResult
 from .decider import MultiprocessDecider
+from .worker import CacheQueryResult
 from .types import DeciderParams
 
 _decider = None
@@ -33,10 +34,12 @@ async def get_decision(params: DeciderParams) -> ADMResult:
     return await process_manager.get_decision(params)
 
 
-async def is_model_cached(resolved_config: Dict[str, Any]) -> bool:
-    """Check if model for this config is already loaded in worker."""
+async def get_model_cache_status(
+    resolved_config: Dict[str, Any]
+) -> CacheQueryResult | None:
+    """Get best-effort model cache status (memory + disk)."""
     process_manager = _get_process_manager()
-    return await process_manager.is_model_cached(resolved_config)
+    return await process_manager.get_model_cache_status(resolved_config)
 
 
 def cleanup():

align_app/adm/decider/decider.py

@@ -14,11 +14,13 @@ class MultiprocessDecider:
     def __init__(self):
         self.worker: WorkerHandle = create_worker(decider_worker_func)
 
-    async def is_model_cached(self, resolved_config: Dict[str, Any]) -> bool:
+    async def get_model_cache_status(
+        self, resolved_config: Dict[str, Any]
+    ) -> CacheQueryResult | None:
         self.worker, result = await send(self.worker, CacheQuery(resolved_config))
         if isinstance(result, CacheQueryResult):
-            return result.is_cached
-        return False
+            return result
+        return None
 
     async def get_decision(self, params: DeciderParams) -> ADMResult:
         self.worker, result = await send(self.worker, params)

align_app/adm/decider/worker.py

@@ -2,9 +2,10 @@
 import hashlib
 import json
 import logging
+import os
 import traceback
 from dataclasses import dataclass
-from typing import Dict, Tuple, Callable, Any
+from typing import Dict, Tuple, Callable, Any, Optional
 from multiprocessing import Queue
 from align_utils.models import ADMResult
 from .executor import instantiate_adm
@@ -24,11 +25,57 @@ class CacheQuery:
 @dataclass
 class CacheQueryResult:
     is_cached: bool
+    is_downloaded: Optional[bool]
+
+
+def _extract_model_name(resolved_config: Dict[str, Any]) -> Optional[str]:
+    if not isinstance(resolved_config, dict):
+        return None
+
+    if isinstance(resolved_config.get("model_name"), str):
+        return resolved_config["model_name"]
+
+    structured = resolved_config.get("structured_inference_engine")
+    if isinstance(structured, dict) and isinstance(structured.get("model_name"), str):
+        return structured["model_name"]
+
+    for value in resolved_config.values():
+        if isinstance(value, dict):
+            found = _extract_model_name(value)
+            if found:
+                return found
+        elif isinstance(value, list):
+            for item in value:
+                if isinstance(item, dict):
+                    found = _extract_model_name(item)
+                    if found:
+                        return found
+    return None
+
+
+def _is_model_downloaded(model_name: Optional[str]) -> Optional[bool]:
+    if not model_name:
+        return None
+
+    if os.path.exists(model_name):
+        return True
+
+    try:
+        from huggingface_hub import snapshot_download
+    except Exception:
+        return None
+
+    try:
+        snapshot_download(model_name, local_files_only=True)
+        return True
+    except Exception:
+        return False
 
 
 def decider_worker_func(task_queue: Queue, result_queue: Queue):
     root_logger = logging.getLogger()
     root_logger.setLevel("WARNING")
+    logger = logging.getLogger(__name__)
 
     model_cache: Dict[str, Tuple[Callable, Callable]] = {}
 
@@ -37,8 +84,15 @@ def decider_worker_func(task_queue: Queue, result_queue: Queue):
             try:
                 if isinstance(task, CacheQuery):
                     cache_key = extract_cache_key(task.resolved_config)
+                    is_cached = cache_key in model_cache
+                    is_downloaded = True if is_cached else _is_model_downloaded(
+                        _extract_model_name(task.resolved_config)
+                    )
                     result_queue.put(
-                        CacheQueryResult(is_cached=cache_key in model_cache)
+                        CacheQueryResult(
+                            is_cached=is_cached,
+                            is_downloaded=is_downloaded,
+                        )
                     )
                     continue
 
@@ -72,11 +126,30 @@ def decider_worker_func(task_queue: Queue, result_queue: Queue):
             except (KeyboardInterrupt, SystemExit):
                 break
             except Exception as e:
-                error_msg = f"{str(e)}\n{traceback.format_exc()}"
+                logger.error("Worker error:\n%s", traceback.format_exc())
+                error_msg = _format_worker_error(e)
                 result_queue.put(Exception(error_msg))
     finally:
         for _, (_, cleanup_func) in model_cache.items():
             try:
                 cleanup_func()
             except Exception:
                 pass
+
+
+def _format_worker_error(error: Exception) -> str:
+    error_text = str(error)
+    gated_tokens = (
+        "GatedRepoError",
+        "gated repo",
+        "401 Client Error",
+        "Access to model",
+        "restricted",
+        "Please log in",
+    )
+    if any(token in error_text for token in gated_tokens):
+        return (
+            "Model access denied. Authenticate with Hugging Face or request access "
+            "to the gated repo."
+        )
+    return f"{error_text}\n{traceback.format_exc()}"

align_app/app/runs_state_adapter.py

@@ -7,7 +7,7 @@
 from .runs_registry import RunsRegistry
 from .runs_table_filter import RunsTableFilter
 from ..adm.decider.types import DeciderParams
-from ..adm.decider import is_model_cached
+from ..adm.decider import get_model_cache_status
 from ..adm.system_adm_discovery import discover_system_adms
 from ..utils.utils import get_id
 from .runs_presentation import extract_base_scenarios
@@ -614,16 +614,18 @@ async def _execute_run_decision(self, run_id: str):
 
         run = self.runs_registry.get_run(run_id)
         is_cached_decision = self.runs_registry.has_cached_decision(run_id)
-        is_model_loaded = False
+        status = None
         if run:
-            is_model_loaded = await is_model_cached(run.decider_params.resolved_config)
+            status = await get_model_cache_status(run.decider_params.resolved_config)
 
-        if is_cached_decision or is_model_loaded:
-            alert_id = self._alerts.create_info_alert(title="Deciding...", timeout=0)
+        if is_cached_decision or (status and status.is_cached):
+            alert_title = "Deciding..."
+        elif status and status.is_downloaded is False:
+            alert_title = "Downloading model and deciding..."
         else:
-            alert_id = self._alerts.create_info_alert(
-                title="Loading model and deciding...", timeout=0
-            )
+            alert_title = "Loading model and deciding..."
+
+        alert_id = self._alerts.create_info_alert(title=alert_title, timeout=0)
         await self.server.network_completion
 
         try:
@@ -632,7 +634,15 @@ async def _execute_run_decision(self, run_id: str):
             self._alerts.create_info_alert(title="Decision complete", timeout=3000)
         except Exception as e:
             self._alerts.remove_alert(alert_id)
-            self._alerts.create_info_alert(title=f"Decision failed: {e}", timeout=5000)
+            error_text = str(e)
+            if "Model access denied" in error_text:
+                message = (
+                    "Decision failed: Model access denied. "
+                    "Authenticate with Hugging Face or request access to the model."
+                )
+            else:
+                message = f"Decision failed: {e}"
+            self._alerts.create_info_alert(title=message, timeout=8000)
 
         with self.state:
             self._rebuild_comparison_runs()

align_app/app/ui.py

@@ -1564,6 +1564,9 @@ def __init__(
                         ".drop-zone-active { outline: 3px dashed #1976d2 !important; outline-offset: -3px; }"
                         ".alert-popup-container { left: auto; right: 0; transform: none; width: fit-content; }"
                         ".alert-popup-container .v-alert { --v-theme-info: 66, 66, 66; }"
+                        ".alert-popup-container .v-alert__icon { display: none; }"
+                        ".alert-popup-container .v-alert__prepend { display: none; }"
+                        ".alert-popup-container .v-alert__prepend .v-icon { display: none; }"
                         "</style>'"
                     )
                 )