Fix resolution of models image in the agent

GitOrigin-RevId: 859d532533bf5a13227210d615b81b420581e80e

src/gretel_client/agents/drivers/k8s.py

@@ -599,7 +599,9 @@ def _resolve_image(self, job_config: Job) -> str:
             # We'll instead release the images with gretelai/image-separated-by-dashes
             if registry_host and self._worker.uses_dockerhub_format:
                 start_index = 0
-                if "." in image_parts[0] or ":" in image_parts[0]:
+                if len(image_parts) != 1 and any(
+                    ["." in image_parts[0], ":" in image_parts[0]]
+                ):
                     start_index = 1
                 image_name_with_dashes = "-".join(image_parts[start_index:])
                 image = f"{registry_host}/{image_name_with_dashes}"

tests/gretel_client/test_k8s.py

@@ -1110,6 +1110,18 @@ def test_special_case_gretelai(self):
 
         assert image == "shiny-new-reg.example.ai/gretelai/gcc-my-workflow-image:latest"
 
+    @patch_auth
+    @patch_image_registry("shiny-new-reg.example.ai/gretelai")
+    def test_special_case_gretelai_models(self):
+        job = Job.from_dict(get_mock_job("gpu-standard"), self.config)
+        original_image = "models:latest"
+        job.container_image = original_image
+        daemon = KubernetesDriverDaemon(self.config, self.core_api)
+        self.driver._worker = daemon
+        image = self.driver._resolve_image(job)
+
+        assert image == "shiny-new-reg.example.ai/gretelai/models:latest"
+
     @patch_autoscaler_env_var("true")
     def test_annotation_set_true(self):
         job = Job.from_dict(get_mock_job(), self.config)