fix(trainer): Fix empty image for Runtime trainer

kubeflow/trainer/backends/container/backend.py

@@ -209,8 +209,7 @@ def get_runtime_packages(self, runtime: types.Runtime):
         """
         Spawn a short-lived container to report Python version, pip list, and nvidia-smi.
         """
-        image = container_utils.resolve_image(runtime)
-        container_utils.maybe_pull_image(self._adapter, image, self.cfg.pull_policy)
+        container_utils.maybe_pull_image(self._adapter, runtime.trainer.image, self.cfg.pull_policy)
 
         command = [
             "bash",
@@ -220,14 +219,17 @@ def get_runtime_packages(self, runtime: types.Runtime):
             "(nvidia-smi || echo 'nvidia-smi not found')",
         ]
 
-        logs = self._adapter.run_oneoff_container(image=image, command=command)
+        logs = self._adapter.run_oneoff_container(image=runtime.trainer.image, command=command)
         print(logs)
 
     def train(
         self,
         runtime: Optional[types.Runtime] = None,
         initializer: Optional[types.Initializer] = None,
-        trainer: Optional[Union[types.CustomTrainer, types.BuiltinTrainer]] = None,
+        trainer: Optional[
+            Union[types.CustomTrainer, types.CustomTrainerContainer, types.BuiltinTrainer]
+        ] = None,
+        options: Optional[list] = None,
     ) -> str:
         if runtime is None:
             runtime = self.get_runtime("torch-distributed")
@@ -249,11 +251,12 @@ def train(
             logger.debug("Generated training script code")
 
             # Resolve image and pull if needed
-            image = container_utils.resolve_image(runtime)
-            logger.debug(f"Using image: {image}")
+            logger.debug(f"Using image: {runtime.trainer.image}")
 
-            container_utils.maybe_pull_image(self._adapter, image, self.cfg.pull_policy)
-            logger.debug(f"Image ready: {image}")
+            container_utils.maybe_pull_image(
+                self._adapter, runtime.trainer.image, self.cfg.pull_policy
+            )
+            logger.debug(f"Image ready: {runtime.trainer.image}")
 
             # Build base environment
             env = container_utils.build_environment(trainer)
@@ -368,7 +371,7 @@ def train(
                 logger.debug(f"Creating container {rank}/{num_nodes}: {container_name}")
 
                 container_id = self._adapter.create_and_start_container(
-                    image=image,
+                    image=runtime.trainer.image,
                     command=full_cmd,
                     name=container_name,
                     network_id=network_id,

kubeflow/trainer/backends/container/runtime_loader.py

@@ -327,9 +327,8 @@ def _create_default_runtimes() -> list[base_types.Runtime]:
                 trainer_type=base_types.TrainerType.CUSTOM_TRAINER,
                 framework=framework,
                 num_nodes=1,
+                image=image,
             ),
-            pretrained_model=None,
-            image=image,
         )
         default_runtimes.append(runtime)
         logger.debug(f"Created default runtime: {runtime.name} with image {image}")
@@ -414,9 +413,8 @@ def _parse_runtime_yaml(data: dict[str, Any], source: str = "unknown") -> base_t
             trainer_type=base_types.TrainerType.CUSTOM_TRAINER,
             framework=framework,
             num_nodes=num_nodes,
+            image=image,
         ),
-        pretrained_model=None,
-        image=image,
     )
 
 

kubeflow/trainer/backends/container/runtime_loader_test.py

@@ -283,6 +283,7 @@ def test_list_training_runtimes_from_sources(test_case):
                         trainer_type=base_types.TrainerType.CUSTOM_TRAINER,
                         framework="torch",
                         num_nodes=1,
+                        image="example.com/container",
                     ),
                 )
                 deepspeed_runtime = base_types.Runtime(
@@ -291,6 +292,7 @@ def test_list_training_runtimes_from_sources(test_case):
                         trainer_type=base_types.TrainerType.CUSTOM_TRAINER,
                         framework="deepspeed",
                         num_nodes=1,
+                        image="example.com/container",
                     ),
                 )
                 mock_github.side_effect = [[torch_runtime], [deepspeed_runtime]]
@@ -303,6 +305,7 @@ def test_list_training_runtimes_from_sources(test_case):
                         trainer_type=base_types.TrainerType.CUSTOM_TRAINER,
                         framework="torch",
                         num_nodes=1,
+                        image="example.com/container",
                     ),
                 )
                 torch_runtime_2 = base_types.Runtime(
@@ -311,6 +314,7 @@ def test_list_training_runtimes_from_sources(test_case):
                         trainer_type=base_types.TrainerType.CUSTOM_TRAINER,
                         framework="torch",
                         num_nodes=2,
+                        image="example.com/container",
                     ),
                 )
                 mock_github.side_effect = [[torch_runtime_1], [torch_runtime_2]]
@@ -324,6 +328,7 @@ def test_list_training_runtimes_from_sources(test_case):
                         trainer_type=base_types.TrainerType.CUSTOM_TRAINER,
                         framework="torch",
                         num_nodes=1,
+                        image="example.com/container",
                     ),
                 )
                 mock_defaults.return_value = [default_runtime]
@@ -358,7 +363,7 @@ def test_create_default_runtimes():
     assert torch_runtimes[0].trainer.trainer_type == base_types.TrainerType.CUSTOM_TRAINER
     assert torch_runtimes[0].trainer.num_nodes == 1
     # Verify default image is set
-    assert torch_runtimes[0].image == constants.DEFAULT_FRAMEWORK_IMAGES["torch"]
+    assert torch_runtimes[0].trainer.image == constants.DEFAULT_FRAMEWORK_IMAGES["torch"]
     print("test execution complete")
 
 
@@ -620,72 +625,10 @@ def test_parse_runtime_yaml_extracts_image(test_case):
         runtime = runtime_loader._parse_runtime_yaml(runtime_yaml, "test")
 
         # Verify image is extracted and stored
-        assert runtime.image == test_case.config["custom_image"]
         assert runtime.name == test_case.config["runtime_name"]
         assert runtime.trainer.framework == test_case.config["framework"]
         assert runtime.trainer.num_nodes == test_case.config["num_nodes"]
-
-        assert test_case.expected_status == SUCCESS
-
-    except Exception as e:
-        assert type(e) is test_case.expected_error
-    print("test execution complete")
-
-
-@pytest.mark.parametrize(
-    "test_case",
-    [
-        TestCase(
-            name="resolve image uses custom image",
-            expected_status=SUCCESS,
-            config={
-                "custom_image": "my-registry.io/pytorch-custom:arm64",
-                "framework": "torch",
-                "expect_custom": True,
-            },
-        ),
-        TestCase(
-            name="resolve image falls back to default when no custom image",
-            expected_status=SUCCESS,
-            config={
-                "custom_image": None,
-                "framework": "torch",
-                "expect_custom": False,
-            },
-        ),
-    ],
-)
-def test_resolve_image_uses_custom_image(test_case):
-    """
-    Test that resolve_image prioritizes runtime.image over default framework images.
-    This ensures custom images from ClusterTrainingRuntimes are actually used.
-    """
-    print("Executing test:", test_case.name)
-    try:
-        from kubeflow.trainer.backends.container import utils
-
-        # Create runtime with or without custom image
-        runtime = base_types.Runtime(
-            name="test-runtime",
-            trainer=base_types.RuntimeTrainer(
-                trainer_type=base_types.TrainerType.CUSTOM_TRAINER,
-                framework=test_case.config["framework"],
-                num_nodes=1,
-            ),
-            image=test_case.config["custom_image"],
-        )
-
-        resolved_image = utils.resolve_image(runtime)
-
-        if test_case.config["expect_custom"]:
-            # Should use custom image
-            assert resolved_image == test_case.config["custom_image"]
-        else:
-            # Should fall back to default
-            assert (
-                resolved_image == constants.DEFAULT_FRAMEWORK_IMAGES[test_case.config["framework"]]
-            )
-            assert "pytorch/pytorch" in resolved_image
+        assert runtime.trainer.image == test_case.config["custom_image"]
 
         assert test_case.expected_status == SUCCESS
 

kubeflow/trainer/backends/container/utils.py

@@ -126,7 +126,7 @@ def container_status_to_trainjob_status(status: str, exit_code: int) -> str:
     if status == "exited":
         # Exit code 0 -> complete, else failed
         return constants.TRAINJOB_COMPLETE if exit_code == 0 else constants.TRAINJOB_FAILED
-    return constants.UNKNOWN
+    return UNKNOWN
 
 
 def aggregate_status_from_containers(container_statuses: list[str]) -> str:
@@ -150,38 +150,6 @@ def aggregate_status_from_containers(container_statuses: list[str]) -> str:
     return UNKNOWN
 
 
-def resolve_image(runtime: types.Runtime) -> str:
-    """
-    Resolve the container image for a runtime.
-
-    Priority:
-    1. Use runtime.image if specified in the ClusterTrainingRuntime
-    2. Fall back to DEFAULT_FRAMEWORK_IMAGES based on framework
-
-    Args:
-        runtime: Runtime object.
-
-    Returns:
-        Container image name.
-
-    Raises:
-        ValueError: If no image is found for the runtime's framework.
-    """
-    # Use image from runtime if specified
-    if runtime.image:
-        return runtime.image
-
-    # Fall back to default framework images
-    framework = runtime.trainer.framework
-    if framework in constants.DEFAULT_FRAMEWORK_IMAGES:
-        return constants.DEFAULT_FRAMEWORK_IMAGES[framework]
-
-    raise ValueError(
-        f"No default image found for framework '{framework}'. "
-        f"Supported frameworks: {list(constants.DEFAULT_FRAMEWORK_IMAGES.keys())}"
-    )
-
-
 def maybe_pull_image(adapter, image: str, pull_policy: str):
     """
     Pull image based on pull policy.
@@ -227,7 +195,7 @@ def get_container_status(adapter, container_id: str) -> str:
         status, exit_code = adapter.container_status(container_id)
         return container_status_to_trainjob_status(status, exit_code)
     except Exception:
-        return constants.UNKNOWN
+        return UNKNOWN
 
 
 def aggregate_container_statuses(adapter, containers: list[dict]) -> str:

kubeflow/trainer/backends/kubernetes/backend_test.py

@@ -527,7 +527,7 @@ def get_replicated_job() -> models.JobsetV1alpha2ReplicatedJob:
 def get_container() -> models.IoK8sApiCoreV1Container:
     return models.IoK8sApiCoreV1Container(
         name="node",
-        image="image",
+        image="example.com/test-runtime",
         command=["echo", "Hello World"],
         resources=get_resource_requirements(),
     )
@@ -543,11 +543,11 @@ def create_runtime_type(
         num_nodes=2,
         device="gpu",
         device_count=RUNTIME_DEVICES,
+        image="example.com/test-runtime",
     )
     trainer.set_command(constants.TORCH_COMMAND)
     return types.Runtime(
         name=name,
-        pretrained_model=None,
         trainer=trainer,
     )
 
@@ -564,14 +564,14 @@ def get_train_job_data_type(
         device="gpu",
         device_count=RUNTIME_DEVICES,
         num_nodes=2,
+        image="example.com/test-runtime",
     )
     trainer.set_command(constants.TORCH_COMMAND)
     return types.TrainJob(
         name=train_job_name,
         creation_timestamp=datetime.datetime(2025, 6, 1, 10, 30, 0),
         runtime=types.Runtime(
             name=runtime_name,
-            pretrained_model=None,
             trainer=trainer,
         ),
         steps=[
@@ -696,6 +696,7 @@ def test_list_runtimes(kubernetes_backend, test_case):
                         num_nodes=1,
                         device="cpu",
                         device_count="1",
+                        image="example.com/image",
                     ),
                 )
             },

kubeflow/trainer/backends/kubernetes/utils.py

@@ -103,6 +103,7 @@ def get_runtime_trainer(
             else types.TrainerType.CUSTOM_TRAINER
         ),
         framework=framework,
+        image=trainer_container.image,
     )
 
     # Get the container devices.

kubeflow/trainer/backends/kubernetes/utils_test.py

@@ -26,6 +26,7 @@ def _build_runtime() -> types.Runtime:
         framework="torch",
         device="cpu",
         device_count="1",
+        image="example.com/image",
     )
     runtime_trainer.set_command(constants.DEFAULT_COMMAND)
     return types.Runtime(name="test-runtime", trainer=runtime_trainer)

kubeflow/trainer/backends/localprocess/backend.py

@@ -282,6 +282,6 @@ def __convert_local_runtime_to_runtime(self, local_runtime) -> types.Runtime:
                 num_nodes=local_runtime.trainer.num_nodes,
                 device_count=local_runtime.trainer.device_count,
                 device=local_runtime.trainer.device,
+                image=local_runtime.trainer.image,
             ),
-            pretrained_model=local_runtime.pretrained_model,
         )

kubeflow/trainer/backends/localprocess/constants.py

@@ -22,6 +22,9 @@
 
 TORCH_FRAMEWORK_TYPE = "torch"
 
+# Image name for the local runtime.
+LOCAL_RUNTIME_IMAGE = "local"
+
 local_runtimes = [
     base_types.Runtime(
         name=constants.TORCH_RUNTIME,
@@ -32,6 +35,7 @@
             device_count=common_constants.UNKNOWN,
             device=common_constants.UNKNOWN,
             packages=["torch"],
+            image=LOCAL_RUNTIME_IMAGE,
         ),
     )
 ]

kubeflow/trainer/backends/localprocess/utils.py

@@ -132,6 +132,7 @@ def get_local_runtime_trainer(
         trainer_type=types.TrainerType.CUSTOM_TRAINER,
         framework=framework,
         packages=local_runtime.trainer.packages,
+        image=local_exec_constants.LOCAL_RUNTIME_IMAGE,
     )
 
     # set command to run from venv

kubeflow/trainer/types/types.py

@@ -232,6 +232,7 @@ class TrainerType(Enum):
 class RuntimeTrainer:
     trainer_type: TrainerType
     framework: str
+    image: str
     num_nodes: int = 1  # The default value is set in the APIs.
     device: str = common_constants.UNKNOWN
     device_count: str = common_constants.UNKNOWN
@@ -251,7 +252,6 @@ class Runtime:
     name: str
     trainer: RuntimeTrainer
     pretrained_model: Optional[str] = None
-    image: Optional[str] = None
 
 
 # Representation for the TrainJob steps.