feat(login): populate resource spec cpu/memory + GPU counts nvidia.com/gpu (#668)

Signed-off-by: Gyuho Lee <[email protected]>

---------

Signed-off-by: Gyuho Lee <[email protected]>

Author: gyuho

api/v1/login.go

@@ -2,13 +2,14 @@ package v1
 
 // LoginRequest is the request for the login request.
 type LoginRequest struct {
-	Token       string             `json:"token"`
-	Network     MachineNetworkSpec `json:"network"`
-	Location    MachineLocation    `json:"location"`
-	CPUInfo     MachineCPUInfo     `json:"cpuInfo"`
-	GPUInfo     MachineGPUInfo     `json:"gpuInfo"`
-	Provider    string             `json:"provider"`
-	MachineInfo MachineInfo        `json:"machineInfo"`
+	Token        string             `json:"token"`
+	Network      MachineNetworkSpec `json:"network"`
+	Location     MachineLocation    `json:"location"`
+	CPUInfo      MachineCPUInfo     `json:"cpuInfo"`
+	GPUInfo      MachineGPUInfo     `json:"gpuInfo"`
+	Provider     string             `json:"provider"`
+	MachineInfo  MachineInfo        `json:"machineInfo"`
+	ResourceSpec map[string]string  `json:"resourceSpec"`
 }
 
 type MachineNetworkSpec struct {

cmd/gpud/command/login.go

@@ -6,12 +6,16 @@ import (
 	"time"
 
 	"github.com/urfave/cli"
+	corev1 "k8s.io/api/core/v1"
 
 	apiv1 "github.com/leptonai/gpud/api/v1"
 	client "github.com/leptonai/gpud/client/v1"
 	"github.com/leptonai/gpud/pkg/config"
+	pkgcpu "github.com/leptonai/gpud/pkg/cpu"
 	gpudstate "github.com/leptonai/gpud/pkg/gpud-state"
 	"github.com/leptonai/gpud/pkg/login"
+	pkgmemory "github.com/leptonai/gpud/pkg/memory"
+	nvidiaquery "github.com/leptonai/gpud/pkg/nvidia-query"
 	"github.com/leptonai/gpud/pkg/server"
 	"github.com/leptonai/gpud/pkg/sqlite"
 )
@@ -59,9 +63,32 @@ func cmdLogin(cliContext *cli.Context) error {
 
 	endpoint := cliContext.String("endpoint")
 
+	req := apiv1.LoginRequest{
+		Token:        token,
+		ResourceSpec: map[string]string{},
+	}
+
+	cpu, err := pkgcpu.GetSystemResourceLogicalCores()
+	if err != nil {
+		return fmt.Errorf("failed to get system resource logical cores: %w", err)
+	}
+	req.ResourceSpec[string(corev1.ResourceCPU)] = cpu
+
+	memory, err := pkgmemory.GetSystemResourceMemoryTotal()
+	if err != nil {
+		return fmt.Errorf("failed to get system resource memory total: %w", err)
+	}
+	req.ResourceSpec[string(corev1.ResourceMemory)] = memory
+
+	gpuCnt, err := nvidiaquery.GetSystemResourceGPUCount()
+	if err != nil {
+		return fmt.Errorf("failed to get system resource gpu count: %w", err)
+	}
+	req.ResourceSpec["nvidia.com/gpu"] = gpuCnt
+
 	// machine ID has not been assigned yet
 	// thus request one and blocks until the login request is processed
-	loginResp, err := login.SendRequest(rootCtx, endpoint, apiv1.LoginRequest{Token: token})
+	loginResp, err := login.SendRequest(rootCtx, endpoint, req)
 	if err != nil {
 		return err
 	}

pkg/cpu/doc.go

@@ -0,0 +1,2 @@
+// Package cpu provides utilities for CPU usage.
+package cpu

pkg/cpu/utils.go

@@ -0,0 +1,28 @@
+package cpu
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/shirou/gopsutil/v4/cpu"
+	"k8s.io/apimachinery/pkg/api/resource"
+)
+
+// GetSystemResourceLogicalCores returns the system CPU resource of the machine
+// with the logical core counts, using the type defined in "corev1.ResourceName"
+// in https://pkg.go.dev/k8s.io/api/core/v1#ResourceName.
+// It represents the CPU, in cores (500m = .5 cores).
+// Must be parsed using the "resource.ParseQuantity" function in https://pkg.go.dev/k8s.io/apimachinery/pkg/api/resource.
+func GetSystemResourceLogicalCores() (string, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	counts, err := cpu.CountsWithContext(ctx, true)
+	if err != nil {
+		return "", fmt.Errorf("failed to get CPU cores count: %w", err)
+	}
+
+	qty := resource.NewQuantity(int64(counts), resource.DecimalSI)
+	return qty.String(), nil
+}

pkg/cpu/utils_test.go

@@ -0,0 +1,17 @@
+package cpu
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"k8s.io/apimachinery/pkg/api/resource"
+)
+
+func TestGetSystemResource(t *testing.T) {
+	cpu, err := GetSystemResourceLogicalCores()
+	assert.NoError(t, err)
+
+	cpuQty, err := resource.ParseQuantity(cpu)
+	assert.NoError(t, err)
+	t.Logf("cpu: %s", cpuQty.String())
+}

pkg/memory/utils.go

@@ -0,0 +1,28 @@
+package memory
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/shirou/gopsutil/v4/mem"
+	"k8s.io/apimachinery/pkg/api/resource"
+)
+
+// GetSystemResourceMemoryTotal returns the system memory resource of the machine
+// for the total memory size, using the type defined in "corev1.ResourceName"
+// in https://pkg.go.dev/k8s.io/api/core/v1#ResourceName.
+// It represents the Memory, in bytes (500Gi = 500GiB = 500 * 1024 * 1024 * 1024).
+// Must be parsed using the "resource.ParseQuantity" function in https://pkg.go.dev/k8s.io/apimachinery/pkg/api/resource.
+func GetSystemResourceMemoryTotal() (string, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	vm, err := mem.VirtualMemoryWithContext(ctx)
+	if err != nil {
+		return "", fmt.Errorf("failed to get memory: %w", err)
+	}
+
+	qty := resource.NewQuantity(int64(vm.Total), resource.DecimalSI)
+	return qty.String(), nil
+}

pkg/memory/utils_test.go

@@ -0,0 +1,17 @@
+package memory
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"k8s.io/apimachinery/pkg/api/resource"
+)
+
+func TestGetSystemResource(t *testing.T) {
+	mem, err := GetSystemResourceMemoryTotal()
+	assert.NoError(t, err)
+
+	memQty, err := resource.ParseQuantity(mem)
+	assert.NoError(t, err)
+	t.Logf("mem: %s", memQty.String())
+}

pkg/nvidia-query/device_count.go

@@ -1,13 +1,18 @@
 package query
 
 import (
+	"errors"
 	"os"
 	"path/filepath"
 	"regexp"
 )
 
 func CountAllDevicesFromDevDir() (int, error) {
-	return countAllDevicesFromDir("/dev")
+	dir := "/dev"
+	if _, err := os.Stat(dir); errors.Is(err, os.ErrNotExist) {
+		return 0, nil
+	}
+	return countAllDevicesFromDir(dir)
 }
 
 // "checkPermissions" in "nvvs/plugin_src/software/Software.cpp"

pkg/nvidia-query/device_count_test.go

@@ -84,3 +84,11 @@ func TestCountAllDevicesFromDir(t *testing.T) {
 		t.Errorf("expected %d devices, but got %d", devCnt, count)
 	}
 }
+
+func TestCountAllDevicesFromDevDir(t *testing.T) {
+	devCnt, err := CountAllDevicesFromDevDir()
+	if err != nil {
+		t.Fatalf("CountAllDevicesFromDevDir returned an error: %v", err)
+	}
+	t.Logf("CountAllDevicesFromDevDir: %d", devCnt)
+}

pkg/nvidia-query/utils.go

@@ -0,0 +1,34 @@
+package query
+
+import (
+	"github.com/leptonai/gpud/pkg/log"
+	nvidianvml "github.com/leptonai/gpud/pkg/nvidia-query/nvml"
+
+	"k8s.io/apimachinery/pkg/api/resource"
+)
+
+// GetSystemResourceLogicalCores returns the system GPU resource of the machine
+// with the GPU count, using the type defined in "corev1.ResourceName"
+// in https://pkg.go.dev/k8s.io/api/core/v1#ResourceName.
+// It represents the GPU count with the key "nvidia.com/gpu" or "nvidia.com/gpu.count".
+// Must be parsed using the "resource.ParseQuantity" function in https://pkg.go.dev/k8s.io/apimachinery/pkg/api/resource.
+//
+// This is different than the device count in DCGM.
+// ref. "CountDevEntry" in "nvvs/plugin_src/software/Software.cpp"
+// ref. https://github.com/NVIDIA/DCGM/blob/903d745504f50153be8293f8566346f9de3b3c93/nvvs/plugin_src/software/Software.cpp#L220-L249
+func GetSystemResourceGPUCount() (string, error) {
+	nvmlInstance, err := nvidianvml.NewInstanceV2()
+	if err != nil {
+		return "", err
+	}
+	defer func() {
+		if err := nvmlInstance.Shutdown(); err != nil {
+			log.Logger.Warnw("failed to shutdown nvml instance", "error", err)
+		}
+	}()
+
+	deviceCount := len(nvmlInstance.Devices())
+
+	qty := resource.NewQuantity(int64(deviceCount), resource.DecimalSI)
+	return qty.String(), nil
+}