Add XGBoost plugin unit tests and update framework test registry

Signed-off-by: Krishna-kg732 <2405732@kiit.ac.in>

Author: Krishna-kg732

pkg/runtime/framework/core/framework_test.go

@@ -56,6 +56,7 @@ import (
 	"github.com/kubeflow/trainer/v2/pkg/runtime/framework/plugins/plainml"
 	"github.com/kubeflow/trainer/v2/pkg/runtime/framework/plugins/torch"
 	"github.com/kubeflow/trainer/v2/pkg/runtime/framework/plugins/volcano"
+	"github.com/kubeflow/trainer/v2/pkg/runtime/framework/plugins/xgboost"
 	index "github.com/kubeflow/trainer/v2/pkg/runtime/indexer"
 	testingutil "github.com/kubeflow/trainer/v2/pkg/util/testing"
 )
@@ -84,12 +85,14 @@ func TestNew(t *testing.T) {
 					torch.Name:        &torch.Torch{},
 					jobset.Name:       &jobset.JobSet{},
 					jax.Name:          &jax.Jax{},
+					xgboost.Name:      &xgboost.XGBoost{},
 				},
 				enforceMLPlugins: []framework.EnforceMLPolicyPlugin{
 					&mpi.MPI{},
 					&plainml.PlainML{},
 					&torch.Torch{},
 					&jax.Jax{},
+					&xgboost.XGBoost{},
 				},
 				enforcePodGroupPolicyPlugins: []framework.EnforcePodGroupPolicyPlugin{
 					&coscheduling.CoScheduling{},
@@ -101,6 +104,7 @@ func TestNew(t *testing.T) {
 					&jobset.JobSet{},
 					&volcano.Volcano{},
 					&jax.Jax{},
+					&xgboost.XGBoost{},
 				},
 				watchExtensionPlugins: []framework.WatchExtensionPlugin{
 					&coscheduling.CoScheduling{},
@@ -137,7 +141,7 @@ func TestNew(t *testing.T) {
 	}
 	cmpOpts := []cmp.Option{
 		cmp.AllowUnexported(Framework{}),
-		cmpopts.IgnoreUnexported(coscheduling.CoScheduling{}, volcano.Volcano{}, mpi.MPI{}, plainml.PlainML{}, torch.Torch{}, jobset.JobSet{}),
+		cmpopts.IgnoreUnexported(coscheduling.CoScheduling{}, volcano.Volcano{}, mpi.MPI{}, plainml.PlainML{}, torch.Torch{}, jobset.JobSet{}, xgboost.XGBoost{}),
 		cmpopts.IgnoreFields(coscheduling.CoScheduling{}, "client"),
 		cmpopts.IgnoreFields(volcano.Volcano{}, "client"),
 		cmpopts.IgnoreFields(jobset.JobSet{}, "client", "restMapper", "scheme", "logger"),

pkg/runtime/framework/plugins/xgboost/xgboost.go

@@ -5,7 +5,7 @@ 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
+    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,
@@ -18,32 +18,118 @@ package xgboost
 
 import (
 	"context"
-	///"fmt"
+	"fmt"
 
+	"k8s.io/apimachinery/pkg/util/validation/field"
+	corev1ac "k8s.io/client-go/applyconfigurations/core/v1"
+	"k8s.io/utils/ptr"
 	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
 
 	trainer "github.com/kubeflow/trainer/v2/pkg/apis/trainer/v1alpha1"
+	"github.com/kubeflow/trainer/v2/pkg/apply"
+	"github.com/kubeflow/trainer/v2/pkg/constants"
 	"github.com/kubeflow/trainer/v2/pkg/runtime"
 	"github.com/kubeflow/trainer/v2/pkg/runtime/framework"
 )
 
-// XGBoost implements the EnforceMLPolicyPlugin interface for distributed
-// XGBoost training using Rabit coordination.
 type XGBoost struct{}
 
 var _ framework.EnforceMLPolicyPlugin = (*XGBoost)(nil)
+var _ framework.CustomValidationPlugin = (*XGBoost)(nil)
 
 const Name = "XGBoost"
 
 func New(context.Context, client.Client, client.FieldIndexer) (framework.Plugin, error) {
 	return &XGBoost{}, nil
 }
+
 func (x *XGBoost) Name() string {
 	return Name
 }
 
-// TODO: Inject DMLC_* Rabit environment variables for
-// distributed XGBoost training. See KEP for env var specification.
+func (x *XGBoost) Validate(_ context.Context, runtimeInfo *runtime.Info, _, newObj *trainer.TrainJob) (admission.Warnings, field.ErrorList) {
+	var allErrs field.ErrorList
+	if newObj.Spec.Trainer != nil {
+		specPath := field.NewPath("spec", "trainer", "env")
+		for i, env := range newObj.Spec.Trainer.Env {
+			if constants.XGBoostReservedEnvNames.Has(env.Name) {
+				allErrs = append(allErrs, field.Forbidden(
+					specPath.Index(i),
+					fmt.Sprintf("%s is reserved for the XGBoost runtime", env.Name),
+				))
+			}
+		}
+	}
+	return nil, allErrs
+}
+
 func (x *XGBoost) EnforceMLPolicy(info *runtime.Info, trainJob *trainer.TrainJob) error {
+	// Guard: Return early if XGBoost policy not configured.
+	if info == nil || info.RuntimePolicy.MLPolicySource == nil ||
+		info.RuntimePolicy.MLPolicySource.XGBoost == nil {
+		return nil
+	}
+
+	// Find the trainer PodSet.
+	trainerPS := info.FindPodSetByAncestor(constants.AncestorTrainer)
+
+	// Set the number of nodes from TrainJob if specified.
+	if trainerPS.Count != nil &&
+		trainJob.Spec.Trainer != nil && trainJob.Spec.Trainer.NumNodes != nil {
+		*trainerPS.Count = *trainJob.Spec.Trainer.NumNodes
+	}
+
+	// Find the trainer container and inject environment variables.
+	var trainerContainer *runtime.Container
+	if trainJob.Spec.Trainer != nil {
+		if trainerContainer = info.FindContainerByPodSetAncestorContainerName(
+			constants.AncestorTrainer, constants.Node,
+		); trainerContainer != nil {
+			numNodes := ptr.Deref(ptr.Deref(trainerPS, runtime.PodSet{}).Count, 1)
+
+			// Auto-derive numWorkersPerNode from GPU resources.
+			// GPU training: 1 worker per GPU | CPU training: 1 worker per node.
+			numWorkersPerNode := int32(1)
+			if res := runtime.ExtractResourcePerNodeFromRuntime(info); res != nil {
+				if gpuCount := runtime.GetNumGPUPerNode(res); gpuCount > 0 {
+					numWorkersPerNode = int32(gpuCount)
+				}
+			}
+			totalWorkers := numNodes * numWorkersPerNode
+
+			// Build tracker URI: <trainjob-name>-node-0-0.<trainjob-name>
+			trackerURI := fmt.Sprintf("%s-%s-0-0.%s",
+				trainJob.Name, constants.Node, trainJob.Name)
+
+			// Inject DMLC_* environment variables.
+			apply.UpsertEnvVars(&trainerContainer.Env,
+				// DMLC_TRACKER_URI - DNS name for rank-0 worker running tracker.
+				*corev1ac.EnvVar().
+					WithName(constants.XGBoostEnvTrackerURI).
+					WithValue(trackerURI),
+				// DMLC_TRACKER_PORT - Default tracker port.
+				*corev1ac.EnvVar().
+					WithName(constants.XGBoostEnvTrackerPort).
+					WithValue(fmt.Sprintf("%d", constants.XGBoostDefaultTrackerPort)),
+				// DMLC_TASK_ID - Worker rank from Job completion index.
+				*corev1ac.EnvVar().
+					WithName(constants.XGBoostEnvTaskID).
+					WithValueFrom(corev1ac.EnvVarSource().
+						WithFieldRef(corev1ac.ObjectFieldSelector().
+							WithFieldPath(constants.JobCompletionIndexFieldPath))),
+				// DMLC_NUM_WORKER - Total number of workers.
+				*corev1ac.EnvVar().
+					WithName(constants.XGBoostEnvNumWorker).
+					WithValue(fmt.Sprintf("%d", totalWorkers)),
+			)
+
+			// Add container port for tracker communication.
+			apply.UpsertPort(&trainerContainer.Ports,
+				*corev1ac.ContainerPort().
+					WithContainerPort(constants.XGBoostDefaultTrackerPort))
+		}
+	}
+
 	return nil
 }