feat: Run dataset and model initializers in parallel by priyank766 · Pull Request #313 · kubeflow/sdk

priyank766 · 2026-02-22T09:00:09Z

What this PR does / why we need it:
Currently, the ContainerBackend (Docker/Podman) runs dataset and model initializers sequentially. For Large Language Model (LLM) fine-tuning or heavy data workloads, downloading both a massive dataset and a base model separately can add significant overhead to the job startup time.

This PR refactors the _run_initializers logic to use concurrent.futures.ThreadPoolExecutor. If both a dataset and a model are configured, they are now initialized in parallel threads. This reduces the total initialization time to the duration of the longest single download, rather than the sum of both.

Key technical changes:

Refactored kubeflow/trainer/backends/container/backend.py to use a thread pool for parallel initializer dispatch.
Added import concurrent.futures to the backend.
Ensured thread safety and proper error propagation (if one initializer fails, the main thread correctly identifies the failure and cleans up).
Verified that shared volume mounts work correctly during concurrent writes from separate containers.

Which issue(s) this PR fixes:
Fixes #290

Checklist:

Docs included if any changes are user facing (Internal performance improvement, no change to public API surface)
Unit tests pass: uv run pytest kubeflow/trainer/backends/container/backend_test.py
Code adheres to Ruff formatting/linting standards verified via make verify.

Verification Results:
Ran the specific backend tests to confirm the threading logic works as expected without race conditions:

uv run pytest kubeflow/trainer/backends/container/backend_test.py
# Result: 19 passed

github-actions · 2026-02-22T09:00:17Z

🎉 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! 🙏

Copilot

Pull request overview

This PR refactors the ContainerBackend to run dataset and model initializers in parallel using Python's ThreadPoolExecutor, reducing total initialization time from sequential (sum of both downloads) to parallel (duration of the longest download). This is particularly beneficial for LLM fine-tuning workloads with large datasets and models.

Changes:

Refactored _run_initializers to use concurrent.futures.ThreadPoolExecutor with max_workers=2
Updated method documentation to reflect parallel execution
Modified logging messages to indicate queueing and parallel completion

priyank766 · 2026-02-23T14:08:50Z

I wanted to know this issue is Open right i submitted PR but I had issue for PR Title can anyone can help me as I didn't get the format for PR Title issue
@astefanutti

astefanutti · 2026-02-23T14:11:58Z

/retitle feat: Run dataset and model initializers in parallel

astefanutti

/ok-to-test

…ow#290) Signed-off-by: priyank <[email protected]>

priyank766 · 2026-02-23T16:40:59Z

@astefanutti E2E Test are solved here

Fiona-Waters

Thanks @priyank766 this looks good. Just a couple of suggestions.

Signed-off-by: priyank <[email protected]>

priyank766 · 2026-03-04T13:39:34Z

Thanks for the feedback! Addressed both points:

Refactored the dataset and model initializer dispatch into a loop over an initializer_configs list, which also makes it straightforward to add new initializer types in the future.
Updated the log message as you said to include the names of the initializers that completed, e.g. Initializers completed in parallel: ['dataset', 'model'].
@Fiona-Waters

Fiona-Waters

@priyank766 This looks great now. I would like to test it on my machine, will do that as soon as I can and will get back to you. Thanks

priyank766 · 2026-03-04T14:15:18Z

Sure I would love to hear your feedback and thoughts
@Fiona-Waters

Fiona-Waters

Was able to run this successfully on my local machine using Docker. Thanks @priyank766
/lgtm

andreyvelich · 2026-03-16T17:08:50Z

/lgtm
/approve
/hold in case @Fiona-Waters or @kramaranya has any additional comments.

google-oss-prow · 2026-03-16T17:09:00Z

[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

andreyvelich · 2026-03-18T11:54:28Z

/hold cancel

Add docs/source/train/initializers.rst documenting the dataset and model initializer types (HuggingFaceDatasetInitializer, S3DatasetInitializer, DataCacheInitializer, HuggingFaceModelInitializer, S3ModelInitializer) that were added to the container backend in PRs kubeflow#188 and kubeflow#313. The guide covers: concept overview, per-initializer code examples, combined usage, ContainerBackendConfig options (images, timeout), log-based debugging, and backend limitations. Co-Authored-By: Claude Sonnet 4.6 <[email protected]>

Add docs/source/train/initializers.rst covering dataset and model initializers for the container backend (added in kubeflow#188, parallelised in kubeflow#313). Includes per-type code examples, combined usage, ContainerBackendConfig options, and debugging via get_job_logs().

Add docs/source/train/initializers.rst covering dataset and model initializers for the container backend (added in kubeflow#188, parallelised in kubeflow#313). Includes per-type code examples, combined usage, ContainerBackendConfig options, and debugging via get_job_logs(). Signed-off-by: Ayush Petwal <[email protected]>

* chore(trainer): add data and model initializers guide Add docs/source/train/initializers.rst covering dataset and model initializers for the container backend (added in #188, parallelised in #313). Includes per-type code examples, combined usage, ContainerBackendConfig options, and debugging via get_job_logs(). Signed-off-by: Ayush Petwal <[email protected]> * docs(trainer): address PR review comments on initializers guide - Fix model output path from /workspace/model-weights to /workspace/model to match the MODEL_PATH constant in constants.py - Clarify DataCacheInitializer is Kubernetes-only in the backend note and annotate the Available Initializers table row accordingly - Add DataCacheInitializer usage example with required fields (storage_uri, metadata_loc, num_data_nodes) and backend constraint note Signed-off-by: 1Ayush-Petwal <[email protected]> * docs(trainer): use latest tag for initializer images in guide Signed-off-by: 1Ayush-Petwal <[email protected]> --------- Signed-off-by: Ayush Petwal <[email protected]> Signed-off-by: 1Ayush-Petwal <[email protected]>

Copilot AI review requested due to automatic review settings February 22, 2026 09:00

google-oss-prow Bot requested review from Electronic-Waste, kramaranya and szaher February 22, 2026 09:00

google-oss-prow Bot added the size/M label Feb 22, 2026

Copilot started reviewing on behalf of priyank766 February 22, 2026 09:00 View session

Copilot AI reviewed Feb 22, 2026

View reviewed changes

priyank766 force-pushed the feat/parallel-initializers-290 branch 2 times, most recently from af5276e to be3c0a8 Compare February 22, 2026 09:07

priyank766 changed the title ~~feat(container): Run dataset and model initializers in parallel (#290)~~ feat : Run dataset and model initializers in parallel (#290) Feb 22, 2026

google-oss-prow Bot changed the title ~~feat : Run dataset and model initializers in parallel (#290)~~ feat: Run dataset and model initializers in parallel Feb 23, 2026

astefanutti reviewed Feb 23, 2026

View reviewed changes

google-oss-prow Bot added the ok-to-test label Feb 23, 2026

feat(trainer): Run dataset and model initializers in parallel (kubefl…

c05ed89

…ow#290) Signed-off-by: priyank <[email protected]>

priyank766 force-pushed the feat/parallel-initializers-290 branch from bb80231 to c05ed89 Compare February 23, 2026 16:19

priyank766 requested a review from astefanutti February 23, 2026 17:04

priyank766 mentioned this pull request Mar 4, 2026

fix(trainer): return TRAINJOB_COMPLETE when all steps are done #340

Merged

1 task

Fiona-Waters reviewed Mar 4, 2026

View reviewed changes

Comment thread kubeflow/trainer/backends/container/backend.py Outdated

Comment thread kubeflow/trainer/backends/container/backend.py Outdated

priyank766 added 2 commits March 4, 2026 19:06

feat(container): refactor initializer dispatch into a loop

3b5fdad

Signed-off-by: priyank <[email protected]>

style(container): fix ruff formatting

0bb3bc2

Signed-off-by: priyank <[email protected]>

priyank766 requested a review from Fiona-Waters March 4, 2026 14:01

Fiona-Waters reviewed Mar 4, 2026

View reviewed changes

Fiona-Waters reviewed Mar 5, 2026

View reviewed changes

google-oss-prow Bot assigned Fiona-Waters Mar 5, 2026

google-oss-prow Bot added the lgtm label Mar 5, 2026

kramaranya reviewed Mar 5, 2026

View reviewed changes

Comment thread kubeflow/trainer/backends/container/backend.py

priyank766 requested a review from kramaranya March 8, 2026 18:53

google-oss-prow Bot added the do-not-merge/hold label Mar 16, 2026

google-oss-prow Bot assigned andreyvelich Mar 16, 2026

google-oss-prow Bot added the approved label Mar 16, 2026

google-oss-prow Bot removed the do-not-merge/hold label Mar 18, 2026

google-oss-prow Bot merged commit 6d1d5a2 into kubeflow:main Mar 18, 2026
18 of 19 checks passed

google-oss-prow Bot added this to the v0.4 milestone Mar 18, 2026

priyank766 deleted the feat/parallel-initializers-290 branch March 18, 2026 12:43

Fiona-Waters mentioned this pull request Mar 21, 2026

Document data and model initializers #413

Closed

Uh oh!

Conversation

priyank766 commented Feb 22, 2026

Uh oh!

github-actions Bot commented Feb 22, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

priyank766 commented Feb 23, 2026

Uh oh!

astefanutti commented Feb 23, 2026

Uh oh!

astefanutti left a comment

Choose a reason for hiding this comment

Uh oh!

priyank766 commented Feb 23, 2026

Uh oh!

Fiona-Waters left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

priyank766 commented Mar 4, 2026

Uh oh!

Fiona-Waters left a comment

Choose a reason for hiding this comment

Uh oh!

priyank766 commented Mar 4, 2026

Uh oh!

Fiona-Waters left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

andreyvelich commented Mar 16, 2026

Uh oh!

google-oss-prow Bot commented Mar 16, 2026

Uh oh!

andreyvelich commented Mar 18, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants