feat(api): introduce OptimizationJob CRD under Trainer by aniket2405 · Pull Request #3552 · kubeflow/trainer

aniket2405 · 2026-05-26T13:55:11Z

What this PR does / why we need it:

This PR introduces the core API schema for the OptimizationJob CRD. This is a foundational step in the consolidation of Katib hyperparameter tuning into the Trainer v2 ecosystem.

By integrating OptimizationJob natively under the trainer.kubeflow.org API group, we unify the user experience and enable shared observability with TrainJobs.

Key features of this API schema:

CEL-Native Validation: Implements server-side CEL validation for SearchSpace and TrialConfig constraints, ensuring high-quality hyperparameter search configurations at the API layer.

Declarative HPO: Enables the OptimizationJob -> TrainJob hierarchy, leveraging TrainJob's native progress tracking to eliminate the legacy Katib DB dependency.

Manifests & Tooling: Scaffolds all necessary deepcopy, clientset, RBAC, and CRD manifests.

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

Checklist:

API design documentation (KEP) submitted in feat: KEP-3562: Introduce OptimizationJob CRD #3565.
Docs included if any changes are user facing.

Special notes for reviewers:

Autogenerated Code: Please note that ~18k lines of this PR consist of autogenerated manifests, zz_generated.deepcopy.go, and clientset code.

Core Logic: Focus your review efforts on api/v1alpha1/optimizationjob_types.go (the API schema), pkg/webhooks/ and the test/integration/webhooks directory (the validation and defaulting logic).

github-actions · 2026-05-26T13:56:45Z

🎉 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

Note

Copilot was unable to run its full agentic suite in this review.

This PR introduces the OptimizationJob v1alpha1 API surface and wires it through the generated Kubernetes clients (clientset, informer, lister), CRDs, OpenAPI, applyconfigurations, and the generated Python SDK models.

Changes:

Add OptimizationJob CRD + Go API types (OptimizationJobSpec/Status and related structs like Algorithm, SearchSpace, TrialConfig).
Generate/extend Go clientset, informers, listers, applyconfiguration helpers, and OpenAPI definitions for the new resource.
Generate Python OpenAPI client models for the new types.

Reviewed changes

Copilot reviewed 33 out of 43 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
pkg/client/listers/trainer/v1alpha1/optimizationjob.go	Adds lister interfaces/implementations for `OptimizationJob`.
pkg/client/listers/trainer/v1alpha1/expansion_generated.go	Adds lister expansion hooks for `OptimizationJob`.
pkg/client/informers/externalversions/trainer/v1alpha1/optimizationjob.go	Adds shared informer for `OptimizationJob`.
pkg/client/informers/externalversions/trainer/v1alpha1/interface.go	Exposes `OptimizationJobs()` informer on the v1alpha1 interface.
pkg/client/informers/externalversions/generic.go	Enables generic informer lookup for `optimizationjobs`.
pkg/client/clientset/versioned/typed/trainer/v1alpha1/trainer_client.go	Exposes typed client method `OptimizationJobs(namespace)`.
pkg/client/clientset/versioned/typed/trainer/v1alpha1/optimizationjob.go	Adds typed REST client for `OptimizationJob` including Apply/Status.
pkg/client/clientset/versioned/typed/trainer/v1alpha1/generated_expansion.go	Adds typed client expansion hook for `OptimizationJob`.
pkg/client/clientset/versioned/typed/trainer/v1alpha1/fake/fake_trainer_client.go	Adds fake typed client method `OptimizationJobs(namespace)`.
pkg/client/clientset/versioned/typed/trainer/v1alpha1/fake/fake_optimizationjob.go	Adds fake typed client implementation for `OptimizationJob`.
pkg/client/applyconfiguration/utils.go	Registers applyconfiguration `ForKind` mappings for new types.
pkg/client/applyconfiguration/trainer/v1alpha1/trialconfig.go	Adds applyconfiguration builder for `TrialConfig`.
pkg/client/applyconfiguration/trainer/v1alpha1/settingkv.go	Adds applyconfiguration builder for `SettingKV`.
pkg/apis/trainer/v1alpha1/zz_generated.openapi.go	Adds OpenAPI definitions for `OptimizationJob` and related types.
pkg/apis/trainer/v1alpha1/optimizationjob_types.go	Introduces the `OptimizationJob` API types and validation markers.
pkg/apis/trainer/v1alpha1/groupversion_info.go	Registers `OptimizationJob` types into the scheme.
manifests/base/crds/trainer.kubeflow.org_optimizationjobs.yaml	Adds CRD manifest for `OptimizationJob` (base manifests).
charts/kubeflow-trainer/crds/trainer.kubeflow.org_optimizationjobs.yaml	Adds CRD manifest for `OptimizationJob` (Helm chart).
hack/violation_exception_v1alpha1.list	Adds API rule violation exceptions for missing list-type markers.
api/openapi-spec/swagger.json	Updates swagger spec with new schemas.
api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_trial_config.py	Adds Python model for `TrialConfig`.
api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_setting_kv.py	Adds Python model for `SettingKV`.
api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_search_space.py	Adds Python model for `SearchSpace`.
api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_parameter_assignment.py	Adds Python model for `ParameterAssignment`.
api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_parameter.py	Adds Python model for `Parameter`.
api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_optimization_job_status.py	Adds Python model for `OptimizationJobStatus`.
api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_optimization_job_spec.py	Adds Python model for `OptimizationJobSpec`.
api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_optimization_job_list.py	Adds Python model for `OptimizationJobList`.
api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_optimization_job.py	Adds Python model for `OptimizationJob`.
api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_objective.py	Adds Python model for `Objective`.
api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_best_trial.py	Adds Python model for `BestTrial`.
api/python_api/kubeflow_trainer_api/models/trainer_v1alpha1_algorithm.py	Adds Python model for `Algorithm`.
api/python_api/kubeflow_trainer_api/models/init.py	Exports the newly generated Python models.

Files not reviewed (10)

pkg/apis/trainer/v1alpha1/zz_generated.deepcopy.go: Language not supported
pkg/client/applyconfiguration/trainer/v1alpha1/algorithm.go: Language not supported
pkg/client/applyconfiguration/trainer/v1alpha1/besttrial.go: Language not supported
pkg/client/applyconfiguration/trainer/v1alpha1/objective.go: Language not supported
pkg/client/applyconfiguration/trainer/v1alpha1/optimizationjob.go: Language not supported
pkg/client/applyconfiguration/trainer/v1alpha1/optimizationjobspec.go: Language not supported
pkg/client/applyconfiguration/trainer/v1alpha1/optimizationjobstatus.go: Language not supported
pkg/client/applyconfiguration/trainer/v1alpha1/parameter.go: Language not supported
pkg/client/applyconfiguration/trainer/v1alpha1/parameterassignment.go: Language not supported
pkg/client/applyconfiguration/trainer/v1alpha1/searchspace.go: Language not supported

Comments suppressed due to low confidence (1)

pkg/client/informers/externalversions/trainer/v1alpha1/interface.go:1

Grammar: change 'a OptimizationJobInformer' to 'an OptimizationJobInformer'.

@@ -0,0 +1,170 @@
+/*
+Copyright 2026 The Kubeflow Authors.
+


+API rule violation: list_type_missing,github.com/kubeflow/trainer/v2/pkg/apis/trainer/v1alpha1,OptimizationJobSpec,Objectives
+API rule violation: list_type_missing,github.com/kubeflow/trainer/v2/pkg/apis/trainer/v1alpha1,OptimizationJobSpec,Parameters
+API rule violation: list_type_missing,github.com/kubeflow/trainer/v2/pkg/apis/trainer/v1alpha1,OptimizationJobStatus,Conditions
+API rule violation: list_type_missing,github.com/kubeflow/trainer/v2/pkg/apis/trainer/v1alpha1,SearchSpace,List


akshaychitneni · 2026-06-01T23:41:52Z

+// +kubebuilder:validation:XValidation:rule="self.type != 'categorical' || has(self.list)",message="list must be provided when type is categorical"
+// +kubebuilder:validation:XValidation:rule="self.type == 'categorical' || (has(self.min) && has(self.max))",message="min and max must be provided for int or double types"
+type SearchSpace struct {
+	// +kubebuilder:validation:Enum=int;double;categorical
+	Type string `json:"type"` // e.g., int, double, categorical
+
+	Max string `json:"max,omitempty"`
+	Min string `json:"min,omitempty"`
+
+	// +optional


Could you look into above validations?

Addressed. Updated the CEL rules to check list sizes and added MinLength=1 / MinItems=1 markers to ensure empty strings and arrays are blocked.

+            "name": obj.get("name") if obj.get("name") is not None else '',
+            "value": obj.get("value") if obj.get("value") is not None else ''


akshaychitneni · 2026-06-01T22:58:11Z

+	// TrialTemplate acts as a generic wrapper for the underlying workload.
+	// Parameters are injected via native Kubernetes Environment Variables, replacing regex.
+	// +kubebuilder:pruning:PreserveUnknownFields
+	TrialTemplate runtime.RawExtension `json:"trialTemplate"`


Should we just use trainJobSpec here?

Since our MVP focuses on orchestrating TrainJobs, embedding the TrainJobSpec directly makes sense. But do you think this should be out long-term thought as well? I believe we discussed we want to keep it tightly coupled with TrainJob, but should we not consider keeping our option open for other types of Jobs?

akshaychitneni · 2026-06-02T00:02:04Z

+	// optimalParameters is a list of the hyperparameter assignments that won.
+	// +listType=atomic
+	// +optional
+	OptimalParameters []ParameterAssignment `json:"optimalParameters,omitempty"`


Should we just say parameters here? it should implicitly be optimal as it is part of besttrial

Makes sense, it was redundant.

akshaychitneni · 2026-06-02T00:12:14Z

Thanks for your work @aniket2405. Added a few comments on the api spec. please take a look

akshaychitneni · 2026-06-02T16:01:00Z

+// SearchSpace defines the type and exact boundaries for the algorithm to search.
+// +kubebuilder:validation:XValidation:rule="self.type != 'categorical' || (has(self.list) && size(self.list) > 0)",message="list must be provided and contain at least one item when type is categorical"
+// +kubebuilder:validation:XValidation:rule="self.type == 'categorical' || (has(self.min) && has(self.max) && self.min != ” && self.max != ”)",message="min and max must be provided and be non-empty for int or double types"
+type SearchSpace struct {


Explore concrete types that can be extended

google-oss-prow · 2026-06-29T14:54:34Z

@andreyvelich: GitHub didn't allow me to request PR reviews from the following users: mimowo.

Note that only kubeflow members and repo collaborators can review this PR, and authors cannot review their own PRs.

Details

In response to this:

Yeah, we can introduce the suspend fields in the later phases, once we integrate OptimizationJob with Kueue!

/cc @tenzen-y @kannon92 @mimowo @kaisoz

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes/test-infra repository.

Krishna-kg732 · 2026-08-06T09:58:41Z

Was reading through the status types and had a few thoughts, all in the "cheaper now than after v1alpha1" bucket. Push back if these were already considered.

status.result as a single struct — with multi-objective on the roadmap (#3799) that eventually becomes a Pareto set. Worth making it Results []Result now so the change stays additive?

Should Result also carry the objective value? Katib's OptimalTrial has Observation. Right now the winning score lives only on the TrainJob, and if that gets GC'd — which is the premise of #3744 — it's gone. Could reuse Metric from trainjob_types.go:593 so it costs no new schema.

Status *OptimizationJobStatus (L49) — any reason for the pointer? TrainJob uses a value with omitzero (trainjob_types.go:53), which I think is the more usual shape for a status subresource, but I may be missing context.

Smaller ones: no +kubebuilder:printcolumn (cf. trainjob_types.go:34), so kubectl get optimizationjobs won't show progress or outcome. And the condition type/reason constants — #3828 declares them in pkg/controller while TrainJob keeps them in the API package (trainjob_types.go:56); worth pulling them here? That'd also be a natural moment to reconcile the Running condition #3828 sets with the Created one KEP §7.4 describes, since the two currently disagree.

Probably just a doc nit: Objective.Metric says Defaults to "loss" (L111) but it's +required with no default marker — stale comment, or missing marker?

First two overlap with what @Sridhar1030 is proposing in #3856, cross-posting there so we don't end up with two status designs.

Signed-off-by: aniket2405 <aniketshaha2001@gmail.com>

andreyvelich · 2026-08-06T18:44:18Z

Thanks for the updates @aniket2405!
/lgtm
/assign @akshaychitneni @robert-bell @Sridhar1030

robert-bell

/lgtm

Thanks everyone! Excited to see this feature!

akshaychitneni · 2026-08-07T15:51:32Z

Great work @aniket2405
/lgtm

andreyvelich · 2026-08-07T16:08:16Z

Super excited to see this moving forward!
/lgtm
/approve

google-oss-prow · 2026-08-07T16:08:26Z

[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

Signed-off-by: aniket2405 <aniketshaha2001@gmail.com> (cherry picked from commit 0cd558c)

Grid search enumerates the full cartesian product of the search space up front, so numTrials has to fit that product. The OptimizationJob API does not enforce this, so a grid job asking for more trials than the search space can produce passes admission and then fails from inside the suggestion container, with nothing pointing at the offending field. The companion constraint, that grid requires a categorical search space, already landed as a CEL rule on OptimizationJobSpec in kubeflow#3552. The budget check cannot follow it into CEL: computing the number of combinations is a fold over the parameter list (a product across choice-set sizes), which CEL cannot express. It therefore has to live in a validating webhook, and OptimizationJob had none. Add an OptimizationJobValidator and register it in the webhook setup: - ValidateCreate rejects a grid job whose numTrials exceeds the number of grid combinations, mirroring Katib's ValidateAlgorithmSettings, which rejects the Experiment before any trial is scheduled. Equal counts are allowed, matching Katib. Random search is unbounded and is skipped. - The combination product terminates early once it reaches numTrials. Choice sets have MinItems=1 so the product is monotonic, which makes the early return correct, and it keeps the product from overflowing int64: parameters and choices are each bounded only by MaxItems=100, so a naive product can reach 100^100 and wrap to a value that would spuriously reject a valid config. - A non-categorical parameter under grid defers to the CEL rule rather than guessing a combination count. - ValidateUpdate is a no-op because the spec is already immutable via CEL, so a create-time budget check cannot be bypassed on update. No mutating webhook is added. The defaults the API needs (numTrials, parallelTrials, searchAlgorithm) are all schema defaults on the CRD, so a Go defaulter would only duplicate them and risk drifting from the schema. Add the OptimizationJobKind constant used by the webhook setup, and regenerate the webhook manifest. The Helm ValidatingWebhookConfiguration template gains the matching entry, with a chart unit test covering it. The CRDs are unchanged: the new marker is a webhook marker, not a schema marker. Add unit tests for the budget check, covering the equal-to-combinations boundary, the single-choice case, the wide-parameter overflow regression, and the non-categorical deferral, plus envtest integration cases covering both the CEL path and the webhook path. Signed-off-by: Sanskar Singh <sanskarsinghty1234@gmail.com>

Copilot AI review requested due to automatic review settings May 26, 2026 13:55

google-oss-prow Bot requested review from akshaychitneni and jinchihe May 26, 2026 13:56

google-oss-prow Bot added the size/XL label May 26, 2026

Copilot AI reviewed May 26, 2026

View reviewed changes

akshaychitneni reviewed Jun 1, 2026

View reviewed changes

Comment thread pkg/apis/trainer/v1alpha1/optimizationjob_types.go Outdated

akshaychitneni reviewed Jun 1, 2026

View reviewed changes

Comment thread pkg/apis/trainer/v1alpha1/optimizationjob_types.go Outdated

akshaychitneni reviewed Jun 2, 2026

View reviewed changes

Comment thread pkg/apis/trainer/v1alpha1/optimizationjob_types.go

akshaychitneni reviewed Jun 2, 2026

View reviewed changes

google-oss-prow Bot added size/XXL and removed size/XL labels Jun 8, 2026

aniket2405 changed the title ~~Introduce OptimizationJob API types under Trainer group~~ feat(api): introduce OptimizationJob CRD under Trainer Jun 10, 2026

aniket2405 commented Jun 23, 2026

View reviewed changes

Comment thread pkg/apis/trainer/v1alpha1/optimizationjob_types.go Outdated

akshaychitneni reviewed Jun 28, 2026

View reviewed changes

Comment thread pkg/apis/trainer/v1alpha1/optimizationjob_types.go Outdated

akshaychitneni reviewed Jun 28, 2026

View reviewed changes

Comment thread pkg/apis/trainer/v1alpha1/optimizationjob_types.go Outdated

akshaychitneni reviewed Jun 29, 2026

View reviewed changes