Update main.py

main.py

@@ -24,13 +24,11 @@ ARTIFACT_URL = os.getenv("ARTIFACT_URL", "https://console.gmicloud.ai/api/v1/ie/
 REFRESH_INTERVAL_SECONDS = int(os.getenv("REFRESH_INTERVAL_SECONDS", "30"))
 
 # --- Shared Cache File Configuration ---
-# Using /dev/shm is faster as it's a RAM disk on Linux. Fallback to /tmp.
 CACHE_DIR = "/dev/shm" if os.path.exists("/dev/shm") else "/tmp"
 CACHE_FILE_PATH = os.path.join(CACHE_DIR, "gmi_routing_table.json")
 LOCK_FILE_PATH = os.path.join(CACHE_DIR, "gmi_routing_table.lock")
 
 # --- In-Memory State for each Worker ---
-# These are global variables *per worker process*.
 worker_model_routing_table = {}
 last_cache_check_time = 0
 
@@ -48,28 +46,22 @@ async def load_or_refresh_models():
     global last_cache_check_time, worker_model_routing_table
 
     now = time.monotonic()
-    # 1. Quick check: If in-memory cache is fresh, do nothing.
     if (now - last_cache_check_time) < REFRESH_INTERVAL_SECONDS:
         return
 
-    # 2. In-memory cache is stale, acquire a lock to check the shared file cache.
-    #    The lock prevents all workers from hitting the API at once.
     lock = FileLock(LOCK_FILE_PATH)
     try:
-        with lock.acquire(timeout=5): # Wait max 5s for the lock
-            # Re-check inside the lock, another process might have just updated the file.
+        with lock.acquire(timeout=5):
             if os.path.exists(CACHE_FILE_PATH):
                 mtime = os.path.getmtime(CACHE_FILE_PATH)
                 if (time.time() - mtime) < REFRESH_INTERVAL_SECONDS:
-                    # File is fresh, just load it into this worker's memory
                     with open(CACHE_FILE_PATH, 'r') as f:
                         worker_model_routing_table = json.load(f)
                     last_cache_check_time = now
-                    logger.info(f"Loaded fresh cache from file. {len(worker_model_routing_table)} models.")
+                    logger.debug(f"Loaded fresh cache from file. {len(worker_model_routing_table)} models.")
                     return
 
-            # 3. We have the lock and the file cache is stale. This worker will be the updater.
-            logger.warning("Cache is stale. This worker is refreshing the model list...")
+            logger.warning("Cache is stale. This worker is refreshing the model list from the API...")
             try:
                 async with httpx.AsyncClient() as client:
                     response = await client.get(ARTIFACT_URL, timeout=30.0)
@@ -78,39 +70,43 @@ async def load_or_refresh_models():
 
                 new_routing_table = {}
                 for artifact in artifacts:
-                    model_name = artifact.get("artifact_metadata", {}).get("artifact_name")
+                    # CORRECTLY get the API model identifier from `model_price.modelName`
+                    model_price_info = artifact.get("model_price")
+                    api_model_id = model_price_info.get("modelName") if model_price_info else None
+                    
+                    display_name = artifact.get("artifact_metadata", {}).get("artifact_name", "Unknown")
                     endpoints = artifact.get("endpoints", [])
-                    if model_name and endpoints and endpoints[0].get("endpoint_url"):
-                        new_routing_table[model_name] = endpoints[0]["endpoint_url"]
+                    
+                    # Condition: Must have a valid API model ID and at least one active endpoint URL
+                    if api_model_id and endpoints and endpoints[0].get("endpoint_url"):
+                        endpoint_url = endpoints[0]["endpoint_url"]
+                        new_routing_table[api_model_id] = endpoint_url
+                        logger.debug(f"Mapped model ID '{api_model_id}' to endpoint '{endpoint_url}'")
+                    else:
+                        logger.debug(f"Skipping model '{display_name}': Missing API model ID or active endpoint.")
 
-                # Write to a temporary file and then atomically rename it
                 temp_path = CACHE_FILE_PATH + f".{os.getpid()}"
                 with open(temp_path, 'w') as f:
                     json.dump(new_routing_table, f)
                 os.rename(temp_path, CACHE_FILE_PATH)
 
                 worker_model_routing_table = new_routing_table
-                logger.info(f"Successfully refreshed cache file with {len(worker_model_routing_table)} models.")
+                logger.info(f"Successfully refreshed cache file with {len(worker_model_routing_table)} active models.")
 
             except Exception as e:
                 logger.error(f"Failed to refresh model cache: {e}. Will use stale data if available.")
     
     except Timeout:
-        logger.warning("Could not acquire lock to refresh cache, another process is likely updating. Reading from file.")
-    
+        logger.warning("Could not acquire lock, another process is updating. Reading from file.")
     except Exception as e:
         logger.error(f"An unexpected error occurred in cache management: {e}")
-
     finally:
-        # 4. Ensure this worker's memory is up-to-date from the file,
-        #    especially if it failed to get the lock or an error occurred.
         if os.path.exists(CACHE_FILE_PATH):
             try:
                 with open(CACHE_FILE_PATH, 'r') as f:
                     worker_model_routing_table = json.load(f)
             except (json.JSONDecodeError, FileNotFoundError):
                  logger.error("Could not read cache file. Routing table may be empty.")
-
         last_cache_check_time = now
 
 
@@ -118,9 +114,7 @@ async def load_or_refresh_models():
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    """Manages the app's lifecycle for startup and shutdown."""
     app.state.http_client = httpx.AsyncClient(timeout=None)
-    # Perform an initial fetch on startup for the first worker that starts.
     await load_or_refresh_models()
     yield
     await app.state.http_client.aclose()
@@ -132,7 +126,6 @@ app = FastAPI(docs_url=None, redoc_url=None, lifespan=lifespan)
 
 @app.get("/")
 async def health_check():
-    """Provides a basic health check endpoint."""
     return JSONResponse({
         "status": "ok",
         "active_models_in_memory": len(worker_model_routing_table)
@@ -146,7 +139,7 @@ async def list_models():
     """
     model_list = [
         { "id": model_id, "object": "model", "owned_by": "gmi-serving" }
-        for model_id in worker_model_routing_table.keys()
+        for model_id in sorted(worker_model_routing_table.keys()) # Sort for consistency
     ]
     return JSONResponse(content={"object": "list", "data": model_list})
 
@@ -177,7 +170,6 @@ async def chat_completions_proxy(request: Request):
     client: httpx.AsyncClient = request.app.state.http_client
     target_url = f"https://{target_host}{request.url.path}"
     
-    # --- Prepare and Forward Request (logic is the same as before) ---
     request_headers = {k: v for k, v in request.headers.items() if k.lower() != 'host'}
     random_ip = ".".join(str(random.randint(1, 254)) for _ in range(4))
     request_headers.update({
@@ -207,5 +199,4 @@ async def chat_completions_proxy(request: Request):
             if attempt == MAX_RETRIES - 1:
                 raise HTTPException(status_code=502, detail=f"Bad Gateway: Error connecting to model backend. {e}")
     
-    # This part should ideally not be reached, but as a fallback:
     raise HTTPException(status_code=502, detail="Bad Gateway: Request failed after all retries.")