Merge pull request #149 from NillionNetwork/feat/env_model_settings

feat: added model settings

Author: jcabrero

nilai-models/src/nilai_models/daemon.py

@@ -6,6 +6,7 @@
 
 from nilai_common import (  # Model service discovery and host settings
     SETTINGS,
+    MODEL_SETTINGS,
     ModelServiceDiscovery,
     ModelEndpoint,
     ModelMetadata,
@@ -14,7 +15,7 @@
 logger = logging.getLogger(__name__)
 
 
-async def get_metadata(num_retries=30):
+async def get_metadata():
     """Fetch model metadata from model
     service and return as ModelMetadata object"""
     current_retries = 0
@@ -46,9 +47,13 @@ async def get_metadata(num_retries=30):
             else:
                 logger.warning(f"Failed to fetch model metadata from {url}: {e}")
             current_retries += 1
-            if current_retries >= num_retries:
+            if (
+                MODEL_SETTINGS.num_retries
+                != -1  # If num_retries == -1 then we do infinite number of retries
+                and current_retries >= MODEL_SETTINGS.num_retries
+            ):
                 raise e
-            await asyncio.sleep(10)
+            await asyncio.sleep(MODEL_SETTINGS.timeout)
 
 
 async def run_service(discovery_service, model_endpoint):

packages/nilai-common/src/nilai_common/__init__.py

@@ -16,7 +16,7 @@
     Message,
     MessageAdapter,
 )
-from nilai_common.config import SETTINGS
+from nilai_common.config import SETTINGS, MODEL_SETTINGS
 from nilai_common.discovery import ModelServiceDiscovery
 from openai.types.completion_usage import CompletionUsage as Usage
 
@@ -36,6 +36,7 @@
     "AMDAttestationToken",
     "NVAttestationToken",
     "SETTINGS",
+    "MODEL_SETTINGS",
     "SearchResult",
     "Source",
     "WebSearchEnhancedMessages",

packages/nilai-common/src/nilai_common/config.py

@@ -1,5 +1,5 @@
 import os
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 
 class HostSettings(BaseModel):
@@ -13,6 +13,11 @@ class HostSettings(BaseModel):
     attestation_port: int = 8081
 
 
+class ModelSettings(BaseModel):
+    num_retries: int = Field(default=30, ge=-1)
+    timeout: int = Field(default=10, ge=1)
+
+
 SETTINGS: HostSettings = HostSettings(
     host=str(os.getenv("SVC_HOST", "localhost")),
     port=int(os.getenv("SVC_PORT", 8000)),
@@ -23,3 +28,8 @@ class HostSettings(BaseModel):
     attestation_host=str(os.getenv("ATTESTATION_HOST", "localhost")),
     attestation_port=int(os.getenv("ATTESTATION_PORT", 8081)),
 )
+
+MODEL_SETTINGS: ModelSettings = ModelSettings(
+    num_retries=int(os.getenv("MODEL_NUM_RETRIES", 30)),
+    timeout=int(os.getenv("MODEL_RETRY_TIMEOUT", 10)),
+)