Adds terminationGracePeriodSeconds field to PodSpecPatch so users can
configure pod termination grace period per TrainJob via RuntimePatches.
This is needed for distributed training with PyTorch Elastic (torchrun)
where large models (70B+ parameters) require more than the default 30s
to complete JIT checkpointing before SIGKILL on node drain or TrainJob
pause.
No changes to merge logic in trainingruntime.go are required since the
existing StrategicMergePatch applied at batchv1.JobTemplateSpec level
already handles this field automatically.
Closes#3285
Signed-off-by: krishdef7 <[email protected]>
* feat(api): add progress types to api
Signed-off-by: Rob Bell <[email protected]>
* feat: progress server stub
Signed-off-by: Rob Bell <[email protected]>
* feat: setup tls for server
Signed-off-by: Rob Bell <[email protected]>
* refactor: pull out middleware
Signed-off-by: Rob Bell <[email protected]>
* refactor: remove logger message noise from tests
Signed-off-by: Rob Bell <[email protected]>
* fix: send server error logs to main logger
Signed-off-by: Rob Bell <[email protected]>
* feat: update the trainer status
Signed-off-by: Rob Bell <[email protected]>
* feat: use separate client to avoid impacting main reconciler client-side rate limits
Signed-off-by: Rob Bell <[email protected]>
* feat: configure progress server in helm chart and configmap
Signed-off-by: Rob Bell <[email protected]>
* fix: prevent race condition
Signed-off-by: Rob Bell <[email protected]>
* feat: inject progress runtime config
Signed-off-by: Rob Bell <[email protected]>
* fix: progress client config
Signed-off-by: Rob Bell <[email protected]>
* fix: update generated config
Signed-off-by: Rob Bell <[email protected]>
* feat: validate auth tokens in endpoint
Signed-off-by: Rob Bell <[email protected]>
* test: validate auth tokens in endpoint
Signed-off-by: Rob Bell <[email protected]>
* refactor: tidy tests to use same format as other tests
Signed-off-by: Rob Bell <[email protected]>
* fix: look up the cluster issuer from the service account token
Different clusters can use different issuers so look up the issuer from the trusted in-cluster service account token.
Signed-off-by: Rob Bell <[email protected]>
* refactor: rename authn middleware
Signed-off-by: Rob Bell <[email protected]>
* feat: add authn to status update endpoint
Signed-off-by: Rob Bell <[email protected]>
* refactor: tidy up reading in-cluster service account token
Signed-off-by: Rob Bell <[email protected]>
* fix: rbac permissions
Signed-off-by: Rob Bell <[email protected]>
* fix: better error response
Signed-off-by: Rob Bell <[email protected]>
* refactor: move kubebuilder annotation higher up
Signed-off-by: Rob Bell <[email protected]>
* fix: make lastUpdateTime required
Signed-off-by: Rob Bell <[email protected]>
* chore: fix config tests after rebase
Signed-off-by: Rob Bell <[email protected]>
* test: update config test with new progress server config
Signed-off-by: Rob Bell <[email protected]>
* test: add progress tracking e2e test
Signed-off-by: Rob Bell <[email protected]>
* fix: allow feature gates to be configurable from command line argument
So they can be overridden in e2e tests
Signed-off-by: Rob Bell <[email protected]>
* refactor: reuse shared progress StatusUrl
Signed-off-by: Rob Bell <[email protected]>
* test: tidy up tests
Signed-off-by: Rob Bell <[email protected]>
* refactor: allow plugins to use config
Signed-off-by: Rob Bell <[email protected]>
* refactor: centralise shared constants
Signed-off-by: Rob Bell <[email protected]>
* fix: disable http2 in progress server tls opts if disabled by config
Signed-off-by: Rob Bell <[email protected]>
* fix: minor improvements suggested by copilot
Signed-off-by: Rob Bell <[email protected]>
* chore: rename progress server => status server
Signed-off-by: Rob Bell <[email protected]>
* chore: update type doc
Signed-off-by: Rob Bell <[email protected]>
* chore: update name of feature gate
Signed-off-by: Rob Bell <[email protected]>
* fix: after rebase
Signed-off-by: Rob Bell <[email protected]>
* chore: fix typo
Signed-off-by: Rob Bell <[email protected]>
* chore: rename handler
Signed-off-by: Rob Bell <[email protected]>
* fix: run status server on all replicas
Signed-off-by: Rob Bell <[email protected]>
* fix: deep copy the rest config
Signed-off-by: Rob Bell <[email protected]>
* refactor: use ssa
Signed-off-by: Rob Bell <[email protected]>
* refactor: separate token audience per train job
Signed-off-by: Rob Bell <[email protected]>
* chore: rename status server
Signed-off-by: Rob Bell <[email protected]>
* fix: revert adding status pkg to manifest gen paths
Signed-off-by: Rob Bell <[email protected]>
* chore: update log messages
Signed-off-by: Rob Bell <[email protected]>
* chore: update api docs
Signed-off-by: Rob Bell <[email protected]>
* chore: update format
Signed-off-by: Rob Bell <[email protected]>
* fix: revert inadvertent change
Signed-off-by: Rob Bell <[email protected]>
* chore: update api name and type docs
Signed-off-by: Rob Bell <[email protected]>
* refactor: initialise oidc provider in manager startup
Signed-off-by: Rob Bell <[email protected]>
* fix: allow empty update request
Signed-off-by: Rob Bell <[email protected]>
* fix: wait for server shut down
Signed-off-by: Rob Bell <[email protected]>
* fix: handle missing train job
Signed-off-by: Rob Bell <[email protected]>
* refactor: update env var names
Signed-off-by: Rob Bell <[email protected]>
* chore: rename packages, plugins, feature gate
Signed-off-by: Rob Bell <[email protected]>
* refactor: update env var names
Signed-off-by: Rob Bell <[email protected]>
* chore: regenerate after rebase
Signed-off-by: Rob Bell <[email protected]>
* chore: fixes after rebase
Signed-off-by: Rob Bell <[email protected]>
* chore: minor doc updates
Signed-off-by: Rob Bell <[email protected]>
* chore: rename status server config
Signed-off-by: Rob Bell <[email protected]>
* chore: replace placeholder config with nil
Signed-off-by: Rob Bell <[email protected]>
* chore: revert separate runtime utils package
Signed-off-by: Rob Bell <[email protected]>
* chore: remove unused config from tests
Signed-off-by: Rob Bell <[email protected]>
---------
Signed-off-by: Rob Bell <[email protected]>
* feat(api): add ActiveDeadlineSeconds and TTLSecondsAfterFinished to Trainer APIs
- Add ActiveDeadlineSeconds *int64 to TrainJobSpec (immutable, min=1)
with reason constant TrainJobDeadlineExceededReason="DeadlineExceeded"
- Add TTLSecondsAfterFinished *int32 to TrainingRuntimeSpec (min=0)
with cross-field CEL rules blocking conflicting lifecycle fields
in the JobSet/Job template
- Add TrainJobDeadlineExceededMessage to constants
- Regenerate zz_generated.deepcopy.go
Part of KEP-2899
Signed-off-by: XploY04 <[email protected]>
* feat(controller): add TTL-based auto-deletion for finished TrainJobs
- Add early-exit guard in Reconcile before prevTrainJob DeepCopy and
removeFailedCondition, mirroring the Kubernetes Job controller pattern.
Once a TrainJob is terminal (Complete or Failed), only TTL cleanup runs.
- Add fetchTTLFromRuntime helper to retrieve TTLSecondsAfterFinished from
the referenced TrainingRuntime or ClusterTrainingRuntime.
- Add reconcileTTL helper to delete finished TrainJobs after their TTL
expires, or requeue at the exact expiry time.
Part of KEP-2899
Signed-off-by: XploY04 <[email protected]>
* feat(controller): enforce ActiveDeadlineSeconds on TrainJobs
- Add reconcileDeadline helper to terminate running TrainJobs that
exceed their ActiveDeadlineSeconds.
- Handle job suspension: StartTime is determined by the LastTransitionTime
of the Suspended=False condition (if it exists) to ensure the deadline
timer resets correctly upon resume.
- Jobs exceeding the deadline are moved to a Failed state with
reason DeadlineExceeded.
Part of KEP-2899
Signed-off-by: XploY04 <[email protected]>
* feat(webhooks): warn on short ttlSecondsAfterFinished
- Add an admission warning to both TrainingRuntime and
ClusterTrainingRuntime webhooks if the TTLSecondsAfterFinished
is configured to less than 60 seconds, as this may cause
unexpected data/log loss for data scientists due to overly
aggressive cleanup.
Part of KEP-2899
Signed-off-by: XploY04 <[email protected]>
* test: add integration tests for TrainJob TTL and deadline lifecycle
- Add integration tests for DeadlineExceeded behavior in TrainJob controller
- Add integration tests for TTL validation and warnings in webhooks
Part of KEP-2899
Signed-off-by: XploY04 <[email protected]>
* cleanup
Signed-off-by: XploY04 <[email protected]>
* fix: ensure JobSet is suspended on expiration and fix webhook test regress
Signed-off-by: XploY04 <[email protected]>
* fix(controller): return errors from fetchTTLFromRuntime for proper requeue
Signed-off-by: XploY04 <[email protected]>
* test: remove unused wantWarning from webhook tests
Signed-off-by: XploY04 <[email protected]>
* fix
Signed-off-by: XploY04 <[email protected]>
* feat: Remove from TrainingRuntime and ClusterTrainingRuntime specifications and associated validation.
Signed-off-by: XploY04 <[email protected]>
* fix(api): restore `activeDeadlineSeconds` validation rule on TrainingRuntime
Signed-off-by: XploY04 <[email protected]>
* refactor: remove unnecessary blank lines and improve logging in ClusterTrainingRuntime validation
Signed-off-by: XploY04 <[email protected]>
* feat: add activeDeadlineSeconds field to TrainJobSpec and update related tests
Signed-off-by: XploY04 <[email protected]>
* Remove unnecessary newline in test file
Signed-off-by: Yash Agarwal <[email protected]>
* Remove unnecessary newline in trainingruntime_webhook_test.go
Signed-off-by: Yash Agarwal <[email protected]>
* refactor: remove unnecessary blank lines in TrainJob reconciliation and deadline tests
Signed-off-by: XploY04 <[email protected]>
* feat: change ActiveDeadlineSeconds from pointer to value type and update related code
Signed-off-by: XploY04 <[email protected]>
* feat: add condition check for TrainJob deadline exceeded before suspending
Signed-off-by: XploY04 <[email protected]>
* feat: enhance TrainJob deadline handling and suspend logic
Signed-off-by: XploY04 <[email protected]>
* feat: enhance TrainJob deadline handling and suspend logic
Signed-off-by: XploY04 <[email protected]>
* fix
Signed-off-by: XploY04 <[email protected]>
* feat: change ActiveDeadlineSeconds to pointer type and update related logic
Signed-off-by: XploY04 <[email protected]>
* feat: ensure requeueAfter is at least 1 second when deadline is exceeded
Signed-off-by: XploY04 <[email protected]>
* feat: change ActiveDeadlineSeconds from pointer to value type and update related logic
Signed-off-by: XploY04 <[email protected]>
* fix: remove early return for finished TrainJobs to prevent status sync issues
The early return was preventing the controller from syncing JobSet
status updates for finished TrainJobs, causing e2e tests to timeout
waiting for conditions that never got updated.
This fix restores the upstream behavior where finished TrainJobs still
go through status reconciliation (though reconcileObjects is skipped).
Signed-off-by: XploY04 <[email protected]>
* fix
Signed-off-by: XploY04 <[email protected]>
* fix: remove x-kubernetes-validations for activeDeadlineSeconds in TrainingRuntime CRDs
Signed-off-by: XploY04 <[email protected]>
* fix: update validation message for ActiveDeadlineSeconds to be more generic
Signed-off-by: XploY04 <[email protected]>
* replacing suspend with delete
Signed-off-by: XploY04 <[email protected]>
* fix: remove SpecLabel and SpecAnnotation methods from TrainJobWrapper
Signed-off-by: XploY04 <[email protected]>
* fix: remove annotations field from TrainJobSpec and CRD definitions
Signed-off-by: XploY04 <[email protected]>
* ensure JobSet deletion in e2e tests
Signed-off-by: XploY04 <[email protected]>
* fix: add JobSet APIs to the scheme in e2e tests
Signed-off-by: XploY04 <[email protected]>
* fix: simplify error handling in Reconcile method of TrainJobReconciler
Signed-off-by: XploY04 <[email protected]>
---------
Signed-off-by: XploY04 <[email protected]>
Signed-off-by: Yash Agarwal <[email protected]>
* Add XGBoost runtime plugin scaffold and register in framework
Signed-off-by: Krishna-kg732 <[email protected]>
Signed-off-by: krishna-kg732 <[email protected]>
* regenerated assets
Signed-off-by: Krishna-kg732 <[email protected]>
Signed-off-by: krishna-kg732 <[email protected]>
* Add XGBoost plugin unit tests and update framework test registry
Signed-off-by: Krishna-kg732 <[email protected]>
Signed-off-by: krishna-kg732 <[email protected]>
* added :Runtime image , unit tests for :(Resources are not set, Resources are set in TrainJob) ,using the ContainerTrainerPort instead of default port
Signed-off-by: krishna-kg732 <[email protected]>
* minor changes
Signed-off-by: krishna-kg732 <[email protected]>
* added plainML fallback(JAX+XGBOOST)
Signed-off-by: krishna-kg732 <[email protected]>
* test: add XGBoost runtime resource resolution unit tests
- Add test case for resources set in Runtime only
- Add test case for resources set in both Runtime and TrainJob
- Add xgboost_distributed.yaml ClusterTrainingRuntime manifest
- Update kustomization.yaml with xgboost-runtime image
Addresses feedback from PR #3200 for complete resource resolution test coverage.
Signed-off-by: krishna-kg732 <[email protected]>
* added autogenerat3ed files
Signed-off-by: krishna-kg732 <[email protected]>
* CI fixes
Signed-off-by: krishna-kg732 <[email protected]>
* chore: fix ci lint errors
Signed-off-by: krishna-kg732 <[email protected]>
* Add XGBoost integration and E2E tests, fix Validate guard clause and nil safety
Signed-off-by: krishna-kg732 <[email protected]>
* trigger ci
Signed-off-by: krishna-kg732 <[email protected]>
* fix: Add XGBoost support to test wrapper, CEL validation, and CI publish workflow
- Add XGBoostPolicy() method to MLPolicySourceWrapper for integration tests
- Add has(self.xgboost) to MLPolicy CEL mutual-exclusivity validation rule
- Regenerate CRD manifests with updated CEL rule
- Add xgboost-runtime to build-and-push-images CI workflow
- Add XGBoost image build step to E2E setup script (commented out due to Kind bug)
Signed-off-by: krishna-kg732 <[email protected]>
* docs: Add XGBoost distributed training example notebook
Add example notebook demonstrating distributed XGBoost training using
the Kubeflow Trainer SDK and the xgboost-distributed ClusterTrainingRuntime.
The notebook uses the modern XGBoost Collective API with RabitTracker
for distributed coordination and synthetic classification data.
Signed-off-by: krishna-kg732 <[email protected]>
* test: Skip XGBoost E2E test until runtime image is published
Signed-off-by: krishna-kg732 <[email protected]>
* chore : address review comments
Signed-off-by: krishna-kg732 <[email protected]>
* chore:resolve merge conflicts
Signed-off-by: krishna-kg732 <[email protected]>
---------
Signed-off-by: Krishna-kg732 <[email protected]>
Signed-off-by: krishna-kg732 <[email protected]>
* feat: support for flux framework as hpc manager
Flux supports the majority of MPI flavors/variants, and can be used
to bootstrap MPI as a plugin. It adds other features for scheduling
and topology that can be used for simulations and ai/ml jobs.
This changeset adds the plugin implementation, including the
plugin module, tests, and an example with a small README to
serve as documentation for the time being.
Signed-off-by: vsoch <[email protected]>
* review: slim down use of annotations
We still need to put the original command in an annotation
to retieve later, but others can be re-derived from the
environment
Signed-off-by: vsoch <[email protected]>
* review: clean up tasks and remove extra FluxPolicySource
Signed-off-by: vsoch <[email protected]>
* review: ensure to unset trainer args
Signed-off-by: vsoch <[email protected]>
* review: better support for nproc per node and nodes
Signed-off-by: vsoch <[email protected]>
* review: add constants types for flux
Signed-off-by: vsoch <[email protected]>
* review: add gpu specification from resources
Signed-off-by: vsoch <[email protected]>
* Fix TrainJob NumProcPerNode
Signed-off-by: Andrey Velichkevich <[email protected]>
* Remove elasticPolicy from kubebuilder:validation
Signed-off-by: Andrey Velichkevich <[email protected]>
* chore: use go embed for long string templates flux plugin
Signed-off-by: vsoch <[email protected]>
* cleanup: removed un-used annotation
Signed-off-by: vsoch <[email protected]>
---------
Signed-off-by: vsoch <[email protected]>
Signed-off-by: Andrey Velichkevich <[email protected]>
Co-authored-by: vsoch <[email protected]>
Co-authored-by: Andrey Velichkevich <[email protected]>
