From cdb5db60db45fb66d1f9914855f26a092dc3e161 Mon Sep 17 00:00:00 2001 From: Sridhar1030 Date: Sun, 9 Aug 2026 13:43:04 +0530 Subject: [PATCH 1/3] docs(proposal): add Phase 2 Kubeflow SDK integration to KEP-3562 Add the SDK client API design for OptimizationJob per the SDK & ML Experience call (2026-07-29): single TrainerClient with optimize(), get_hyperparameters() helper, SDK-to-CRD type mappings, and the migration path from the Katib-backed OptimizerClient. Part of kubeflow/trainer#3794 Signed-off-by: Sridhar1030 --- proposals/2605-optimization-job-crd/README.md | 168 +++++++++++++++++- 1 file changed, 164 insertions(+), 4 deletions(-) diff --git a/proposals/2605-optimization-job-crd/README.md b/proposals/2605-optimization-job-crd/README.md index 74adabe953..bf66d19ed2 100644 --- a/proposals/2605-optimization-job-crd/README.md +++ b/proposals/2605-optimization-job-crd/README.md @@ -14,8 +14,9 @@ 6. [Sample YAML (Phase 1)](#6-sample-yaml-phase-1) 7. [Reconciliation & Architecture (Phase 1)](#7-reconciliation--architecture-phase-1) 8. [Open Discussions](#8-open-discussions) -9. [Implementation History](#9-implementation-history) -10. [Alternatives](#10-alternatives) +9. [Phase 2: Kubeflow SDK Integration](#9-phase-2-kubeflow-sdk-integration) +10. [Implementation History](#10-implementation-history) +11. [Alternatives](#11-alternatives) --- @@ -474,17 +475,176 @@ Metric reporting from the `TrainJob` is strictly asynchronous and relies entirel Pruning decisions are computed controller-side based on this monotonic metric history. A "Stop Signal" is then propagated to the training runtime as a non-blocking annotation or status field, which the Kubeflow SDK periodically polls. Synchronous "kill" calls during metric reporting create tight coupling and latency bottlenecks; by separating reporting from termination, we ensure the controller remains performant under heavy trial loads. -## Implementation History +## 9. Phase 2: Kubeflow SDK Integration + +This section defines the Python SDK (https://github.com/kubeflow/sdk) client APIs for the +`OptimizationJob` CRD. It replaces the "[Separate KEP for this integration]" placeholder from +Story 4 and follows the direction agreed on the Kubeflow SDK & ML Experience call (2026-07-29): + +- **Keep the single `TrainerClient()`.** Data scientists should not switch between disparate + clients to train and to tune. Training and hyperparameter optimization are parts of the same + workflow. +- **APIs must be easy to understand from the ML Engineer's point of view.** The same + `TrainJobTemplate` used for training is reused for optimization, and hyperparameter + consumption inside the training function requires no YAML or CLI argument changes. + +### 9.1. User-Facing API: `TrainerClient().optimize()` + +The `optimize()` method is added to the existing `TrainerClient`, alongside `train()`. It +accepts the same trial template concepts that `train()` already uses, plus the optimization +configuration types that exist in the SDK today (`kubeflow.optimizer.types`): + +```python +from kubeflow.optimizer import Objective, RandomSearch, Search, TrialConfig +from kubeflow.trainer import CustomTrainer, TrainerClient, TrainJobTemplate + + +def train_fn(): + from kubeflow.trainer import get_hyperparameters + + hps = get_hyperparameters() # {"learning_rate": 0.0021, "batch_size": "32"} + ... + + +client = TrainerClient() + +job_name = client.optimize( + trial_template=TrainJobTemplate( + trainer=CustomTrainer(func=train_fn), + runtime="torch-distributed", + ), + search_space={ + "learning_rate": Search.loguniform(0.0001, 0.1), + "batch_size": Search.choice([16, 32, 64]), + }, + objectives=[Objective(metric="val_loss", direction="minimize")], + trial_config=TrialConfig(num_trials=20, parallel_trials=4), + algorithm=RandomSearch(random_state=42), +) + +client.wait_for_optimization_job_status(job_name) +best = client.get_best_results(job_name) # Result(parameters={...}, metrics=[...]) +``` + +`optimize()` builds an `OptimizationJob` custom resource (group `trainer.kubeflow.org`, +version `v1alpha1`) and submits it via the Kubernetes API. The `trainJobTemplate` in the spec +is produced by the exact same internal spec builder that `train()` uses, so runtimes, +initializers, and custom trainers behave identically for training and for trials. + +Lifecycle methods added to `TrainerClient`, mirroring the existing TrainJob methods: + +| Method | Behavior | +| --- | --- | +| `optimize(...) -> str` | Create the OptimizationJob, return its generated name. | +| `get_optimization_job(name)` | Read the CR, map `status.conditions` to `Created`/`Complete`/`Failed`. | +| `list_optimization_jobs()` | List OptimizationJobs in the namespace. | +| `wait_for_optimization_job_status(name, status={...}, timeout, polling_interval, callbacks)` | Poll until a desired status; raise on unexpected `Failed`. | +| `get_best_results(name)` | Read `status.result` (trainJobName + parameter assignments). | +| `get_optimization_job_logs(name, trial_name=None, follow=False)` | Delegate to TrainJob log streaming; defaults to the best trial. | +| `get_optimization_job_events(name)` | Kubernetes events for the OptimizationJob and its TrainJobs. | +| `delete_optimization_job(name)` | Delete the CR. | + +Because trials are plain `TrainJob` resources, every trial is also individually accessible +through the existing `get_job()`, `get_job_logs()`, and `get_job_events()` APIs — no +Trial-specific SDK surface is needed. + +### 9.2. In-Trial Helper: `get_hyperparameters()` + +Per Section 8.2, the controller injects suggestions as `KUBEFLOW_TRAINER_OPT_` +environment variables. The SDK provides the helper promised there: + +```python +def get_hyperparameters() -> dict[str, int | float | str]: + """Read hyperparameter suggestions injected by the OptimizationJob controller.""" +``` + +- Scans the environment for the `KUBEFLOW_TRAINER_OPT_` prefix and strips it, lower-casing + the remaining key back to the parameter name. +- Values from `uniform`/`logUniform` spaces are coerced with the CRD `type` + (`Int` -> `int`, `Float` -> `float`); `categorical` choices remain strings. +- Returns an empty dict when running outside an optimization trial, so the same training + function works unchanged under `train()`. + +### 9.3. SDK-to-CRD Type Mapping + +The SDK reuses the existing search-space and algorithm types from KEP-46 so users don't +learn a second vocabulary. Numeric bounds are serialized as strings to satisfy the CRD +`Double` pattern: + +| SDK (today) | OptimizationJob CRD | +| --- | --- | +| `Search.uniform(min, max)` | `searchSpace.uniform {min: "...", max: "...", type: Float}` | +| `Search.loguniform(min, max)` | `searchSpace.logUniform {min: "...", max: "...", type: Float}` | +| `Search.choice([...])` | `searchSpace.categorical {choices: [...]}` (values stringified) | +| `Objective(metric, direction="minimize")` | `objectives[0] {metric, direction: Minimize}` | +| `RandomSearch(random_state=N)` | `searchAlgorithm.random {seed: N}` | +| `GridSearch()` | `searchAlgorithm.grid {}` | +| `TrialConfig(num_trials, parallel_trials)` | `numTrials`, `parallelTrials` | + +Constraints surfaced client-side as `ValueError` before submission, matching CRD validation: +exactly one objective (multi-objective is a Phase 2/3 CRD goal), and +`TrialConfig.max_failed_trials` is rejected until a `TrialPolicy` exists in the CRD. + +### 9.4. Results and Trial History + +Phase 1 status exposes only `status.result` (best trial name + parameter assignments), so +`get_best_results()` returns those parameters, and metrics/trial listings are limited until +[KEP-3744](https://github.com/kubeflow/trainer/pull/3856) stores trial history natively in +the OptimizationJob status. Once KEP-3744 lands, `get_optimization_job(name).trials` and the +metrics in `Result` are populated from status without additional API calls. + +### 9.5. Migration from the Katib-Backed OptimizerClient + +The existing `OptimizerClient` (KEP-46) maps the same Python types onto Katib `Experiment` +CRs. Transition plan: + +1. **Introduce** `TrainerClient().optimize()` targeting `OptimizationJob` (this section). + The Katib-backed `OptimizerClient` remains untouched and fully supported. +2. **Deprecate** `OptimizerClient` with a warning once the OptimizationJob controller is + released and feature parity for Random/Grid search is verified. +3. **Remove** the Katib backend in a future SDK major/minor release, per the standard + Kubeflow deprecation policy. + +Because both paths share the `Search`/`Objective`/`TrialConfig` types, user migration is +limited to swapping `OptimizerClient().optimize(...)` for `TrainerClient().optimize(...)`. + +### 9.6. Implementation Notes & Open Questions + +- A working prototype of the CR construction, status mapping, and unit-test approach exists + in [kubeflow/sdk#726](https://github.com/kubeflow/sdk/pull/726) (held pending this design). + It validates the type mappings above against a mocked Kubernetes API; E2E validation is + blocked until the OptimizationJob controller merges. +- **Open:** whether `get_optimization_job()`-style names are kept, or the existing + `get_job()`/`wait_for_job_status()` methods gain a job-type abstraction instead. +- **Open:** whether the optimization types (`Search`, `Objective`, `TrialConfig`) are + re-exported from `kubeflow.trainer` so a single import namespace covers the whole workflow. +- **Open:** whether generated Python models for `trainer.kubeflow.org/v1alpha1` types are + published in `kubeflow-trainer-api` before the SDK integration merges, or the SDK builds + plain dictionaries in the interim. + +## 10. Implementation History - **2026-06-01:** Initial KEP draft creation for the `OptimizationJob` CRD. +- **2026-08-09:** Added Phase 2 Kubeflow SDK integration design (`TrainerClient().optimize()`). -## 10. Alternatives +## 11. Alternatives ### Extend Existing Katib Experiment and Trial CRDs Instead of introducing the `OptimizationJob` CRD, we could have updated the existing Katib `Experiment` and `Trial` CRDs to support `TrainJob` references. Katib's current architecture is fundamentally built around unstructured YAML templates and arbitrary CRD support, relying heavily on brittle regex string substitution (e.g., `${searchSpace.lr}`). Fitting this legacy structure to support the strictly typed `TrainJob` v2 API would require massive breaking changes to Katib or result in a disjointed user experience. Introducing a purpose-built `OptimizationJob` ensures tight coupling with `TrainJob` and native Kubernetes validation. +### Standalone OptimizerClient Backend for OptimizationJob (SDK) + +For Phase 2, instead of extending `TrainerClient`, the trainer-native backend could be added +as a second backend of the existing `OptimizerClient`, selected via a dedicated backend +config (prototyped in [kubeflow/sdk#726](https://github.com/kubeflow/sdk/pull/726)). This +was rejected on the SDK & ML Experience call (2026-07-29): it keeps training and tuning in +two separate clients, forces users to choose a backend for what is an implementation detail, +and contradicts the goal of a single, MLE-friendly entry point. The prototype's internals +(spec construction, status mapping, tests) carry over to the `TrainerClient().optimize()` +design unchanged. + ### Stateful Sidecars with Persistent Storage Katib currently deploys a stateful `Suggestion` sidecar and a persistent DB layer for every experiment. We could have replicated this architecture for `OptimizationJob`. From 8d2238770e7d18882e381e9e9f9dd76685e3280e Mon Sep 17 00:00:00 2001 From: Sridhar1030 Date: Sun, 9 Aug 2026 13:46:26 +0530 Subject: [PATCH 2/3] fixed import and Search* parameters Signed-off-by: Sridhar1030 --- proposals/2605-optimization-job-crd/README.md | 247 ++++++++++++++---- 1 file changed, 189 insertions(+), 58 deletions(-) diff --git a/proposals/2605-optimization-job-crd/README.md b/proposals/2605-optimization-job-crd/README.md index bf66d19ed2..a5afd1212d 100644 --- a/proposals/2605-optimization-job-crd/README.md +++ b/proposals/2605-optimization-job-crd/README.md @@ -1,6 +1,7 @@ # KEP-3562: OptimizationJob CRD for Hyperparameter Optimization - **Authors:** Aniket Shaha (@aniket2405) +- **Phase 2 (SDK Integration) Author:** Sridhar Pillai (@Sridhar1030) --- @@ -20,6 +21,8 @@ --- + + ## 1. Background & Motivation Historically, Katib has served as Kubeflow’s general-purpose hyperparameter tuning and Neural Architecture Search (NAS) engine. It uses the generic `Experiment` CRD to orchestrate trials, supporting arbitrary Kubernetes workloads via unstructured YAML templates. @@ -77,6 +80,8 @@ spec: - **As an ML Researcher**, I want to consume hyperparameter suggestions via standard environment variables rather than brittle YAML regex string substitution. - **Motivation:** Using the `KUBEFLOW_TRAINER_OPT_` pattern allows me to cleanly parse tuning suggestions inside my Python scripts using existing SDK helper functions without modifying my container's CLI argument parsing logic. [Separate KEP for this integration]. + + ## 3. Goals - **Tighter TrainJob Integration:** Introduce the `OptimizationJob` CRD focused exclusively on `TrainJobs`, using a structured `TrainJobTemplateSpec`. @@ -95,6 +100,8 @@ spec: - Support for Multi-Objective Optimization - Integration w/ Kueue with `suspend` and `managedBy` APIs in OptimizationJob + + ## 4. Non-Goals / Future Iterations To ensure a stable and reviewable initial release (Phase 1), the following features are explicitly out of scope for now and will be addressed in future iterations: @@ -118,6 +125,8 @@ To ensure a stable and reviewable initial release (Phase 1), the following featu - Support for multiple providers for the same algorithm. - Integration with the legacy Katib UI. + + ## 5. Phase 1 API Design (v1alpha1) The MVP API surface is strongly typed to ensure native API server validation via OpenAPI schemas and CEL rules. Mathematical parameters like standard deviations and interval boundaries utilize `string` types to prevent float precision rounding, protected by K8s CEL type-casting. @@ -325,6 +334,8 @@ type Result struct { } ``` + + ## 6. Sample YAML (Phase 1) The `TrainJobTemplate` utilizes a structured approach. Hyperparameters are dynamically injected by the controller directly into the Pod as prefixed environment variables (e.g., `KUBEFLOW_TRAINER_OPT_`) and appended as annotations on the `TrainJob` metadata. @@ -384,12 +395,18 @@ status: value: "32" ``` + + ## 7. Reconciliation & Architecture (Phase 1) + + ### 7.1 Prerequisites - TrainJob Feature Flag (Hard Dependency): The unified `TrainJob` API feature flag MUST be enabled in the cluster/controller environment. + + ### 7.2. gRPC API Strategy & Adapter Pattern To accelerate the MVP and reduce risk, the evolution of the gRPC contract between the Go controller and the Python suggestion engines is divided into two phases: @@ -436,8 +453,12 @@ The core successful conditions for the Phase 1 MVP are: - **Complete**: The conditions of the `TrialPolicy` have been satisfied (e.g., the desired `NumTrials` have successfully finished) and the best result has been recorded. - **Failed**: The `OptimizationJob` encountered a terminal error preventing further execution (e.g., the backend suggestion gRPC service crashed). + + ## 8. Design Decisions & Open Discussions + + ### 8.1. Decision: Decoupling the gRPC Contract **Status: Deferred to Phase 2** @@ -451,6 +472,8 @@ We have deprecated string templating (`{{.param}}`). To pass parameters to the t - **The Design:** The controller injects `KUBEFLOW_TRAINER_OPT_` as environment variables directly into the `trainJob.spec.trainer.env` array. It simultaneously stores the raw JSON parameter assignment as an Annotation on the TrainJob metadata. - **The "Why":** This aligns well with the unified Kubeflow Python SDK (KEP-46). Data scientists can use SDK helper functions (e.g., `get_hyperparameters()`) to cleanly parse the environment variables inside their training scripts without modifying YAML command arguments. The metadata annotations allow the controller to reconstruct trial history purely from the Kubernetes API without requiring Katib DB. + + ### 8.3. Decision: Explicit Separation of Search vs. Pruning **Status: Resolved (Phase 2 Roadmap)** @@ -477,42 +500,68 @@ Pruning decisions are computed controller-side based on this monotonic metric hi ## 9. Phase 2: Kubeflow SDK Integration -This section defines the Python SDK (https://github.com/kubeflow/sdk) client APIs for the +This section defines the Python SDK ([https://github.com/kubeflow/sdk](https://github.com/kubeflow/sdk)) client APIs for the `OptimizationJob` CRD. It replaces the "[Separate KEP for this integration]" placeholder from Story 4 and follows the direction agreed on the Kubeflow SDK & ML Experience call (2026-07-29): -- **Keep the single `TrainerClient()`.** Data scientists should not switch between disparate - clients to train and to tune. Training and hyperparameter optimization are parts of the same - workflow. +- **Keep the single** `TrainerClient()`**.** Data scientists should not switch between disparate +clients to train and to tune. Training and hyperparameter optimization are parts of the same +workflow. - **APIs must be easy to understand from the ML Engineer's point of view.** The same - `TrainJobTemplate` used for training is reused for optimization, and hyperparameter - consumption inside the training function requires no YAML or CLI argument changes. +`TrainJobTemplate` used for training is reused for optimization, and hyperparameter +consumption inside the training function requires no YAML or CLI argument changes. + + ### 9.1. User-Facing API: `TrainerClient().optimize()` The `optimize()` method is added to the existing `TrainerClient`, alongside `train()`. It accepts the same trial template concepts that `train()` already uses, plus the optimization -configuration types that exist in the SDK today (`kubeflow.optimizer.types`): +configuration types that exist in the SDK today (`kubeflow.optimizer.types`). + +All optimization types are importable from `kubeflow.trainer`, so the entire workflow — +training, tuning, and consuming suggestions — lives in one namespace (the types remain +defined in `kubeflow.optimizer` and are re-exported for backward compatibility). + +The example below intentionally mirrors the Kubeflow SDK README: **one** `TrainJobTemplate` +is defined, used first to train, then reused unchanged to optimize. The training function +consumes hyperparameters through its ordinary `func_args` — no function changes are needed +to move from training to tuning (see 9.2 for the mechanism): ```python -from kubeflow.optimizer import Objective, RandomSearch, Search, TrialConfig -from kubeflow.trainer import CustomTrainer, TrainerClient, TrainJobTemplate +from kubeflow.trainer import ( + CustomTrainer, + Objective, + RandomSearch, + Search, + TrainerClient, + TrainJobTemplate, + TrialConfig, +) -def train_fn(): - from kubeflow.trainer import get_hyperparameters +def train_fn(learning_rate: str, batch_size: str): + from kubeflow.trainer import report_metrics - hps = get_hyperparameters() # {"learning_rate": 0.0021, "batch_size": "32"} + lr = float(learning_rate) + bs = int(batch_size) ... + report_metrics({"val_loss": val_loss}) # feeds the Objective below (see 9.3) +# The same template used by TrainerClient().train(**template). +template = TrainJobTemplate( + runtime="torch-distributed", + trainer=CustomTrainer( + func=train_fn, + func_args={"learning_rate": "0.01", "batch_size": "32"}, + ), +) + client = TrainerClient() job_name = client.optimize( - trial_template=TrainJobTemplate( - trainer=CustomTrainer(func=train_fn), - runtime="torch-distributed", - ), + trial_template=template, search_space={ "learning_rate": Search.loguniform(0.0001, 0.1), "batch_size": Search.choice([16, 32, 64]), @@ -533,25 +582,45 @@ initializers, and custom trainers behave identically for training and for trials Lifecycle methods added to `TrainerClient`, mirroring the existing TrainJob methods: -| Method | Behavior | -| --- | --- | -| `optimize(...) -> str` | Create the OptimizationJob, return its generated name. | -| `get_optimization_job(name)` | Read the CR, map `status.conditions` to `Created`/`Complete`/`Failed`. | -| `list_optimization_jobs()` | List OptimizationJobs in the namespace. | -| `wait_for_optimization_job_status(name, status={...}, timeout, polling_interval, callbacks)` | Poll until a desired status; raise on unexpected `Failed`. | -| `get_best_results(name)` | Read `status.result` (trainJobName + parameter assignments). | -| `get_optimization_job_logs(name, trial_name=None, follow=False)` | Delegate to TrainJob log streaming; defaults to the best trial. | -| `get_optimization_job_events(name)` | Kubernetes events for the OptimizationJob and its TrainJobs. | -| `delete_optimization_job(name)` | Delete the CR. | + +| Method | Behavior | +| -------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------- | +| `optimize(...) -> str` | Create the OptimizationJob, return its generated name. | +| `get_optimization_job(name)` | Read the CR, map `status.conditions` to `Created`/`Complete`/`Failed`. | +| `list_optimization_jobs()` | List OptimizationJobs in the namespace. | +| `wait_for_optimization_job_status(name, status={...}, timeout, polling_interval, callbacks)` | Poll until a desired status; raise on unexpected `Failed`. | +| `get_best_results(name)` | Read `status.result` (trainJobName + parameter assignments). | +| `get_optimization_job_logs(name, trial_name=None, follow=False)` | Delegate to TrainJob log streaming; defaults to the best trial. | +| `get_optimization_job_events(name)` | Kubernetes events for the OptimizationJob and its TrainJobs. | +| `delete_optimization_job(name)` | Delete the CR. | + Because trials are plain `TrainJob` resources, every trial is also individually accessible through the existing `get_job()`, `get_job_logs()`, and `get_job_events()` APIs — no Trial-specific SDK surface is needed. -### 9.2. In-Trial Helper: `get_hyperparameters()` +### 9.2. Consuming Suggestions in the Training Code Per Section 8.2, the controller injects suggestions as `KUBEFLOW_TRAINER_OPT_` -environment variables. The SDK provides the helper promised there: +environment variables. The SDK surfaces them to user code through two mechanisms: + +**1. Automatic** `func_args` **override (CustomTrainer).** The launcher script that the SDK +already generates to marshal `func_args` into the container gains one step: before invoking +the user function, it scans the environment for the `KUBEFLOW_TRAINER_OPT_` prefix and +overrides any matching `func_args` key with the suggested value. This preserves the SDK's +core promise — the *same* training function and the *same* `TrainJobTemplate` work for both +`train()` and `optimize()`: + +- Under `train()`, `train_fn` receives the `func_args` defaults (`learning_rate="0.01"`). +- Under `optimize()`, each trial's `train_fn` receives the per-trial suggestion +(`learning_rate="0.0021"`), with no `${trialParameters.*}`-style template substitution +involved. +- Search-space parameter names are validated client-side against `func_args` keys at +`optimize()` time, so typos fail before submission. + +**2. Explicit helper (custom containers and builtin trainers).** Where the SDK does not +control the entrypoint (`CustomTrainerContainer`, `BuiltinTrainer`, or plain scripts), user +code reads suggestions with the helper promised in Section 8.2: ```python def get_hyperparameters() -> dict[str, int | float | str]: @@ -559,33 +628,87 @@ def get_hyperparameters() -> dict[str, int | float | str]: ``` - Scans the environment for the `KUBEFLOW_TRAINER_OPT_` prefix and strips it, lower-casing - the remaining key back to the parameter name. +the remaining key back to the parameter name. - Values from `uniform`/`logUniform` spaces are coerced with the CRD `type` - (`Int` -> `int`, `Float` -> `float`); `categorical` choices remain strings. +(`Int` -> `int`, `Float` -> `float`); `categorical` choices remain strings. - Returns an empty dict when running outside an optimization trial, so the same training - function works unchanged under `train()`. +code works unchanged under `train()`. + + + +### 9.3. Reporting the Objective Metric: `report_metrics()` -### 9.3. SDK-to-CRD Type Mapping +The suggestion flow above is one half of the trial contract; the objective metric flowing +back is the other. Per Section 8.6, the OptimizationJob controller consumes metrics from the +standardized `TrainJob` status fields defined in +[KEP-2779](https://github.com/kubeflow/trainer/tree/master/proposals/2779-trainjob-progress), +which are populated through the push endpoint whose credentials +(`KUBEFLOW_TRAINER_SERVER_URL` / `_CA_CERT` / `_TOKEN`) the control plane injects into +training pods. -The SDK reuses the existing search-space and algorithm types from KEP-46 so users don't -learn a second vocabulary. Numeric bounds are serialized as strings to satisfy the CRD +KEP-2779 defines the transport and shows raw-`urllib` instrumentation; the SDK completes the +user experience with a thin wrapper, symmetric to `get_hyperparameters()`: + +```python +def report_metrics(metrics: dict[str, float], step: int | None = None) -> None: + """Push training metrics (e.g. the optimization objective) to the TrainJob status.""" +``` + +- Reads the `KUBEFLOW_TRAINER_SERVER_*` environment variables and POSTs the payload to the +KEP-2779 endpoint; implemented with the standard library only. +- No-ops (with a debug log) when the environment variables are absent, so instrumented code +runs unchanged outside the cluster and under plain `train()` without the feature gate. +- The metric name used for the `Objective` (e.g. `val_loss`) must match a key reported here; +`optimize()` documentation makes this contract explicit. +- Framework callbacks planned in KEP-2779 (e.g. a Transformers `KubeflowTrainerCallback`) +are expected to call this same helper, so manually instrumented and callback-instrumented +trials behave identically. + +The transport, authentication, and status schema remain owned by KEP-2779; this section only +fixes the user-facing function signature. + +### 9.4. Backend-Neutral Search Space Types & SDK-to-CRD Mapping + +Today the `Search.uniform()`/`Search.loguniform()`/`Search.choice()` helpers return Katib +generated models (`kubeflow_katib_api.models.V1beta1ParameterSpec`). Reusing those for the +trainer-native path would keep a hard dependency on Katib API models in a flow whose purpose +is to remove Katib coupling. + +Phase 2 therefore makes the `Search` helpers **backend-neutral**: they return the plain SDK +dataclasses that already exist in `kubeflow.optimizer.types.search_types` +(`ContinuousSearchSpace`, `CategoricalSearchSpace`), and each backend owns its own +translation: + +- **Trainer-native backend:** SDK dataclass -> `OptimizationJob` `spec.parameters` entries. +- **Katib backend (during the transition):** SDK dataclass -> `V1beta1ParameterSpec`, i.e. +the conversion that lives inside `Search` today moves into the Katib backend. + +User-facing call sites (`Search.uniform(0.01, 0.1)`) are unchanged; only the return type +becomes SDK-owned. This removes `kubeflow-katib-api` from the trainer-native import path +entirely and lets the Katib dependency drop out with the backend in the final migration step +(9.6). The same pattern applies to algorithms: `RandomSearch`/`GridSearch` dataclasses stay +backend-neutral, and the Katib `_to_katib_spec()` conversion moves behind the Katib backend. + +The trainer-native mapping — numeric bounds serialized as strings to satisfy the CRD `Double` pattern: -| SDK (today) | OptimizationJob CRD | -| --- | --- | -| `Search.uniform(min, max)` | `searchSpace.uniform {min: "...", max: "...", type: Float}` | -| `Search.loguniform(min, max)` | `searchSpace.logUniform {min: "...", max: "...", type: Float}` | -| `Search.choice([...])` | `searchSpace.categorical {choices: [...]}` (values stringified) | -| `Objective(metric, direction="minimize")` | `objectives[0] {metric, direction: Minimize}` | -| `RandomSearch(random_state=N)` | `searchAlgorithm.random {seed: N}` | -| `GridSearch()` | `searchAlgorithm.grid {}` | -| `TrialConfig(num_trials, parallel_trials)` | `numTrials`, `parallelTrials` | + +| SDK type | OptimizationJob CRD | +| ------------------------------------------ | --------------------------------------------------------------- | +| `Search.uniform(min, max)` | `searchSpace.uniform {min: "...", max: "...", type: Float}` | +| `Search.loguniform(min, max)` | `searchSpace.logUniform {min: "...", max: "...", type: Float}` | +| `Search.choice([...])` | `searchSpace.categorical {choices: [...]}` (values stringified) | +| `Objective(metric, direction="minimize")` | `objectives[0] {metric, direction: Minimize}` | +| `RandomSearch(random_state=N)` | `searchAlgorithm.random {seed: N}` | +| `GridSearch()` | `searchAlgorithm.grid {}` | +| `TrialConfig(num_trials, parallel_trials)` | `numTrials`, `parallelTrials` | + Constraints surfaced client-side as `ValueError` before submission, matching CRD validation: exactly one objective (multi-objective is a Phase 2/3 CRD goal), and `TrialConfig.max_failed_trials` is rejected until a `TrialPolicy` exists in the CRD. -### 9.4. Results and Trial History +### 9.5. Results and Trial History Phase 1 status exposes only `status.result` (best trial name + parameter assignments), so `get_best_results()` returns those parameters, and metrics/trial listings are limited until @@ -593,42 +716,50 @@ Phase 1 status exposes only `status.result` (best trial name + parameter assignm the OptimizationJob status. Once KEP-3744 lands, `get_optimization_job(name).trials` and the metrics in `Result` are populated from status without additional API calls. -### 9.5. Migration from the Katib-Backed OptimizerClient +### 9.6. Migration from the Katib-Backed OptimizerClient The existing `OptimizerClient` (KEP-46) maps the same Python types onto Katib `Experiment` CRs. Transition plan: 1. **Introduce** `TrainerClient().optimize()` targeting `OptimizationJob` (this section). - The Katib-backed `OptimizerClient` remains untouched and fully supported. + The Katib-backed `OptimizerClient` remains untouched and fully supported. 2. **Deprecate** `OptimizerClient` with a warning once the OptimizationJob controller is - released and feature parity for Random/Grid search is verified. + released and feature parity for Random/Grid search is verified. 3. **Remove** the Katib backend in a future SDK major/minor release, per the standard - Kubeflow deprecation policy. + Kubeflow deprecation policy. Because both paths share the `Search`/`Objective`/`TrialConfig` types, user migration is limited to swapping `OptimizerClient().optimize(...)` for `TrainerClient().optimize(...)`. -### 9.6. Implementation Notes & Open Questions +### 9.7. Implementation Notes & Open Questions - A working prototype of the CR construction, status mapping, and unit-test approach exists - in [kubeflow/sdk#726](https://github.com/kubeflow/sdk/pull/726) (held pending this design). - It validates the type mappings above against a mocked Kubernetes API; E2E validation is - blocked until the OptimizationJob controller merges. +in [kubeflow/sdk#726](https://github.com/kubeflow/sdk/pull/726) (held pending this design). +It validates the type mappings above against a mocked Kubernetes API; E2E validation is +blocked until the OptimizationJob controller merges. - **Open:** whether `get_optimization_job()`-style names are kept, or the existing - `get_job()`/`wait_for_job_status()` methods gain a job-type abstraction instead. -- **Open:** whether the optimization types (`Search`, `Objective`, `TrialConfig`) are - re-exported from `kubeflow.trainer` so a single import namespace covers the whole workflow. +`get_job()`/`wait_for_job_status()` methods gain a job-type abstraction instead. +- **Open:** convergence with the proposed `TrainJobTemplate.train()`/`.optimize()` +convenience methods ([kubeflow/sdk#347](https://github.com/kubeflow/sdk/pull/347), +[kubeflow/sdk#366](https://github.com/kubeflow/sdk/pull/366)); those would delegate to the +`TrainerClient` APIs defined here, so the two efforts are complementary. - **Open:** whether generated Python models for `trainer.kubeflow.org/v1alpha1` types are - published in `kubeflow-trainer-api` before the SDK integration merges, or the SDK builds - plain dictionaries in the interim. +published in `kubeflow-trainer-api` before the SDK integration merges, or the SDK builds +plain dictionaries in the interim. + + ## 10. Implementation History - **2026-06-01:** Initial KEP draft creation for the `OptimizationJob` CRD. - **2026-08-09:** Added Phase 2 Kubeflow SDK integration design (`TrainerClient().optimize()`). + + ## 11. Alternatives + + ### Extend Existing Katib Experiment and Trial CRDs Instead of introducing the `OptimizationJob` CRD, we could have updated the existing Katib `Experiment` and `Trial` CRDs to support `TrainJob` references. @@ -648,4 +779,4 @@ design unchanged. ### Stateful Sidecars with Persistent Storage Katib currently deploys a stateful `Suggestion` sidecar and a persistent DB layer for every experiment. We could have replicated this architecture for `OptimizationJob`. -Deploying dedicated databases for every hyperparameter tuning job introduces severe cluster resource bloat and operational complexity. By reconstructing the trial history directly from completed `TrainJob` annotations and passing it statelessly to the gRPC provider, we eliminate the need for persistent storage and sidecar lifecycle management. +Deploying dedicated databases for every hyperparameter tuning job introduces severe cluster resource bloat and operational complexity. By reconstructing the trial history directly from completed `TrainJob` annotations and passing it statelessly to the gRPC provider, we eliminate the need for persistent storage and sidecar lifecycle management. \ No newline at end of file From 16a6e0d1f68156c73b4f7621a65cd85df5e1c348 Mon Sep 17 00:00:00 2001 From: Sridhar1030 Date: Sun, 9 Aug 2026 14:25:02 +0530 Subject: [PATCH 3/3] fix end of file newline for pre-commit Signed-off-by: Sridhar1030 --- proposals/2605-optimization-job-crd/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/proposals/2605-optimization-job-crd/README.md b/proposals/2605-optimization-job-crd/README.md index a5afd1212d..68cea4193c 100644 --- a/proposals/2605-optimization-job-crd/README.md +++ b/proposals/2605-optimization-job-crd/README.md @@ -779,4 +779,4 @@ design unchanged. ### Stateful Sidecars with Persistent Storage Katib currently deploys a stateful `Suggestion` sidecar and a persistent DB layer for every experiment. We could have replicated this architecture for `OptimizationJob`. -Deploying dedicated databases for every hyperparameter tuning job introduces severe cluster resource bloat and operational complexity. By reconstructing the trial history directly from completed `TrainJob` annotations and passing it statelessly to the gRPC provider, we eliminate the need for persistent storage and sidecar lifecycle management. \ No newline at end of file +Deploying dedicated databases for every hyperparameter tuning job introduces severe cluster resource bloat and operational complexity. By reconstructing the trial history directly from completed `TrainJob` annotations and passing it statelessly to the gRPC provider, we eliminate the need for persistent storage and sidecar lifecycle management.