removing timezone tzinfo in favor of utc (#257)

* removing timezone tzinfo in favor of utc

* adding utc everywhere in last_updated

* adding variable assignment

* fixing timezone assignment to utc

* fixing logging

Author: ian-scale

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

@@ -1,6 +1,7 @@
 from datetime import datetime
 from typing import Dict, NamedTuple, Tuple
 
+import pytz
 from model_engine_server.common.config import hmi_config
 from model_engine_server.common.env_vars import GIT_TAG
 from model_engine_server.core.config import infra_config
@@ -108,7 +109,11 @@ async def execute(self, endpoint_infra_states: Dict[str, Tuple[bool, ModelEndpoi
             if record is None:
                 continue
 
-            last_updated_at = record.last_updated_at or datetime.min
+            last_updated_at = (
+                record.last_updated_at.replace(tzinfo=pytz.utc)
+                if record.last_updated_at is not None
+                else datetime.min.replace(tzinfo=pytz.utc)
+            )
             has_no_available_workers = int(state.deployment_state.available_workers == 0)
             is_high_priority = int(state.high_priority is True)
 
@@ -125,36 +130,36 @@ async def execute(self, endpoint_infra_states: Dict[str, Tuple[bool, ModelEndpoi
 
             image_repository_and_tag = state.image.split("/", 1)[1]
             repository_name, image_tag = image_repository_and_tag.split(":")
-            if state.resource_state.gpus == 0 and (
-                (
-                    state.image not in images_to_cache_priority["cpu"]
-                    or last_updated_at.replace(
-                        tzinfo=images_to_cache_priority["cpu"][state.image].last_updated_at.tzinfo
+            try:
+                if state.resource_state.gpus == 0 and (
+                    (
+                        state.image not in images_to_cache_priority["cpu"]
+                        or last_updated_at.replace(tzinfo=pytz.utc)
+                        > images_to_cache_priority["cpu"][state.image].last_updated_at
                     )
-                    > images_to_cache_priority["cpu"][state.image].last_updated_at
-                )
-                and self.docker_repository.image_exists(image_tag, repository_name)
-            ):
-                images_to_cache_priority["cpu"][state.image] = cache_priority
-            elif state.resource_state.gpus > 0:
-                for gpu_type, key in [
-                    (GpuType.NVIDIA_AMPERE_A10, "a10"),
-                    (GpuType.NVIDIA_AMPERE_A100, "a100"),
-                    (GpuType.NVIDIA_TESLA_T4, "t4"),
-                ]:
-                    if state.resource_state.gpu_type == gpu_type and (
-                        (
-                            state.image not in images_to_cache_priority[key]
-                            or last_updated_at.replace(
-                                tzinfo=images_to_cache_priority[key][
-                                    state.image
-                                ].last_updated_at.tzinfo
+                    and self.docker_repository.image_exists(image_tag, repository_name)
+                ):
+                    images_to_cache_priority["cpu"][state.image] = cache_priority
+                elif state.resource_state.gpus > 0:
+                    for gpu_type, key in [
+                        (GpuType.NVIDIA_AMPERE_A10, "a10"),
+                        (GpuType.NVIDIA_AMPERE_A100, "a100"),
+                        (GpuType.NVIDIA_TESLA_T4, "t4"),
+                    ]:
+                        if state.resource_state.gpu_type == gpu_type and (
+                            (
+                                state.image not in images_to_cache_priority[key]
+                                or last_updated_at.replace(tzinfo=pytz.utc)
+                                > images_to_cache_priority[key][state.image].last_updated_at
                             )
-                            > images_to_cache_priority[key][state.image].last_updated_at
-                        )
-                        and self.docker_repository.image_exists(image_tag, repository_name)
-                    ):
-                        images_to_cache_priority[key][state.image] = cache_priority
+                            and self.docker_repository.image_exists(image_tag, repository_name)
+                        ):
+                            images_to_cache_priority[key][state.image] = cache_priority
+            except Exception as exc:
+                logger.warning(
+                    f"Endpoint {endpoint_id} had an error. Error message: {exc}. Skipping caching ..."
+                )
+                continue
 
         images_to_cache = CachedImages(cpu=[], a10=[], a100=[], t4=[])
         for key, val in images_to_cache_priority.items():