feat: get tflops from config file

Author: 0x5457

api/v1/gpu_types.go

@@ -26,8 +26,8 @@ type GPUStatus struct {
 	// +kubebuilder:default=Pending
 	Phase TensorFusionGPUPhase `json:"phase"`
 
-	Capacity  Resource `json:"capacity"`
-	Available Resource `json:"available"`
+	Capacity  *Resource `json:"capacity"`
+	Available *Resource `json:"available"`
 
 	UUID string `json:"uuid"`
 

api/v1/zz_generated.deepcopy.go

@@ -9,7 +9,9 @@ import (
 
 	"github.com/NVIDIA/go-nvml/pkg/nvml"
 	tfv1 "github.com/NexusGPU/tensor-fusion-operator/api/v1"
+	"github.com/NexusGPU/tensor-fusion-operator/internal/config"
 	"github.com/NexusGPU/tensor-fusion-operator/internal/reporter"
+	"github.com/samber/lo"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	ctrl "sigs.k8s.io/controller-runtime"
@@ -19,9 +21,11 @@ import (
 func main() {
 	var dryRun bool
 	var hostname string
-
+	var gpuInfoConfig string
 	flag.BoolVar(&dryRun, "dry-run", false, "dry run mode")
 	flag.StringVar(&hostname, "hostname", "", "hostname")
+	flag.StringVar(&gpuInfoConfig, "gpu-info-config", "", "specify the path to gpuInfoConfig file")
+
 	if hostname == "" {
 		hostname = os.Getenv("HOSTNAME")
 	}
@@ -34,6 +38,12 @@ func main() {
 	flag.Parse()
 	ctrl.SetLogger(zap.New(zap.UseFlagOptions(&opts)))
 
+	gpuinfos, err := config.LoadGpuInfoFromFile(gpuInfoConfig)
+	if err != nil {
+		ctrl.Log.Error(err, "unable to read gpuInfoConfig file")
+		os.Exit(1)
+	}
+
 	ret := nvml.Init()
 	if ret != nvml.SUCCESS {
 		ctrl.Log.Error(errors.New(nvml.ErrorString(ret)), "unable to initialize NVML")
@@ -90,15 +100,21 @@ func main() {
 			ctrl.Log.Error(errors.New(nvml.ErrorString(ret)), "unable to get memory info of device", "index", i)
 			os.Exit(1)
 		}
+		info, ok := lo.Find(gpuinfos, func(info config.GpuInfo) bool {
+			return info.FullModelName == deviceName
+		})
+		tflops := info.Fp16TFlops
+		if !ok {
+			tflops = resource.MustParse("0")
+		}
 		gpu := &tfv1.GPU{
 			ObjectMeta: metav1.ObjectMeta{
 				Name: uuid,
 			},
 			Status: tfv1.GPUStatus{
-				Capacity: tfv1.Resource{
-					Vram: resource.MustParse(fmt.Sprintf("%dKi", memInfo.Total)),
-					// TODO: compute Tflops based on GPU model
-					Tflops: resource.MustParse("100"),
+				Capacity: &tfv1.Resource{
+					Vram:   resource.MustParse(fmt.Sprintf("%dKi", memInfo.Total)),
+					Tflops: tflops,
 				},
 				UUID:     uuid,
 				GPUModel: deviceName,
@@ -112,7 +128,11 @@ func main() {
 			// keep Available field
 			available := gpu.Status.Available
 			gpu.Status = gpuCopy.Status
-			gpu.Status.Available = available
+			if available != nil {
+				gpu.Status.Available = available
+			} else {
+				gpu.Status.Available = gpu.Status.Capacity
+			}
 			return nil
 		}); err != nil {
 			ctrl.Log.Error(err, "failed to report GPU", "gpu", gpu)

cmd/nodediscovery/main.go

@@ -0,0 +1,29 @@
+package config
+
+import (
+	"os"
+
+	"k8s.io/apimachinery/pkg/api/resource"
+	"sigs.k8s.io/yaml"
+)
+
+type GpuInfo struct {
+	Model         string            `json:"model"`
+	Vendor        string            `json:"vendor"`
+	CostPerHour   float64           `json:"costPerHour"`
+	Fp16TFlops    resource.Quantity `json:"fp16TFlops"`
+	FullModelName string            `json:"fullModelName"`
+}
+
+func LoadGpuInfoFromFile(filename string) ([]GpuInfo, error) {
+	infos := make([]GpuInfo, 0)
+	data, err := os.ReadFile(filename)
+	if err != nil {
+		return infos, err
+	}
+	err = yaml.Unmarshal(data, &infos)
+	if err != nil {
+		return nil, err
+	}
+	return infos, nil
+}

internal/config/gpu_info.go

@@ -12,8 +12,8 @@ const (
 
 	LabelKeyOwner = Domain + "/managed-by"
 
-	GPUNodePoolIdentifierLabelPrefix = Domain + "/pool/"
-	GPUNodePoolIdentifierLabelFormat = Domain + "/pool/%s"
+	GPUNodePoolIdentifierLabelPrefix = Domain + "/pool-"
+	GPUNodePoolIdentifierLabelFormat = Domain + "/pool-%s"
 	NodeDeletionMark                 = Domain + "/should-delete"
 
 	TensorFusionEnabledLabelKey = Domain + "/enabled"

internal/constants/constants.go

@@ -86,7 +86,9 @@ func (r *GPUPoolReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ct
 	}
 
 	// For provisioning mode, check if need to scale up GPUNodes upon AvailableCapacity changed
-	isProvisioningMode := pool.Spec.NodeManagerConfig.NodeProvisioner != nil && pool.Spec.NodeManagerConfig.NodeProvisioner.Mode == tfv1.NodeProvisionerModeNative
+	isProvisioningMode := pool.Spec.NodeManagerConfig != nil &&
+		pool.Spec.NodeManagerConfig.NodeProvisioner != nil &&
+		pool.Spec.NodeManagerConfig.NodeProvisioner.Mode == tfv1.NodeProvisionerModeNative
 	if isProvisioningMode {
 		if err := r.reconcilePoolCapacityWithProvisioner(ctx, pool); err != nil {
 			return ctrl.Result{}, err

internal/controller/gpupool_controller.go

@@ -93,7 +93,7 @@ func (r *NodeReconciler) generateGPUNode(ctx context.Context, node *corev1.Node,
 		ObjectMeta: metav1.ObjectMeta{
 			Name: node.Name,
 			Labels: map[string]string{
-				fmt.Sprint(constants.GPUNodePoolIdentifierLabelFormat, poolName): "true",
+				fmt.Sprintf(constants.GPUNodePoolIdentifierLabelFormat, poolName): "true",
 			},
 		},
 		Spec: tfv1.GPUNodeSpec{

internal/controller/node_controller.go

@@ -86,11 +86,11 @@ var _ = Describe("TensorFusionConnection Controller", func() {
 				NodeSelector: map[string]string{
 					"kubernetes.io/hostname": "mock-node",
 				},
-				Capacity: tfv1.Resource{
+				Capacity: &tfv1.Resource{
 					Tflops: resource.MustParse("2"),
 					Vram:   resource.MustParse("2Gi"),
 				},
-				Available: tfv1.Resource{
+				Available: &tfv1.Resource{
 					Tflops: resource.MustParse("2"),
 					Vram:   resource.MustParse("2Gi"),
 				},

internal/controller/tensorfusionconnection_controller_test.go

@@ -27,6 +27,9 @@ func NewDryRunReporter() Reporter {
 
 func (r *DryRunReporter) Report(ctx context.Context, obj client.Object, f controllerutil.MutateFn) error {
 	log := log.FromContext(ctx)
+	if err := f(); err != nil {
+		return err
+	}
 	objYaml, err := yaml.Marshal(obj)
 	if err != nil {
 		return err

internal/reporter/reporter.go

@@ -14,7 +14,7 @@ func createGPU(name string, tflops, vram string) *tfv1.GPU {
 			Name: name,
 		},
 		Status: tfv1.GPUStatus{
-			Available: tfv1.Resource{
+			Available: &tfv1.Resource{
 				Tflops: resource.MustParse(tflops),
 				Vram:   resource.MustParse(vram),
 			},
@@ -189,11 +189,11 @@ func TestNaiveScheduler_Release(t *testing.T) {
 					Name: "gpu1",
 				},
 				Status: tfv1.GPUStatus{
-					Capacity: tfv1.Resource{
+					Capacity: &tfv1.Resource{
 						Tflops: resource.MustParse("100"),
 						Vram:   resource.MustParse("16Gi"),
 					},
-					Available: tfv1.Resource{
+					Available: &tfv1.Resource{
 						Tflops: resource.MustParse("100"),
 						Vram:   resource.MustParse("16Gi"),
 					},
@@ -218,11 +218,11 @@ func TestNaiveScheduler_Release(t *testing.T) {
 					Name: "gpu1",
 				},
 				Status: tfv1.GPUStatus{
-					Capacity: tfv1.Resource{
+					Capacity: &tfv1.Resource{
 						Tflops: resource.MustParse("100"),
 						Vram:   resource.MustParse("16Gi"),
 					},
-					Available: tfv1.Resource{
+					Available: &tfv1.Resource{
 						Tflops: resource.MustParse("100"),
 						Vram:   resource.MustParse("16Gi"),
 					},