fix: add limiter for worker (#61)

Author: 0x5457

internal/constants/constants.go

@@ -44,9 +44,11 @@ const (
 	ConnectionNameEnv      = "TENSOR_FUSION_CONNECTION_NAME"
 	ConnectionNamespaceEnv = "TENSOR_FUSION_CONNECTION_NAMESPACE"
 
-	WorkerPortEnv       = "TENSOR_FUSION_WORKER_PORT"
-	NamespaceEnv        = "OPERATOR_NAMESPACE"
-	NamespaceDefaultVal = "tensor-fusion"
+	WorkerPortEnv        = "TENSOR_FUSION_WORKER_PORT"
+	WokerCudaUpLimitEnv  = "TENSOR_FUSION_CUDA_UP_LIMIT"
+	WokerCudaMemLimitEnv = "TENSOR_FUSION_CUDA_MEM_LIMIT"
+	NamespaceEnv         = "OPERATOR_NAMESPACE"
+	NamespaceDefaultVal  = "tensor-fusion"
 )
 
 const (

internal/controller/tensorfusionworkload_controller.go

@@ -165,7 +165,7 @@ func (r *TensorFusionWorkloadReconciler) tryStartWorker(
 		if errors.IsNotFound(err) {
 			// Pod doesn't exist, create a new one
 			port := workerGenerator.AllocPort()
-			pod, err = workerGenerator.GenerateWorkerPod(gpu, namespacedName, port)
+			pod, err = workerGenerator.GenerateWorkerPod(gpu, namespacedName, port, workload.Spec.Resources.Limits)
 			if err != nil {
 				return nil, fmt.Errorf("generate worker pod %w", err)
 			}

internal/worker/worker.go

@@ -48,6 +48,7 @@ func (wg *WorkerGenerator) GenerateWorkerPod(
 	gpu *tfv1.GPU,
 	namespacedName types.NamespacedName,
 	port int,
+	limits tfv1.Resource,
 ) (*corev1.Pod, error) {
 	podTmpl := &corev1.PodTemplate{}
 	err := json.Unmarshal(wg.WorkerConfig.PodTemplate.Raw, podTmpl)
@@ -78,6 +79,14 @@ func (wg *WorkerGenerator) GenerateWorkerPod(
 	}, corev1.EnvVar{
 		Name:  constants.WorkerPortEnv,
 		Value: strconv.Itoa(port),
+	}, corev1.EnvVar{
+		Name: constants.WokerCudaUpLimitEnv,
+		// TODO: convert tflops to percent
+		Value: "100",
+	}, corev1.EnvVar{
+		Name: constants.WokerCudaMemLimitEnv,
+		// bytesize
+		Value: strconv.FormatInt(limits.Vram.Value(), 10),
 	})
 
 	return &corev1.Pod{