kubeflow / sdk / 21941961384

73%

push

github

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 s... (continued)

160 of 195 new or added lines in 7 files covered. (82.05%)

2977 of 4327 relevant lines covered (68.8%)

1.38 hits per line