Skip to content

feat(api): introduce OptimizationJob CRD under Trainer - #3552

Merged
google-oss-prow[bot] merged 1 commit into
kubeflow:masterfrom
aniket2405:optimization-job-crd
Aug 7, 2026
Merged

feat(api): introduce OptimizationJob CRD under Trainer#3552
google-oss-prow[bot] merged 1 commit into
kubeflow:masterfrom
aniket2405:optimization-job-crd

Conversation

@aniket2405

@aniket2405 aniket2405 commented May 26, 2026

Copy link
Copy Markdown
Member

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:

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).

Copilot AI review requested due to automatic review settings May 26, 2026 13:55
@github-actions

Copy link
Copy Markdown
Contributor

🎉 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:

Feel free to ask questions in the comments if you need any help or clarification!
Thanks again for contributing to Kubeflow! 🙏

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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'.

Comment thread pkg/apis/trainer/v1alpha1/optimizationjob_types.go Outdated
Comment thread pkg/apis/trainer/v1alpha1/optimizationjob_types.go Outdated
Comment thread pkg/apis/trainer/v1alpha1/optimizationjob_types.go Outdated
@@ -0,0 +1,170 @@
/*
Copyright 2026 The Kubeflow Authors.

Comment thread hack/violation_exception_v1alpha1.list Outdated
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
Comment on lines +51 to +60
// +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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you look into above validations?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment on lines +84 to +85
"name": obj.get("name") if obj.get("name") is not None else '',
"value": obj.get("value") if obj.get("value") is not None else ''
// 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"`

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we just use trainJobSpec here?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Comment thread pkg/apis/trainer/v1alpha1/optimizationjob_types.go Outdated
Comment thread pkg/apis/trainer/v1alpha1/optimizationjob_types.go Outdated
// optimalParameters is a list of the hyperparameter assignments that won.
// +listType=atomic
// +optional
OptimalParameters []ParameterAssignment `json:"optimalParameters,omitempty"`

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Makes sense, it was redundant.

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

Copy link
Copy Markdown
Contributor

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

// 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 {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Explore concrete types that can be extended

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ack.

@google-oss-prow google-oss-prow Bot added size/XXL and removed size/XL labels Jun 8, 2026
@aniket2405 aniket2405 changed the title Introduce OptimizationJob API types under Trainer group feat(api): introduce OptimizationJob CRD under Trainer Jun 10, 2026
Comment thread pkg/apis/trainer/v1alpha1/optimizationjob_types.go Outdated
Comment thread pkg/apis/trainer/v1alpha1/optimizationjob_types.go Outdated
Comment thread pkg/apis/trainer/v1alpha1/optimizationjob_types.go Outdated
Comment thread pkg/apis/trainer/v1alpha1/optimizationjob_types.go Outdated
Comment thread pkg/apis/trainer/v1alpha1/optimizationjob_types.go Outdated
Comment thread pkg/apis/trainer/v1alpha1/optimizationjob_types.go Outdated
Comment thread pkg/apis/trainer/v1alpha1/optimizationjob_types.go Outdated
Comment thread pkg/apis/trainer/v1alpha1/optimizationjob_types.go
@google-oss-prow

Copy link
Copy Markdown
Contributor

@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.

Comment thread pkg/apis/trainer/v1alpha1/optimizationjob_types.go Outdated
Comment thread pkg/apis/trainer/v1alpha1/optimizationjob_types.go
Comment thread hack/e2e-setup-cluster.sh
Direction ObjectiveDirection `json:"direction,omitempty"`
}

// +kubebuilder:validation:ExactlyOneOf=random;grid

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

+1 , one doubt: since the current MLPolicy rule is <= 1 (zero policies allowed), wouldn't AtMostOneOf be the right fit here instead of ExactlyOneOf? Or do we want to tighten it?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right.
ExactlyOneOf will be breaking change, given MLPolicy rule is <= 1. AtMostOneOf should be the right way forward imo.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point.

// +kubebuilder:validation:MinLength=1
// +kubebuilder:validation:MaxLength=64
// +required
Value string `json:"value,omitempty"`

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point, @aniket2405 let's define // +kubebuilder:validation:items:MaxLength=64 for categorical choices too.

@Krishna-kg732

Copy link
Copy Markdown
Member

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>
@aniket2405
aniket2405 force-pushed the optimization-job-crd branch from 12eaf1e to 992209b Compare August 6, 2026 18:25
@andreyvelich

Copy link
Copy Markdown
Member

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

@robert-bell robert-bell left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

/lgtm

Thanks everyone! Excited to see this feature!

@akshaychitneni

Copy link
Copy Markdown
Contributor

Great work @aniket2405
/lgtm

@andreyvelich

Copy link
Copy Markdown
Member

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

@google-oss-prow

Copy link
Copy Markdown
Contributor

[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:

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

@google-oss-prow
google-oss-prow Bot merged commit 0cd558c into kubeflow:master Aug 7, 2026
4 checks passed
@google-oss-prow google-oss-prow Bot added this to the v2.4 milestone Aug 7, 2026
github-actions Bot pushed a commit to tariq-hasan/trainer that referenced this pull request Aug 8, 2026
Signed-off-by: aniket2405 <aniketshaha2001@gmail.com>
(cherry picked from commit 0cd558c)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

KEP-3562: OptimizationJob API and Validation Framework

8 participants