fix(operator): guard against nil trainer PodSet in JAX EnforceMLPolicy

[immanuwell]

What this PR does / why we need it:

EnforceMLPolicy in the JAX plugin calls info.FindPodSetByAncestor(constants.AncestorTrainer) which returns *PodSet and can be nil if no replicatedJob carries the trainer ancestor label. The next line dereferences it directly (trainerPS.Count != nil), causing a nil pointer dereference panic at runtime.

The fix adds a trainerPS != nil guard, matching the pattern already used in the Torch plugin. The XGBoost plugin uses a similar early-return guard too -- the JAX plugin was just missing it.

Reproduce:

Create a ClusterTrainingRuntime with JAX mlPolicy but without a replicatedJob labeled trainer.kubeflow.org/trainjob-ancestor-step: trainer, then submit a TrainJob against it. The controller panics when processing EnforceMLPolicy.

Or run the new unit test case no panic when JAX policy is set but trainer PodSet is absent after reverting the one-line fix -- it will panic.

Which issue(s) this PR fixes:
Fixes #

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

[google-oss-prow]

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by:
Once this PR has been reviewed and has the lgtm label, please assign terrytangyuan for approval. For more information see the Kubernetes Code Review Process.

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

Details

Needs approval from an approver in each of these files:

• OWNERS

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