feat: Add training progression tracking feature for experimental implementation

Signed-off-by: abhijeet-dhumal <abhijeetdhumal652@gmail.com>

Author: abhijeet-dhumal

Makefile

@@ -254,7 +254,7 @@ deploy-rhoai: ## Deploy operator using RHOAI manifests with kustomize
 		crd/trainingruntimes.trainer.kubeflow.org \
 		crd/trainjobs.trainer.kubeflow.org 2>/dev/null || sleep 5
 	@echo "Applying operator and resources..."
-	@$(KUBECTL) apply --server-side=true -k $(RHOAI_MANIFESTS_DIR)
+	@$(KUBECTL) apply -k $(RHOAI_MANIFESTS_DIR) --server-side=true --force-conflicts
 	@echo "Waiting for deployment to be ready..."
 	@$(KUBECTL) wait --for=condition=available --timeout=300s \
 		deployment/kubeflow-trainer-controller-manager -n $(NAMESPACE) 2>/dev/null || true

cmd/trainer-controller-manager/main.go

@@ -40,6 +40,7 @@ import (
 	trainer "github.com/kubeflow/trainer/v2/pkg/apis/trainer/v1alpha1"
 	"github.com/kubeflow/trainer/v2/pkg/config"
 	"github.com/kubeflow/trainer/v2/pkg/controller"
+	rhaisetup "github.com/kubeflow/trainer/v2/pkg/rhai"
 	"github.com/kubeflow/trainer/v2/pkg/runtime"
 	runtimecore "github.com/kubeflow/trainer/v2/pkg/runtime/core"
 	"github.com/kubeflow/trainer/v2/pkg/util/cert"
@@ -150,9 +151,47 @@ func setupControllers(mgr ctrl.Manager, runtimes map[string]runtime.Runtime, cer
 	<-certsReady
 	setupLog.Info("Certs ready")
 
-	if failedCtrlName, err := controller.SetupControllers(mgr, runtimes, ctrlpkg.Options{}); err != nil {
-		setupLog.Error(err, "Could not create controller", "controller", failedCtrlName)
-		os.Exit(1)
+	// Setup controllers with optional RHAI features (decorator pattern for midstream extensions)
+	rhaiEnabled := os.Getenv("ENABLE_RHAI_FEATURES") == "true"
+	if rhaiEnabled {
+		setupLog.Info("RHAI features enabled")
+
+		// Setup TrainingRuntime controller
+		runtimeRec := controller.NewTrainingRuntimeReconciler(
+			mgr.GetClient(),
+			mgr.GetEventRecorderFor("trainer-trainingruntime-controller"),
+		)
+		if err := runtimeRec.SetupWithManager(mgr, ctrlpkg.Options{}); err != nil {
+			setupLog.Error(err, "Could not create controller", "controller", "TrainingRuntime")
+			os.Exit(1)
+		}
+
+		// Setup ClusterTrainingRuntime controller
+		clRuntimeRec := controller.NewClusterTrainingRuntimeReconciler(
+			mgr.GetClient(),
+			mgr.GetEventRecorderFor("trainer-clustertrainingruntime-controller"),
+		)
+		if err := clRuntimeRec.SetupWithManager(mgr, ctrlpkg.Options{}); err != nil {
+			setupLog.Error(err, "Could not create controller", "controller", "ClusterTrainingRuntime")
+			os.Exit(1)
+		}
+
+		// Wrap base TrainJob reconciler with RHAI progression tracking
+		baseReconciler := controller.NewTrainJobReconciler(
+			mgr.GetClient(),
+			mgr.GetEventRecorderFor("trainer-trainjob-controller"),
+			runtimes,
+			controller.WithWatchers(runtimeRec, clRuntimeRec),
+		)
+		if err := rhaisetup.SetupWithManager(mgr, baseReconciler); err != nil {
+			setupLog.Error(err, "Could not setup RHAI features")
+			os.Exit(1)
+		}
+	} else {
+		if failedCtrlName, err := controller.SetupControllers(mgr, runtimes, ctrlpkg.Options{}); err != nil {
+			setupLog.Error(err, "Could not create controller", "controller", failedCtrlName)
+			os.Exit(1)
+		}
 	}
 	if failedWebhook, err := webhooks.Setup(mgr, runtimes); err != nil {
 		setupLog.Error(err, "Could not create webhook", "webhook", failedWebhook)

manifests/rhoai/kustomization.yaml

@@ -54,6 +54,13 @@ patches:
 # through a ComponentConfig type
 - path: manager_config_patch.yaml
 - path: manager_metrics_patch.yaml
+# RHAI-specific RBAC for progression tracking
+- path: rbac_progression_patch.yaml
+  target:
+    group: rbac.authorization.k8s.io
+    version: v1
+    kind: ClusterRole
+    name: kubeflow-trainer-controller-manager
 - patch: |-
     - op: remove
       path: /spec/ports/0

manifests/rhoai/manager_config_patch.yaml

@@ -8,6 +8,9 @@ spec:
       containers:
       - name: manager
         image: $(image)
+        env:
+        - name: ENABLE_RHAI_FEATURES
+          value: "true"
         args:
         - --config=/controller_manager_config.yaml
         - --zap-log-level=2

manifests/rhoai/params.env

@@ -1 +1 @@
-odh-kubeflow-trainer-controller-image=quay.io/opendatahub/trainer:v2.1.0
+odh-kubeflow-trainer-controller-image=quay.io/abdhumal/trainer:v2.1.0-rhai-progression-18Nov-2

manifests/rhoai/rbac_progression_patch.yaml

@@ -0,0 +1,26 @@
+# RHAI-specific: Permissions for progression tracking
+# Allows the controller to read pod IPs and execute commands for metrics polling
+- op: add
+  path: /rules/-
+  value:
+    apiGroups:
+    - ""
+    resources:
+    - pods
+    verbs:
+    - create
+    - delete
+    - get
+    - list
+    - patch
+    - update
+    - watch
+- op: add
+  path: /rules/-
+  value:
+    apiGroups:
+    - ""
+    resources:
+    - pods/exec
+    verbs:
+    - create

pkg/rhai/README.md

@@ -0,0 +1,126 @@
+# RHAI (Red Hat AI) Extensions
+
+This directory contains RHAI-specific extensions for the Kubeflow Trainer operator.
+
+## Purpose
+
+The `rhai/` package provides midstream-specific features that are not part of upstream Kubeflow:
+- **Progression tracking**: Real-time training metrics polling and status updates
+- **Custom annotations**: RHAI-specific metadata for training jobs
+- **Extended RBAC**: Additional permissions
+
+## Structure
+
+```
+pkg/rhai/
+├── README.md                       # This file
+├── setup.go                        # RHAI feature registration
+├── controller/
+│   └── progression_controller.go  # Wraps base controller with progression tracking
+└── progression/
+    ├── progression.go              # Core progression tracking logic
+    └── progression_test.go         # Tests for progression tracking
+```
+
+## How It Works
+
+### 1. Controller Wrapping
+
+The `ProgressionReconciler` wraps the base `TrainJobReconciler` and adds:
+- Metrics polling from training pods
+- Progress annotation updates
+- Automatic requeuing for ongoing polling
+
+### 2. Progression Tracking
+
+When enabled via annotation `trainer.opendatahub.io/progression-tracking: "enabled"`:
+- Controller polls training pod's metrics endpoint (default port: 28080)
+- Updates TrainJob annotations with real-time progress
+- Captures final metrics on job completion/failure
+
+### 3. Manifest Integration
+
+RHAI-specific manifests in `manifests/rhoai/`:
+- `rbac_progression_patch.yaml`: Additional RBAC for pod access
+- `manager_config_patch.yaml`: ConfigMap mounting for feature flags
+
+## Enabling RHAI Features
+
+RHAI features are controlled via the `ENABLE_RHAI_FEATURES` environment variable. When enabled, the operator uses a wrapping controller that adds progression tracking to the base upstream functionality.
+
+### Deployment Configuration
+
+**For Kubernetes/OpenShift deployments**, set the environment variable in your deployment manifest:
+
+```yaml
+# manifests/rhoai/manager_config_patch.yaml
+apiVersion: apps/v1
+kind: Deployment
+spec:
+  template:
+    spec:
+      containers:
+      - name: manager
+        env:
+        - name: ENABLE_RHAI_FEATURES
+          value: "true"
+```
+
+**For local development**, export the variable before running:
+
+```bash
+export ENABLE_RHAI_FEATURES=true
+go run ./cmd/trainer-controller-manager/main.go
+```
+
+## Usage Example
+
+Create a TrainJob with progression tracking:
+
+```yaml
+apiVersion: trainer.kubeflow.org/v1alpha1
+kind: TrainJob
+metadata:
+  name: pytorch-example
+  annotations:
+    trainer.opendatahub.io/progression-tracking: "enabled"
+    trainer.opendatahub.io/metrics-port: "28080"           # optional, default: 28080
+    trainer.opendatahub.io/metrics-poll-interval: "30s"   # optional, default: 30s
+spec:
+  # ... your training job spec ...
+```
+
+The controller will:
+1. Poll the primary pod's metrics endpoint every 30s(default) - configurable
+2. Update the `trainer.opendatahub.io/trainerStatus` annotation with:
+   - Progress percentage
+   - Current step/epoch
+   - Loss and learning rate
+   - Time elapsed/remaining
+   - Custom metrics
+3. Capture final status when job completes
+
+## Development
+
+### Running Tests
+
+```bash
+go test ./pkg/rhai/...
+```
+
+### Adding New RHAI Features
+
+1. Create new package under `pkg/rhai/yourfeature/`
+2. Add controller wrapper if needed in `pkg/rhai/controller/`
+3. Update `pkg/rhai/setup.go` to register the feature
+4. Add manifest patches in `manifests/rhoai/`
+5. Document in this README
+
+## Maintenance
+
+When rebasing from upstream:
+1. Pull upstream changes: `git pull upstream master`
+2. Rebase: `git rebase upstream/master`
+3. `pkg/rhai/` should auto-merge (no conflicts expected)
+4. Review controller integration in `main.go` if base controller changed
+5. Run tests: `make test`

pkg/rhai/constants/constants.go

@@ -0,0 +1,53 @@
+/*
+Copyright 2024 The Kubeflow Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+// Package constants contains shared constants for all RHAI features.
+package constants
+
+const (
+	// Progression tracking feature annotations
+
+	// AnnotationProgressionTracking enables/disables progression tracking for a TrainJob.
+	// Value: "enabled" to enable tracking, any other value or absence disables it.
+	// Example: trainer.opendatahub.io/progression-tracking: "enabled"
+	AnnotationProgressionTracking string = "trainer.opendatahub.io/progression-tracking"
+
+	// AnnotationTrainerStatus stores the JSON-encoded training status/progress.
+	// This annotation is automatically updated by the controller with real-time metrics.
+	// Example: trainer.opendatahub.io/trainerStatus: '{"status":"training","progress":{"percent":45.2},...}'
+	AnnotationTrainerStatus string = "trainer.opendatahub.io/trainerStatus"
+
+	// AnnotationMetricsPort specifies the port where the training pod exposes metrics.
+	// Default: 28080
+	// Example: trainer.opendatahub.io/metrics-port: "8080"
+	AnnotationMetricsPort string = "trainer.opendatahub.io/metrics-port"
+
+	// AnnotationMetricsPollInterval specifies how often to poll metrics (supports duration format).
+	// Accepts: "30s", "1m", or integer seconds "30" (min: 5s, max: 300s)
+	// Default: 30s
+	// Example: trainer.opendatahub.io/metrics-poll-interval: "45s"
+	AnnotationMetricsPollInterval string = "trainer.opendatahub.io/metrics-poll-interval"
+
+	// AnnotationFramework specifies the training framework (for framework-specific handling).
+	// Example: trainer.opendatahub.io/framework: "pytorch"
+	AnnotationFramework string = "trainer.opendatahub.io/framework"
+
+	// DefaultMetricsPort is the default port for metrics endpoints in training pods.
+	DefaultMetricsPort string = "28080"
+
+	// DefaultMetricsPollIntervalSecs is the default interval (in seconds) for polling training metrics.
+	DefaultMetricsPollIntervalSecs int = 30
+)