Merge pull request #732 from rancher/add-workload-services

Add workload services
rancher · Nov 29, 2022 · 6cf1d1f · 6cf1d1f
2 parents ecd998e + 20d11f5
commit 6cf1d1f
Show file tree

Hide file tree

Showing 20 changed files with 580 additions and 166,189 deletions.
diff --git a/apis/ai/v1beta1/opnicluster_meta.go b/apis/ai/v1beta1/opnicluster_meta.go
@@ -14,6 +14,7 @@ const (
 	PayloadReceiverService
 	GPUControllerService
 	MetricsService
+	TrainingControllerService
 	OpensearchUpdateService
 )
 
@@ -42,6 +43,8 @@ func (s ServiceKind) String() string {
 		return "metrics"
 	case OpensearchUpdateService:
 		return "opensearch-update"
+	case TrainingControllerService:
+		return "training-controller"
 	default:
 		return ""
 	}
@@ -76,6 +79,8 @@ func (s ServiceKind) GetImageSpec(opniCluster *OpniCluster) *opnimeta.ImageSpec
 		return &opniCluster.Spec.Services.Metrics.ImageSpec
 	case OpensearchUpdateService:
 		return &opniCluster.Spec.Services.OpensearchUpdate.ImageSpec
+	case TrainingControllerService:
+		return &opniCluster.Spec.Services.TrainingController.ImageSpec
 	default:
 		return nil
 	}
@@ -97,6 +102,8 @@ func (s ServiceKind) GetNodeSelector(opniCluster *OpniCluster) map[string]string
 		return opniCluster.Spec.Services.Metrics.NodeSelector
 	case OpensearchUpdateService:
 		return opniCluster.Spec.Services.OpensearchUpdate.NodeSelector
+	case TrainingControllerService:
+		return opniCluster.Spec.Services.TrainingController.NodeSelector
 	default:
 		return map[string]string{}
 	}
@@ -118,6 +125,8 @@ func (s ServiceKind) GetTolerations(opniCluster *OpniCluster) []corev1.Toleratio
 		return opniCluster.Spec.Services.Metrics.Tolerations
 	case OpensearchUpdateService:
 		return opniCluster.Spec.Services.OpensearchUpdate.Tolerations
+	case TrainingControllerService:
+		return opniCluster.Spec.Services.TrainingController.Tolerations
 	default:
 		return []corev1.Toleration{}
 	}

diff --git a/apis/ai/v1beta1/opnicluster_types.go b/apis/ai/v1beta1/opnicluster_types.go
@@ -87,21 +87,28 @@ type OpniCluster struct {
 }
 
 type ServicesSpec struct {
-	Drain            DrainServiceSpec            `json:"drain,omitempty"`
-	Inference        InferenceServiceSpec        `json:"inference,omitempty"`
-	Preprocessing    PreprocessingServiceSpec    `json:"preprocessing,omitempty"`
-	PayloadReceiver  PayloadReceiverServiceSpec  `json:"payloadReceiver,omitempty"`
-	GPUController    GPUControllerServiceSpec    `json:"gpuController,omitempty"`
-	Metrics          MetricsServiceSpec          `json:"metrics,omitempty"`
-	OpensearchUpdate OpensearchUpdateServiceSpec `json:"opensearchUpdate,omitempty"`
+	Drain              DrainServiceSpec              `json:"drain,omitempty"`
+	Inference          InferenceServiceSpec          `json:"inference,omitempty"`
+	Preprocessing      PreprocessingServiceSpec      `json:"preprocessing,omitempty"`
+	PayloadReceiver    PayloadReceiverServiceSpec    `json:"payloadReceiver,omitempty"`
+	GPUController      GPUControllerServiceSpec      `json:"gpuController,omitempty"`
+	Metrics            MetricsServiceSpec            `json:"metrics,omitempty"`
+	OpensearchUpdate   OpensearchUpdateServiceSpec   `json:"opensearchUpdate,omitempty"`
+	TrainingController TrainingControllerServiceSpec `json:"trainingController,omitempty"`
 }
 
 type DrainServiceSpec struct {
 	opnimeta.ImageSpec `json:",inline,omitempty"`
-	Enabled            *bool               `json:"enabled,omitempty"`
-	NodeSelector       map[string]string   `json:"nodeSelector,omitempty"`
-	Tolerations        []corev1.Toleration `json:"tolerations,omitempty"`
-	Replicas           *int32              `json:"replicas,omitempty"`
+	Enabled            *bool                    `json:"enabled,omitempty"`
+	NodeSelector       map[string]string        `json:"nodeSelector,omitempty"`
+	Tolerations        []corev1.Toleration      `json:"tolerations,omitempty"`
+	Replicas           *int32                   `json:"replicas,omitempty"`
+	Workload           WorkloadDrainServiceSpec `json:"workload,omitempty"`
+}
+
+type WorkloadDrainServiceSpec struct {
+	Enabled  *bool  `json:"enabled,omitempty"`
+	Replicas *int32 `json:"replicas,omitempty"`
 }
 
 type InferenceServiceSpec struct {
@@ -135,6 +142,13 @@ type GPUControllerServiceSpec struct {
 	Tolerations        []corev1.Toleration `json:"tolerations,omitempty"`
 }
 
+type TrainingControllerServiceSpec struct {
+	opnimeta.ImageSpec `json:",inline,omitempty"`
+	Enabled            *bool               `json:"enabled,omitempty"`
+	NodeSelector       map[string]string   `json:"nodeSelector,omitempty"`
+	Tolerations        []corev1.Toleration `json:"tolerations,omitempty"`
+}
+
 type MetricsServiceSpec struct {
 	opnimeta.ImageSpec  `json:",inline,omitempty"`
 	Enabled             *bool                         `json:"enabled,omitempty"`

diff --git a/apis/ai/v1beta1/zz_generated.deepcopy.go b/apis/ai/v1beta1/zz_generated.deepcopy.go
diff --git a/apis/v1beta2/opnicluster_meta.go b/apis/v1beta2/opnicluster_meta.go
@@ -15,6 +15,7 @@ const (
 	GPUControllerService
 	MetricsService
 	OpensearchUpdateService
+	TrainingControllerService
 )
 
 type OpensearchRole string
@@ -42,6 +43,8 @@ func (s ServiceKind) String() string {
 		return "metrics"
 	case OpensearchUpdateService:
 		return "opensearch-update"
+	case TrainingControllerService:
+		return "training-controller"
 	default:
 		return ""
 	}
@@ -76,6 +79,8 @@ func (s ServiceKind) GetImageSpec(spec OpniClusterSpec) *opnimeta.ImageSpec {
 		return &spec.Services.Metrics.ImageSpec
 	case OpensearchUpdateService:
 		return &spec.Services.OpensearchUpdate.ImageSpec
+	case TrainingControllerService:
+		return &spec.Services.TrainingController.ImageSpec
 	default:
 		return nil
 	}
@@ -97,6 +102,8 @@ func (s ServiceKind) GetNodeSelector(spec OpniClusterSpec) map[string]string {
 		return spec.Services.Metrics.NodeSelector
 	case OpensearchUpdateService:
 		return spec.Services.OpensearchUpdate.NodeSelector
+	case TrainingControllerService:
+		return spec.Services.TrainingController.NodeSelector
 	default:
 		return map[string]string{}
 	}
@@ -118,6 +125,8 @@ func (s ServiceKind) GetTolerations(spec OpniClusterSpec) []corev1.Toleration {
 		return spec.Services.Metrics.Tolerations
 	case OpensearchUpdateService:
 		return spec.Services.OpensearchUpdate.Tolerations
+	case TrainingControllerService:
+		return spec.Services.TrainingController.Tolerations
 	default:
 		return []corev1.Toleration{}
 	}

diff --git a/apis/v1beta2/opnicluster_types.go b/apis/v1beta2/opnicluster_types.go
@@ -100,13 +100,14 @@ type OpniCluster struct {
 }
 
 type ServicesSpec struct {
-	Drain            DrainServiceSpec            `json:"drain,omitempty"`
-	Inference        InferenceServiceSpec        `json:"inference,omitempty"`
-	Preprocessing    PreprocessingServiceSpec    `json:"preprocessing,omitempty"`
-	PayloadReceiver  PayloadReceiverServiceSpec  `json:"payloadReceiver,omitempty"`
-	GPUController    GPUControllerServiceSpec    `json:"gpuController,omitempty"`
-	Metrics          MetricsServiceSpec          `json:"metrics,omitempty"`
-	OpensearchUpdate OpensearchUpdateServiceSpec `json:"opensearchUpdate,omitempty"`
+	Drain              DrainServiceSpec              `json:"drain,omitempty"`
+	Inference          InferenceServiceSpec          `json:"inference,omitempty"`
+	Preprocessing      PreprocessingServiceSpec      `json:"preprocessing,omitempty"`
+	PayloadReceiver    PayloadReceiverServiceSpec    `json:"payloadReceiver,omitempty"`
+	GPUController      GPUControllerServiceSpec      `json:"gpuController,omitempty"`
+	Metrics            MetricsServiceSpec            `json:"metrics,omitempty"`
+	OpensearchUpdate   OpensearchUpdateServiceSpec   `json:"opensearchUpdate,omitempty"`
+	TrainingController TrainingControllerServiceSpec `json:"trainingController,omitempty"`
 }
 
 type DrainServiceSpec struct {
@@ -140,6 +141,13 @@ type PayloadReceiverServiceSpec struct {
 	Tolerations        []corev1.Toleration `json:"tolerations,omitempty"`
 }
 
+type TrainingControllerServiceSpec struct {
+	opnimeta.ImageSpec `json:",inline,omitempty"`
+	Enabled            *bool               `json:"enabled,omitempty"`
+	NodeSelector       map[string]string   `json:"nodeSelector,omitempty"`
+	Tolerations        []corev1.Toleration `json:"tolerations,omitempty"`
+}
+
 type GPUControllerServiceSpec struct {
 	opnimeta.ImageSpec `json:",inline,omitempty"`
 	Enabled            *bool               `json:"enabled,omitempty"`

diff --git a/apis/v1beta2/zz_generated.deepcopy.go b/apis/v1beta2/zz_generated.deepcopy.go
diff --git a/config/crd/bases/ai.opni.io_opniclusters.yaml b/config/crd/bases/ai.opni.io_opniclusters.yaml