Azure fixes + additional asks (#468)

Author: squeakymouse

charts/model-engine/templates/_helpers.tpl

@@ -270,8 +270,6 @@ env:
     value: {{ .Values.azure.abs_account_name }}
   - name: SERVICEBUS_NAMESPACE
     value: {{ .Values.azure.servicebus_namespace }}
-  - name: SERVICEBUS_SAS_KEY
-    value: {{ .Values.azure.servicebus_sas_key }}
   {{- end }}
 {{- end }}
 

model-engine/Dockerfile

@@ -1,6 +1,6 @@
 # syntax = docker/dockerfile:experimental
 
-FROM python:3.8.8-slim as model-engine
+FROM python:3.8.18-slim as model-engine
 WORKDIR /workspace
 
 RUN apt-get update && apt-get install -y \

model-engine/model_engine_server/api/dependencies.py

@@ -1,5 +1,6 @@
 import asyncio
 import os
+import time
 from dataclasses import dataclass
 from typing import Callable, Optional
 
@@ -442,6 +443,7 @@ async def verify_authentication(
 def get_or_create_aioredis_pool() -> aioredis.ConnectionPool:
     global _pool
 
-    if _pool is None:
+    expiration_timestamp = hmi_config.cache_redis_url_expiration_timestamp
+    if _pool is None or (expiration_timestamp is not None and time.time() > expiration_timestamp):
         _pool = aioredis.BlockingConnectionPool.from_url(hmi_config.cache_redis_url)
     return _pool

model-engine/model_engine_server/api/files_v1.py

@@ -44,7 +44,7 @@ async def upload_file(
     )
     return await use_case.execute(
         user=auth,
-        filename=file.filename,
+        filename=file.filename or "",
         content=file.file.read(),
     )
 

model-engine/model_engine_server/common/config.py

@@ -30,6 +30,8 @@
 
 SERVICE_CONFIG_PATH = os.environ.get("DEPLOY_SERVICE_CONFIG_PATH", DEFAULT_SERVICE_CONFIG_PATH)
 
+redis_cache_expiration_timestamp = None
+
 
 # duplicated from llm/ia3_finetune
 def get_model_cache_directory_name(model_name: str):
@@ -81,9 +83,17 @@ def cache_redis_url(self) -> str:
 
         assert self.cache_redis_azure_host and infra_config().cloud_provider == "azure"
         username = os.getenv("AZURE_OBJECT_ID")
-        password = DefaultAzureCredential().get_token("https://redis.azure.com/.default").token
+        token = DefaultAzureCredential().get_token("https://redis.azure.com/.default")
+        password = token.token
+        global redis_cache_expiration_timestamp
+        redis_cache_expiration_timestamp = token.expires_on
         return f"rediss://{username}:{password}@{self.cache_redis_azure_host}"
 
+    @property
+    def cache_redis_url_expiration_timestamp(self) -> Optional[int]:
+        global redis_cache_expiration_timestamp
+        return redis_cache_expiration_timestamp
+
     @property
     def cache_redis_host_port(self) -> str:
         # redis://redis.url:6379/<db_index>

model-engine/model_engine_server/db/base.py

@@ -58,7 +58,7 @@ def get_engine_url(env: Optional[str] = None, read_only: bool = True, sync: bool
             user = os.environ.get("AZURE_IDENTITY_NAME")
             password = (
                 DefaultAzureCredential()
-                .get_token("https://ossrdbms-aad.database.windows.net")
+                .get_token("https://ossrdbms-aad.database.windows.net/.default")
                 .token
             )
             logger.info(f"Connecting to db {db} as user {user}")
@@ -81,7 +81,9 @@ def get_engine_url(env: Optional[str] = None, read_only: bool = True, sync: bool
 
     # For async postgres, we need to use an async dialect.
     if not sync:
-        engine_url = engine_url.replace("postgresql://", "postgresql+asyncpg://")
+        engine_url = engine_url.replace("postgresql://", "postgresql+asyncpg://").replace(
+            "sslmode", "ssl"
+        )
     return engine_url
 
 

model-engine/model_engine_server/domain/exceptions.py

@@ -41,6 +41,12 @@ class DockerImageNotFoundException(DomainException):
     tag: str
 
 
+class DockerRepositoryNotFoundException(DomainException):
+    """
+    Thrown when a Docker repository that is trying to be accessed doesn't exist.
+    """
+
+
 class DockerBuildFailedException(DomainException):
     """
     Thrown if the server failed to build a docker image.

model-engine/model_engine_server/infra/gateways/live_model_endpoints_schema_gateway.py

@@ -3,8 +3,9 @@
 from typing import Any, Callable, Dict, Sequence, Set, Type, Union
 
 from fastapi import routing
+from fastapi._compat import GenerateJsonSchema, get_model_definitions
+from fastapi.openapi.constants import REF_TEMPLATE
 from fastapi.openapi.utils import get_openapi_path
-from fastapi.utils import get_model_definitions
 from model_engine_server.common.dtos.tasks import (
     EndpointPredictV1Request,
     GetAsyncTaskV1Response,
@@ -119,8 +120,13 @@ def get_openapi(
             if isinstance(route, routing.APIRoute):
                 prefix = model_endpoint_name
                 model_name_map = LiveModelEndpointsSchemaGateway.get_model_name_map(prefix)
+                schema_generator = GenerateJsonSchema(ref_template=REF_TEMPLATE)
                 result = get_openapi_path(
-                    route=route, model_name_map=model_name_map, operation_ids=operation_ids
+                    route=route,
+                    model_name_map=model_name_map,
+                    operation_ids=operation_ids,
+                    schema_generator=schema_generator,
+                    field_mapping={},
                 )
                 if result:
                     path, security_schemes, path_definitions = result

model-engine/model_engine_server/infra/repositories/acr_docker_repository.py

@@ -6,6 +6,7 @@
 from model_engine_server.common.dtos.docker_repository import BuildImageRequest, BuildImageResponse
 from model_engine_server.core.config import infra_config
 from model_engine_server.core.loggers import logger_name, make_logger
+from model_engine_server.domain.exceptions import DockerRepositoryNotFoundException
 from model_engine_server.domain.repositories import DockerRepository
 
 logger = make_logger(logger_name())
@@ -36,7 +37,11 @@ def get_latest_image_tag(self, repository_name: str) -> str:
         credential = DefaultAzureCredential()
         client = ContainerRegistryClient(endpoint, credential)
 
-        image = client.list_manifest_properties(
-            repository_name, order_by="time_desc", results_per_page=1
-        ).next()
-        return image.tags[0]
+        try:
+            image = client.list_manifest_properties(
+                repository_name, order_by="time_desc", results_per_page=1
+            ).next()
+            # Azure automatically deletes empty ACR repositories, so repos will always have at least one image
+            return image.tags[0]
+        except ResourceNotFoundError:
+            raise DockerRepositoryNotFoundException

model-engine/model_engine_server/infra/services/image_cache_service.py

@@ -7,6 +7,7 @@
 from model_engine_server.core.config import infra_config
 from model_engine_server.core.loggers import logger_name, make_logger
 from model_engine_server.domain.entities import GpuType, ModelEndpointInfraState
+from model_engine_server.domain.exceptions import DockerRepositoryNotFoundException
 from model_engine_server.domain.repositories import DockerRepository
 from model_engine_server.infra.gateways.resources.image_cache_gateway import (
     CachedImages,
@@ -78,11 +79,14 @@ def _cache_finetune_llm_images(
         vllm_image_032 = DockerImage(
             f"{infra_config().docker_repo_prefix}/{hmi_config.vllm_repository}", "0.3.2"
         )
-        latest_tag = (
-            self.docker_repository.get_latest_image_tag(hmi_config.batch_inference_vllm_repository)
-            if not CIRCLECI
-            else "fake_docker_repository_latest_image_tag"
-        )
+        latest_tag = "fake_docker_repository_latest_image_tag"
+        if not CIRCLECI:
+            try:  # pragma: no cover
+                latest_tag = self.docker_repository.get_latest_image_tag(
+                    hmi_config.batch_inference_vllm_repository
+                )
+            except DockerRepositoryNotFoundException:
+                pass
         vllm_batch_image_latest = DockerImage(
             f"{infra_config().docker_repo_prefix}/{hmi_config.batch_inference_vllm_repository}",
             latest_tag,