feat(trainer): add dataset and model initializer support to container backend (#188)

* feat(trainer): add dataset and model initializer support to container backend

Add support for dataset and model initializers in the container backend
to bring it to feature parity with the Kubernetes backend.

Changes:
- Add utility functions for building initializer commands and environment variables
- Implement _run_initializers() and _run_single_initializer() methods in ContainerBackend
- Run initializers sequentially before training containers start
- Download datasets to /workspace/dataset and models to /workspace/model
- Track initializer containers as separate steps in TrainJob
- Support all initializer types: HuggingFace, S3, and DataCache
- Add comprehensive unit tests for all initializer configurations
- Handle initializer failures with proper cleanup and error messages

Fixes #171

Signed-off-by: HKanoje <hrithik.kanoje@gmail.com>

* feat(trainer): address reviewer feedback for initializer support

- Make initializer image configurable via ContainerBackendConfig
- Make initializer timeout configurable (default 600 seconds)
- Implement wait API in adapters instead of polling
- Clean up successful initializer containers after completion
- Clean up network on initializer failure
- Raise ValueError for unsupported initializer types (no datacache fallback)

All tests passing (173/173). Addresses all feedback from PR #188.

Signed-off-by: HKanoje <hrithik.kanoje@gmail.com>

* chore(trainer): add cleanup helper to reduce duplication

Add _cleanup_container_resources() helper method to consolidate
duplicated cleanup logic for stopping/removing containers and
deleting networks. Refactor 5 locations across train(), initializer
handlers, and delete_job() to use this helper.

Signed-off-by: HKanoje <hrithik.kanoje@gmail.com>

* fix(trainer): use correct initializer images and working directory

Address feedback for initializer support in container backend:

- Use separate images for dataset/model initializers:
  - kubeflow/dataset-initializer:latest for datasets
  - kubeflow/model-initializer:latest for models
  (instead of kubeflow/training-operator:latest)

- Update python commands to use pkg.initializers module:
  - python -m pkg.initializers.dataset (for dataset)
  - python -m pkg.initializers.model (for model)

- Change initializer working_dir from /workspace to /app
  per Dockerfile convention

Refs: https://github.com/kubeflow/trainer/tree/master/cmd/initializers
Signed-off-by: HKanoje <hrithik.kanoje@gmail.com>

* fix(container): address PR review comments for initializer support

- Use GHCR images as default for dataset/model initializers
- Replace suppress with try-except blocks
- Refactor initializer utils with ContainerInitializer dataclass
- Add get_dataset_initializer and get_model_initializer functions
- Remove DataCache support (unsupported in container backend)
- Merge initializer tests into test_train() and test_get_job_logs()
- Remove duplicate test functions

Signed-off-by: HKanoje <hrithik.kanoje@gmail.com>

* fix(container): add name field to ContainerInitializer and remove init_type

- Add name field to ContainerInitializer dataclass
- Set name='dataset-initializer' and name='model-initializer' in utils
- Remove init_type parameter from _run_single_initializer()
- Use container_init.name for labels and log messages

Signed-off-by: HKanoje <hrithik.kanoje@gmail.com>

---------

Signed-off-by: HKanoje <hrithik.kanoje@gmail.com>

Author: HKanoje

kubeflow/trainer/backends/container/adapters/base.py

@@ -193,3 +193,22 @@ def get_network(self, network_id: str) -> Optional[dict]:
             Dictionary with network info including labels, or None if not found
         """
         raise NotImplementedError()
+
+    @abc.abstractmethod
+    def wait_for_container(self, container_id: str, timeout: Optional[int] = None) -> int:
+        """
+        Wait for a container to exit and return its exit code.
+
+        This is a blocking call that waits until the container stops.
+
+        Args:
+            container_id: Container ID
+            timeout: Maximum time to wait in seconds, or None to wait indefinitely
+
+        Returns:
+            Container exit code
+
+        Raises:
+            TimeoutError: If timeout is reached before container exits
+        """
+        raise NotImplementedError()

kubeflow/trainer/backends/container/adapters/docker.py

@@ -227,3 +227,31 @@ def get_network(self, network_id: str) -> Optional[dict]:
             }
         except Exception:
             return None
+
+    def wait_for_container(self, container_id: str, timeout: Optional[int] = None) -> int:
+        """
+        Wait for a Docker container to exit and return its exit code.
+
+        Args:
+            container_id: Container ID
+            timeout: Maximum time to wait in seconds, or None to wait indefinitely
+
+        Returns:
+            Container exit code
+
+        Raises:
+            TimeoutError: If timeout is reached before container exits
+        """
+        try:
+            container = self.get_container(container_id)
+            result = container.wait(timeout=timeout)
+            # Docker wait() returns a dict with 'StatusCode' key
+            if isinstance(result, dict):
+                return result.get("StatusCode", 0)
+            return int(result)
+        except Exception as e:
+            if "timeout" in str(e).lower():
+                raise TimeoutError(
+                    f"Container {container_id} did not exit within {timeout} seconds"
+                ) from e
+            raise

kubeflow/trainer/backends/container/adapters/podman.py

@@ -254,3 +254,29 @@ def get_network(self, network_id: str) -> Optional[dict]:
             }
         except Exception:
             return None
+
+    def wait_for_container(self, container_id: str, timeout: Optional[int] = None) -> int:
+        """
+        Wait for a Podman container to exit and return its exit code.
+
+        Args:
+            container_id: Container ID
+            timeout: Maximum time to wait in seconds, or None to wait indefinitely
+
+        Returns:
+            Container exit code
+
+        Raises:
+            TimeoutError: If timeout is reached before container exits
+        """
+        try:
+            container = self.get_container(container_id)
+            result = container.wait(timeout=timeout)
+            # Podman wait() returns exit code directly
+            return int(result)
+        except Exception as e:
+            if "timeout" in str(e).lower():
+                raise TimeoutError(
+                    f"Container {container_id} did not exit within {timeout} seconds"
+                ) from e
+            raise