fix: Enforce single ML policy constraint with CEL validation for Torch, MPI, and JAX

[Krishna-kg732]

What this PR solves

This PR fixes a validation issue where multiple ML runtime policies (Torch, MPI, JAX) could be configured simultaneously in a TrainingRuntime, leading to conflicting runtime configurations.

The previous validation logic : !(has(self.torch) && has(self.mpi)) which only prevented Torch and MPI from being set together, but didn't account for:

• JAX runtime policy
• Scenarios where all three policies could be partially configured
• Future extensibility for additional runtime policies

This allowed invalid configurations where users could set multiple incompatible runtime policies.

Solution

• Added comprehensive CEL validation: Updated the validation rule to [has(self.torch), has(self.mpi), has(self.jax)].filter(x, x).size() <= 1 which:

  • Creates a list of boolean values for each policy field
  • Filters for truthy values (policies that are set)
  • Ensures at most one policy is configured

• Updated PlainML plugin: Modified the EnforceMLPolicy function to check for JAX policy alongside Torch and MPI, ensuring PlainML only applies when no other runtime policy is active

Testing

• Validation occurs at the CRD level via CEL expressions
• Runtime enforcement in PlainML plugin ensures correct fallback behavior

mentioned in PR#3200

[github-actions]

🎉 Welcome to the Kubeflow Trainer! 🎉

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-trainer-team

Join the community:

• Slack: Join our #kubeflow-trainer Slack channel.
• Meetings: Attend the Kubeflow AutoML and Training Working Group 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 tightens TrainingRuntime ML policy validation to prevent configuring multiple incompatible runtime policies (Torch/MPI/JAX) at the same time, and aligns PlainML fallback behavior with that constraint.

Changes:

• Updated CRD CEL validation to enforce “at most one of torch/mpi/jax is set”.
• Updated PlainML’s EnforceMLPolicy to no-op when a JAX policy is configured (matching existing Torch/MPI behavior).

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
pkg/runtime/framework/plugins/plainml/plainml.go	Extends PlainML’s fallback guard to treat JAX as an explicitly selected runtime policy (so PlainML won’t apply).
pkg/apis/trainer/v1alpha1/trainingruntime_types.go	Replaces pairwise Torch/MPI exclusion with a single CEL rule that limits the number of configured ML policies to 1 across Torch/MPI/JAX.

[andreyvelich]

Thank you @Krishna-kg732!
/lgtm
/approve

[astefanutti]

Thanks @Krishna-kg732!

/lgtm

[astefanutti]

/retest

[andreyvelich]

@Krishna-kg732 Please rebase your PR.

[andreyvelich]

One more rebase is needed @Krishna-kg732.

[andreyvelich]

Thank you for this @Krishna-kg732!
/lgtm
/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