feat: add GPU info configuration and TFLOPS-based resource limiting (#84)

* feat: add GPU info configuration and TFLOPS-based resource limiting

* fix lint

Author: 0x5457

charts/tensor-fusion/Chart.yaml

@@ -15,7 +15,7 @@ type: application
 # This is the chart version. This version number should be incremented each time you make changes
 # to the chart and its templates, including the app version.
 # Versions are expected to follow Semantic Versioning (https://semver.org/)
-version: 1.2.13
+version: 1.2.14
 
 # This is the version number of the application being deployed. This version number should be
 # incremented each time you make changes to the application. Versions are not expected to

charts/tensor-fusion/templates/controller-deployment.yaml

@@ -57,6 +57,9 @@ spec:
             - name: cloud-vendor-credentials
               mountPath: /tmp/secret
               readOnly: true
+            - name: gpu-info
+              mountPath: /etc/tensor-fusion
+              readOnly: true
         {{- if .Values.agent.agentId }}
         - name: cluster-agent
           image: "{{ .Values.agent.image.repository }}:{{ .Values.agent.image.tag | default "latest" }}"
@@ -105,6 +108,10 @@ spec:
           secret:
             secretName: tf-cloud-vendor-credentials
             defaultMode: 420
+        - configMap:
+            defaultMode: 420
+            name: {{ .Release.Name }}-public-gpu-info
+          name: gpu-info
         - name: logs
           emptyDir: {}
       {{- with .Values.controller.affinity }}

cmd/main.go

@@ -41,6 +41,7 @@ import (
 
 	tensorfusionaiv1 "github.com/NexusGPU/tensor-fusion/api/v1"
 	tfv1 "github.com/NexusGPU/tensor-fusion/api/v1"
+	"github.com/NexusGPU/tensor-fusion/internal/config"
 	"github.com/NexusGPU/tensor-fusion/internal/controller"
 	"github.com/NexusGPU/tensor-fusion/internal/scheduler"
 	"github.com/NexusGPU/tensor-fusion/internal/server"
@@ -69,6 +70,7 @@ func main() {
 	var secureMetrics bool
 	var enableHTTP2 bool
 	var tlsOpts []func(*tls.Config)
+	var gpuInfoConfig string
 
 	flag.StringVar(&metricsAddr, "metrics-bind-address", "0", "The address the metrics endpoint binds to. "+
 		"Use :8443 for HTTPS or :8080 for HTTP, or leave as 0 to disable the metrics service.")
@@ -80,6 +82,8 @@ func main() {
 		"If set, the metrics endpoint is served securely via HTTPS. Use --metrics-secure=false to use HTTP instead.")
 	flag.BoolVar(&enableHTTP2, "enable-http2", false,
 		"If set, HTTP/2 will be enabled for the metrics and webhook servers")
+	flag.StringVar(&gpuInfoConfig, "gpu-info-config",
+		"/etc/tensor-fusion/gpu-info.yaml", "specify the path to gpuInfoConfig file")
 	opts := zap.Options{
 		Development: true,
 	}
@@ -107,6 +111,12 @@ func main() {
 		TLSOpts: tlsOpts,
 	})
 
+	gpuInfos, err := config.LoadGpuInfoFromFile(gpuInfoConfig)
+	if err != nil {
+		ctrl.Log.Error(err, "unable to read gpuInfoConfig file")
+		gpuInfos = make([]config.GpuInfo, 0)
+	}
+
 	// Metrics endpoint is enabled in 'config/default/kustomization.yaml'. The Metrics options configure the server.
 	// More info:
 	// - https://pkg.go.dev/sigs.k8s.io/[email protected]/pkg/metrics/server
@@ -260,6 +270,7 @@ func main() {
 		Scheme:    mgr.GetScheme(),
 		Scheduler: scheduler,
 		Recorder:  mgr.GetEventRecorderFor("tensorfusionworkload"),
+		GpuInfos:  gpuInfos,
 	}).SetupWithManager(mgr); err != nil {
 		setupLog.Error(err, "unable to create controller", "controller", "TensorFusionWorkload")
 		os.Exit(1)

cmd/nodediscovery/main.go

@@ -109,7 +109,7 @@ func main() {
 
 	allDeviceIDs := make([]string, 0)
 
-	for i := 0; i < count; i++ {
+	for i := range count {
 		device, ret := nvml.DeviceGetHandleByIndex(i)
 		if ret != nvml.SUCCESS {
 			ctrl.Log.Error(errors.New(nvml.ErrorString(ret)), "unable to get device", "index", i)

internal/config/gpu_info.go

@@ -27,3 +27,15 @@ func LoadGpuInfoFromFile(filename string) ([]GpuInfo, error) {
 	}
 	return infos, nil
 }
+
+func MockGpuInfo() []GpuInfo {
+	return []GpuInfo{
+		{
+			Model:         "mock",
+			Vendor:        "mock",
+			CostPerHour:   0.1,
+			Fp16TFlops:    resource.MustParse("1000"),
+			FullModelName: "mock",
+		},
+	}
+}

internal/constants/constants.go

@@ -45,12 +45,13 @@ const (
 	ConnectionNameEnv      = "TENSOR_FUSION_CONNECTION_NAME"
 	ConnectionNamespaceEnv = "TENSOR_FUSION_CONNECTION_NAMESPACE"
 
-	WorkerPortEnv         = "TENSOR_FUSION_WORKER_PORT"
-	WorkerCudaUpLimitEnv  = "TENSOR_FUSION_CUDA_UP_LIMIT"
-	WorkerCudaMemLimitEnv = "TENSOR_FUSION_CUDA_MEM_LIMIT"
-	WorkerPodNameEnv      = "POD_NAME"
-	NamespaceEnv          = "OPERATOR_NAMESPACE"
-	NamespaceDefaultVal   = "tensor-fusion-sys"
+	WorkerPortEnv              = "TENSOR_FUSION_WORKER_PORT"
+	WorkerCudaUpLimitTflopsEnv = "TENSOR_FUSION_CUDA_UP_LIMIT_TFLOPS"
+	WorkerCudaUpLimitEnv       = "TENSOR_FUSION_CUDA_UP_LIMIT"
+	WorkerCudaMemLimitEnv      = "TENSOR_FUSION_CUDA_MEM_LIMIT"
+	WorkerPodNameEnv           = "POD_NAME"
+	NamespaceEnv               = "OPERATOR_NAMESPACE"
+	NamespaceDefaultVal        = "tensor-fusion-sys"
 )
 
 const (

internal/controller/tensorfusionworkload_controller.go

@@ -35,6 +35,7 @@ import (
 
 	tensorfusionaiv1 "github.com/NexusGPU/tensor-fusion/api/v1"
 	tfv1 "github.com/NexusGPU/tensor-fusion/api/v1"
+	"github.com/NexusGPU/tensor-fusion/internal/config"
 	"github.com/NexusGPU/tensor-fusion/internal/constants"
 	"github.com/NexusGPU/tensor-fusion/internal/metrics"
 	scheduler "github.com/NexusGPU/tensor-fusion/internal/scheduler"
@@ -49,6 +50,7 @@ type TensorFusionWorkloadReconciler struct {
 	Scheme    *runtime.Scheme
 	Scheduler scheduler.Scheduler
 	Recorder  record.EventRecorder
+	GpuInfos  []config.GpuInfo
 }
 
 // +kubebuilder:rbac:groups=tensor-fusion.ai,resources=tensorfusionworkloads,verbs=get;list;watch;create;update;patch;delete
@@ -106,7 +108,7 @@ func (r *TensorFusionWorkloadReconciler) Reconcile(ctx context.Context, req ctrl
 	}
 
 	// Create worker generator
-	workerGenerator := &worker.WorkerGenerator{WorkerConfig: pool.Spec.ComponentConfig.Worker}
+	workerGenerator := &worker.WorkerGenerator{WorkerConfig: pool.Spec.ComponentConfig.Worker, GpuInfos: r.GpuInfos}
 
 	podTemplateHash, err := workerGenerator.PodTemplateHash(workload.Spec.Resources.Limits)
 	if err != nil {

internal/controller/tensorfusionworkload_controller_test.go

@@ -33,6 +33,7 @@ import (
 
 	tensorfusionaiv1 "github.com/NexusGPU/tensor-fusion/api/v1"
 	tfv1 "github.com/NexusGPU/tensor-fusion/api/v1"
+	"github.com/NexusGPU/tensor-fusion/internal/config"
 	"github.com/NexusGPU/tensor-fusion/internal/constants"
 	scheduler "github.com/NexusGPU/tensor-fusion/internal/scheduler"
 )
@@ -70,8 +71,9 @@ var _ = Describe("TensorFusionWorkload Controller", func() {
 		}
 		Expect(k8sClient.Create(ctx, gpu)).To(Succeed())
 		gpu.Status = tfv1.GPUStatus{
-			Phase: tfv1.TensorFusionGPUPhaseRunning,
-			UUID:  "mock-gpu",
+			Phase:    tfv1.TensorFusionGPUPhaseRunning,
+			UUID:     "mock-gpu",
+			GPUModel: "mock",
 			NodeSelector: map[string]string{
 				"kubernetes.io/hostname": "mock-node",
 			},
@@ -92,6 +94,7 @@ var _ = Describe("TensorFusionWorkload Controller", func() {
 			Scheme:    k8sClient.Scheme(),
 			Scheduler: scheduler.NewScheduler(k8sClient),
 			Recorder:  record.NewFakeRecorder(3),
+			GpuInfos:  config.MockGpuInfo(),
 		}
 
 		// Clean up any pods from previous tests

internal/worker/worker.go

@@ -4,10 +4,12 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"math"
 	"strconv"
 	"time"
 
 	tfv1 "github.com/NexusGPU/tensor-fusion/api/v1"
+	"github.com/NexusGPU/tensor-fusion/internal/config"
 	"github.com/NexusGPU/tensor-fusion/internal/constants"
 	"github.com/NexusGPU/tensor-fusion/internal/utils"
 	"github.com/samber/lo"
@@ -22,6 +24,7 @@ func init() {
 }
 
 type WorkerGenerator struct {
+	GpuInfos     []config.GpuInfo
 	WorkerConfig *tfv1.WorkerConfig
 }
 
@@ -84,16 +87,25 @@ func (wg *WorkerGenerator) GenerateWorkerPod(
 		SubPathExpr: fmt.Sprintf("${%s}", constants.WorkerPodNameEnv),
 	})
 
+	info, ok := lo.Find(wg.GpuInfos, func(info config.GpuInfo) bool {
+		return info.FullModelName == gpu.Status.GPUModel
+	})
+	if !ok {
+		return nil, "", fmt.Errorf("gpu info(%s) not found", gpu.Status.GPUModel)
+	}
+
 	spec.Containers[0].Env = append(spec.Containers[0].Env, corev1.EnvVar{
 		Name:  "NVIDIA_VISIBLE_DEVICES",
 		Value: gpu.Status.UUID,
 	}, corev1.EnvVar{
 		Name:  constants.WorkerPortEnv,
 		Value: strconv.Itoa(port),
 	}, corev1.EnvVar{
-		Name: constants.WorkerCudaUpLimitEnv,
-		// TODO: convert tflops to percent
-		Value: "100",
+		Name:  constants.WorkerCudaUpLimitTflopsEnv,
+		Value: strconv.FormatInt(info.Fp16TFlops.Value(), 10),
+	}, corev1.EnvVar{
+		Name:  constants.WorkerCudaUpLimitEnv,
+		Value: strconv.FormatInt(int64(math.Ceil(float64(limits.Tflops.Value())/float64(info.Fp16TFlops.Value())*100)), 10),
 	}, corev1.EnvVar{
 		Name: constants.WorkerCudaMemLimitEnv,
 		// bytesize