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

[HKanoje]

What this PR does / why we need it

This PR implements dataset and model initializer support in the container backend, bringing it to feature parity with the Kubernetes backend. This addresses issue #171 by enabling users to automatically download and prepare datasets and models before training starts.

Solution Overview

This implementation adds full initializer support to the container backend by:

1. Running initializer containers before training - Initializers execute sequentially (dataset first, then model) before any training containers start
2. Using shared volumes - All containers (initializers and training nodes) share the same workspace directory on the host
3. Proper lifecycle management - Initializers must complete successfully before training begins, with automatic cleanup on failures
4. Comprehensive error handling - Clear error messages and proper resource cleanup when initialization fails

Detailed Changes

1. New Utility Functions (kubeflow/trainer/backends/container/utils.py)

build_initializer_command(initializer, init_type)

Builds the appropriate container command based on initializer type:

• HuggingFace: Uses kubeflow.storage_initializer.hugging_face module
• S3: Uses kubeflow.storage_initializer.s3 module
• DataCache: Uses kubeflow.storage_initializer.datacache module

build_initializer_env(initializer, init_type)

Constructs environment variables from initializer configuration:

• Sets STORAGE_URI from the initializer config
• Sets OUTPUT_PATH to /workspace/dataset or /workspace/model based on type
• Adds optional fields like ACCESS_TOKEN, ENDPOINT, REGION, etc.
• Handles DataCache-specific variables like CLUSTER_SIZE, METADATA_LOC

get_initializer_image()

Returns the initializer container image (kubeflow/training-operator:latest).
This can be made configurable via backend config in future iterations.

2. Enhanced ContainerBackend (kubeflow/trainer/backends/container/backend.py)

_run_initializers(job_name, initializer, workdir, network_id)

Orchestrates the initialization phase:

• Pulls the initializer image if needed (respects pull_policy)
• Runs dataset initializer if configured
• Runs model initializer if configured
• Ensures proper sequencing and error propagation

_run_single_initializer(job_name, initializer_config, init_type, image, workdir, network_id)

Executes a single initializer container:

• Creates container with proper labels for tracking
• Mounts shared volume to /workspace
• Monitors container status with configurable timeout (default 10 minutes)
• Waits for successful completion (exit code 0)
• Captures and reports logs on failure
• Cleans up failed containers automatically

Updated train() method

• Creates network and working directory first
• Runs initializers before generating training script
• Only proceeds to training if initialization succeeds
• Maintains backward compatibility (initializers are optional)

Updated __get_trainjob_from_containers()

• Correctly counts only training nodes for num_nodes (excludes initializers)
• Includes initializer containers in the steps list
• Properly tracks initializer status

Updated get_job_logs()

• Supports fetching logs from specific initializer steps
• When requesting node-0 logs (default), only shows training container logs
• Can explicitly request initializer logs with step="dataset-initializer" or step="model-initializer"

3. Comprehensive Test Coverage (kubeflow/trainer/backends/container/backend_test.py)

Added 11 new test cases covering:

Initialization Success Scenarios

• HuggingFace dataset initializer - Tests storage_uri parsing and access_token handling
• S3 dataset initializer - Tests endpoint, region, and credential configuration
• HuggingFace model initializer - Tests model downloads with ignore_patterns
• S3 model initializer - Tests S3-compatible storage with authentication
• Both dataset and model - Tests sequential execution of both initializers
• DataCache initializer - Tests distributed cache configuration with metadata_loc and cluster size

Log Retrieval

• Tests getting logs from dataset-initializer step
• Tests getting logs from model-initializer step
• Tests that default node logs exclude initializer logs

Error Handling

• Non-zero exit code - Verifies proper error reporting when initializer fails
• Timeout handling - Ensures timeout errors are caught and reported
• Resource cleanup - Confirms containers and networks are cleaned up on failure

Implementation Details

Initialization Flow

1. User calls train() with initializer parameter

2. ContainerBackend creates:

   ├── Working directory: ~/.kubeflow/trainer/containers/<job-name>/
   └── Network: <job-name>-net

3. If initializer.dataset is set:

   ├── Create dataset-initializer container
   ├── Mount workdir to /workspace
   ├── Run initialization (downloads to /workspace/dataset)
   └── Wait for completion (exit code 0)

4. If initializer.model is set:

   ├── Create model-initializer container
   ├── Mount workdir to /workspace
   ├── Run initialization (downloads to /workspace/model)
   └── Wait for completion (exit code 0)

5. Create training containers:

   ├── Mount same workdir to /workspace
   └── Access data at /workspace/dataset and /workspace/model

Volume Layout

Host: ~/.kubeflow/trainer/containers/<job-name>/

├── dataset/ (from dataset-initializer, if configured)
│   └── <downloaded dataset files>
├── model/ (from model-initializer, if configured)
│   └── <downloaded model files>
└── outputs/ (accessible to all training nodes)
    └── <training outputs>

Container Mount: /workspace/

├── dataset/ (read by training code)
├── model/ (read by training code)
└── outputs/ (written by training code)

Testing Results

All tests pass with no regressions:

• 43/43 container backend tests passed
• 173/173 trainer module tests passed
• make verify (ruff lint + format) passed
• pre-commit hooks passed

Usage Example

from kubeflow.trainer import TrainerClient
from kubeflow.trainer.types import (
    CustomTrainer,
    Initializer,
    HuggingFaceDatasetInitializer,
    HuggingFaceModelInitializer,
)

# Define initializers
initializer = Initializer(
    dataset=HuggingFaceDatasetInitializer(
        storage_uri="hf://username/my-dataset",
        access_token="hf_xxxxx",
    ),
    model=HuggingFaceModelInitializer(
        storage_uri="hf://username/my-model",
        access_token="hf_xxxxx",
    ),
)

# Define trainer
trainer = CustomTrainer(
    func=train_func,
    num_nodes=2,
)

# Train with initializers
client = TrainerClient()
job_name = client.train(
    trainer=trainer,
    initializer=initializer,  # <-- Now works with container backend!
)

# Data is automatically available at:
# - /workspace/dataset (in training containers)
# - /workspace/model (in training containers)

[github-actions]

🎉 Welcome to the Kubeflow SDK! 🎉

Thanks for opening your first PR! We're happy to have you as part of our community 🚀

Here's what happens next:

• If you haven't already, please check out our Contributing Guide for repo-specific guidelines and the Kubeflow Contributor Guide for general community standards
• Our team will review your PR soon! cc @kubeflow/kubeflow-sdk-team

Join the community:

• Slack: Join our #kubeflow-ml-experience and #kubeflow-trainer Slack channels
• Meetings: Attend the Kubeflow SDK and ML Experience bi-weekly meetings

Feel free to ask questions in the comments if you need any help or clarification!
Thanks again for contributing to Kubeflow! 🙏

[Fiona-Waters]

Thanks @HKanoje for this PR! I've left some comments, ptal! Thanks

[Fiona-Waters]

Should we make this configurable rather than hardcoding it?

[HKanoje]

Yes, Definitely. I've made this configurable via ContainerBackendConfig.initializer_image (default: kubeflow/training-operator:latest). Users can now customize it when creating the backend.

[Fiona-Waters]

Should this be configurable, or is 10 minutes always going to be enough time?

[HKanoje]

Added ContainerBackendConfig.initializer_timeout (default: 600 seconds / 10 minutes). This gives users flexibility for large datasets/models that may take longer to download.

[Fiona-Waters]

As well as cleaning up when a failure occurs, should we clean up the initializer containers when they have been successful also?

[HKanoje]

Implemented wait_for_container() in the adapter interface and both Docker/Podman adapters. This replaces the polling loop with a single blocking wait call - much more efficient.

[Fiona-Waters]

If the initializer fails should we clean up the network we have created?

[HKanoje]

Absolutely! Added cleanup for successful initializer containers after completion to prevent accumulation. Also added cleanup for timed-out containers.

[Fiona-Waters]

Great! Would it make sense to add a helper function for the cleanup logic to reduce duplication?

[HKanoje]

Done! Added _cleanup_container_resources() helper method in commit 0b7a952 to consolidate the duplicated cleanup logic across exception handlers and delete_job().

[Fiona-Waters]

Here we are setting datacache as the default/fallback, do we want to do this? In thekubernetes backend we offer 2 options and raise a value error if the type is invalid.

[HKanoje]

You're right - that was inconsistent with the kubernetes backend. Changed to raise ValueError with a clear message listing all supported types instead of defaulting to datacache.

[Fiona-Waters]

Could using wait API be supported?

[kramaranya]

/ok-to-test

[kramaranya]

Hey @HKanoje, could you please sign your commits?

[coveralls]

Pull Request Test Coverage Report for Build 21811391977

Warning: This coverage report may be inaccurate.

This pull request's base commit is no longer the HEAD commit of its target branch. This means it includes changes from outside the original pull request, including, potentially, unrelated coverage changes.

• For more information on this, see Tracking coverage changes with pull request builds.
• To avoid this issue with future PRs, see these Recommended CI Configurations.
• For a quick fix, rebase this PR at GitHub. Your next report should be accurate.

Details

• 160 of 195 (82.05%) changed or added relevant lines in 7 files are covered.
• 21 unchanged lines in 1 file lost coverage.
• Overall coverage increased (+0.9%) to 68.424%

---

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
kubeflow/trainer/backends/container/adapters/base.py	2	3	66.67%
kubeflow/trainer/backends/container/backend_test.py	59	61	96.72%
kubeflow/trainer/backends/container/utils.py	28	30	93.33%
kubeflow/trainer/backends/container/adapters/podman.py	1	9	11.11%
kubeflow/trainer/backends/container/adapters/docker.py	1	11	9.09%
kubeflow/trainer/backends/container/backend.py	66	78	84.62%

Files with Coverage Reduction	New Missed Lines	%
kubeflow/trainer/backends/kubernetes/utils.py	21	77.64%

Totals
Change from base Build 21637340746:	0.9%
Covered Lines:	2921
Relevant Lines:	4269

---

💛 - Coveralls

[HKanoje]

@kramaranya Done! All commits are now signed.

[HKanoje]

@kramaranya @szaher Please Review the changes whenever you get a chance! Thanks!

[Fiona-Waters]

I ran this locally using docker/colima, I have left some comments that should be addressed.
I have created this quick PR in the trainer repo moving a kubernetes import to method level as it was causing an error locally looking for kubeconfig. This PR should be merged after that one. Hope that makes sense. Thanks for your work on this.

[Fiona-Waters]

I tested the initializer functionality with the local Docker/Podman container backend and found some issues relating to the use of this image here that need to be addressed in this PR. This image is a controller image and does not container initializer code. Instead we should use these: kubeflow/dataset-initializer:latest, kubeflow/model-initializer:latest. Functionality should be updated to select whichever one is required for each initialiser.

[Fiona-Waters]

Those images come from here FYI https://github.com/kubeflow/trainer/tree/master/cmd/initializers

[HKanoje]

Thanks for catching this! You're right - I've updated the implementation to use the correct images.

[Fiona-Waters]

Suggested change

	python_cmd = "python -m kubeflow.storage_initializer.hugging_face "
	python_cmd = "python -m pkg.initializers.dataset"

[HKanoje]

Done!

[Fiona-Waters]

This should be /app - https://github.com/kubeflow/trainer/blob/master/cmd/initializers/dataset/Dockerfile#L3

Suggested change

	working_dir=constants.WORKSPACE_PATH,
	working_dir=/app,

[HKanoje]

Fixed! I've changed the working_dir from /workspace to /app to match the Dockerfile convention. Added a comment referencing the Dockerfile for future maintainability.

[Fiona-Waters]

Thanks for this @HKanoje I've re-run with your latest changes (and the trainer changes) and it works as expected.
/lgtm

[andreyvelich]

This looks great @HKanoje!
I left a few thoughts.

[andreyvelich]

Why do we use suppress here? We don't do that in other part of SDK.

[andreyvelich]

Can you add these test cases to the test_train() API:

sdk/kubeflow/trainer/backends/container/backend_test.py

Line 470 in 6241fae

def test_train(container_backend, test_case):

[andreyvelich]

Same question, add it to test_get_job_logs()

[andreyvelich]

This can be added to test_train()

[andreyvelich]

Can you use GHCR image by default: https://github.com/kubeflow/trainer/pkgs/container/trainer%2Fdataset-initializer

[andreyvelich]

Can we simplify this to something like we do here:

sdk/kubeflow/trainer/backends/kubernetes/utils.py

Lines 574 to 587 in 6241fae

	def get_optional_initializer_envs(
	initializer: types.BaseInitializer, required_fields: set
	) -> list[models.IoK8sApiCoreV1EnvVar]:
	"""Get the optional envs from the initializer config"""
	envs = []
	for f in fields(initializer):
	if f.name not in required_fields:
	value = getattr(initializer, f.name)
	if value is not None:
	# Convert list values (like ignore_patterns) to comma-separated strings
	if isinstance(value, list):
	value = ",".join(str(item) for item in value)
	envs.append(models.IoK8sApiCoreV1EnvVar(name=f.name.upper(), value=value))
	return envs

Env variables are always have the same name as fields, but just upper cased.

[andreyvelich]

you don't need the init_type variable since you can understand whether it is dataset or model initializer by checking type of initializer field.

Suggested change

	def build_initializer_command(initializer: types.BaseInitializer, init_type: str) -> list[str]:
	def build_initializer_command(initializer: types.BaseInitializer) -> list[str]:

[andreyvelich]

Let's remove it from now, I am not sure how data cache can be supported in Container backend at the moment.
cc @akshaychitneni

Suggested change

types.DataCacheInitializer,

[andreyvelich]

Same point

Suggested change

	def build_initializer_env(initializer: types.BaseInitializer, init_type: str) -> dict[str, str]:
	def build_initializer_env(initializer: types.BaseInitializer) -> dict[str, str]:

[andreyvelich]

Can you simplify the logic here to be consistent with what we do in Kubernetes backend:

sdk/kubeflow/trainer/backends/kubernetes/backend.py

Line 704 in 6241fae

dataset=utils.get_dataset_initializer(initializer.dataset)

Just simple define two utils function in container utils:

container_utils.get_dataset_initializer()
container_utils.get_model_initializer()

Which returns internal type that you can use in the _adapter.create_and_start_container() API:

@dataclass
class ContainerInitializer:
   image: str
   command: str
   env: dict

WDYT @HKanoje @Fiona-Waters ?

[HKanoje]

Hi @andreyvelich @Fiona-Waters,

I've addressed all the review comments in the latest commit:

• Updated default images to GHCR (ghcr.io/kubeflow/trainer/dataset-initializer:latest and ghcr.io/kubeflow/trainer/model-initializer:latest)
• Replaced suppress with try-except blocks (added # noqa: SIM105 to satisfy linter)
• Refactored initializer utils with ContainerInitializer dataclass
• Added get_dataset_initializer() and get_model_initializer() functions following kubernetes backend pattern
• Simplified env building using get_optional_initializer_envs() helper
• Removed init_type parameter (now determined from initializer type)
• Removed DataCache support (not supported in container backend)
• Merged initializer tests into test_train() and test_get_job_logs()
• Removed duplicate test functions

[andreyvelich]

Instead of passing init_type separately, you can simple add name to the ContainerInitializer type which can be:

name = dataset-initializer
name = model-initializer

Then, just use this name in the f"{self.label_prefix}/step"

[andreyvelich]

cc @HKanoje

[HKanoje]

Done! @andreyvelich

Changes made:

Added name field to ContainerInitializer dataclass
Set name="dataset-initializer" in get_dataset_initializer()
Set name="model-initializer" in get_model_initializer()
Removed init_type parameter from _run_single_initializer()
Now using container_init.name for labels and log messages

[andreyvelich]

Thanks for this @HKanoje!
Please can you create tracking issue for this: #188 (comment)
/lgtm
/assign @Fiona-Waters @kramaranya

[Fiona-Waters]

/lgtm
Thanks @HKanoje !

[HKanoje]

@andreyvelich Done, I have created the issue #290

[andreyvelich]

Thanks! /approve

[andreyvelich]

/approve

[google-oss-prow]

[APPROVALNOTIFIER] This PR is APPROVED

This pull-request has been approved by: andreyvelich

The full list of commands accepted by this bot can be found here.

The pull request process is described here

Details

Needs approval from an approver in each of these files:

• OWNERS [andreyvelich]

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment