feat: add application logic for metadata collector (#278)

Signed-off-by: Ajay Mishra <[email protected]>

Author: XRFXLP

data-models/pkg/model/gpu_metadata.go

@@ -0,0 +1,39 @@
+// Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package model
+
+type GPUMetadata struct {
+	Version       string    `json:"version"`
+	Timestamp     string    `json:"timestamp"`
+	NodeName      string    `json:"node_name"`
+	ChassisSerial *string   `json:"chassis_serial"`
+	GPUs          []GPUInfo `json:"gpus"`
+	NVSwitches    []string  `json:"nvswitches"`
+}
+
+type GPUInfo struct {
+	GPUID        int      `json:"gpu_id"`
+	UUID         string   `json:"uuid"`
+	PCIAddress   string   `json:"pci_address"`
+	SerialNumber string   `json:"serial_number"`
+	DeviceName   string   `json:"device_name"`
+	NVLinks      []NVLink `json:"nvlinks"`
+}
+
+type NVLink struct {
+	LinkID           int    `json:"link_id"`
+	RemotePCIAddress string `json:"remote_pci_address"`
+	RemoteLinkID     int    `json:"remote_link_id"`
+}

metadata-collector/Dockerfile

@@ -17,7 +17,7 @@ FROM golang:1.25-bookworm AS builder
 WORKDIR /workspace
 
 COPY commons/ commons/
-
+COPY data-models/ data-models/
 
 COPY metadata-collector/go.mod metadata-collector/go.sum metadata-collector/
 WORKDIR /workspace/metadata-collector
@@ -27,7 +27,7 @@ COPY metadata-collector/ .
 
 RUN CGO_ENABLED=1 GOOS=linux go build -a -ldflags="-s -w" -o metadata-collector main.go
 
-FROM ubuntu:22.04
+FROM nvidia/cuda:12.3.0-base-ubuntu22.04
 
 RUN groupadd -r nvsentinel && useradd -r -g nvsentinel nvsentinel
 

metadata-collector/go.mod

@@ -5,8 +5,25 @@ go 1.25
 toolchain go1.25.3
 
 require (
-	github.com/NVIDIA/go-nvml v0.12.4-0
+	github.com/NVIDIA/go-nvml v0.13.0-1
 	github.com/nvidia/nvsentinel/commons v0.0.0
+	github.com/nvidia/nvsentinel/data-models v0.0.0
+	github.com/stretchr/testify v1.11.1
 )
 
-replace github.com/nvidia/nvsentinel/commons => ../commons
+require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	golang.org/x/net v0.46.0 // indirect
+	golang.org/x/sys v0.37.0 // indirect
+	golang.org/x/text v0.30.0 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20251014184007-4626949a642f // indirect
+	google.golang.org/grpc v1.76.0 // indirect
+	google.golang.org/protobuf v1.36.10 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
+)
+
+replace (
+	github.com/nvidia/nvsentinel/commons => ../commons
+	github.com/nvidia/nvsentinel/data-models => ../data-models
+)

metadata-collector/go.sum

@@ -1,10 +1,48 @@
-github.com/NVIDIA/go-nvml v0.12.4-0 h1:4tkbB3pT1O77JGr0gQ6uD8FrsUPqP1A/EOEm2wI1TUg=
-github.com/NVIDIA/go-nvml v0.12.4-0/go.mod h1:8Llmj+1Rr+9VGGwZuRer5N/aCjxGuR5nPb/9ebBiIEQ=
+github.com/NVIDIA/go-nvml v0.13.0-1 h1:OLX8Jq3dONuPOQPC7rndB6+iDmDakw0XTYgzMxObkEw=
+github.com/NVIDIA/go-nvml v0.13.0-1/go.mod h1:+KNA7c7gIBH7SKSJ1ntlwkfN80zdx8ovl4hrK3LmPt4=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI=
+github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
+github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
+github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
+github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
+github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
+github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
+github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
+github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
+github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
 github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
+go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64=
+go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y=
+go.opentelemetry.io/otel v1.38.0 h1:RkfdswUDRimDg0m2Az18RKOsnI8UDzppJAtj01/Ymk8=
+go.opentelemetry.io/otel v1.38.0/go.mod h1:zcmtmQ1+YmQM9wrNsTGV/q/uyusom3P8RxwExxkZhjM=
+go.opentelemetry.io/otel/metric v1.37.0 h1:mvwbQS5m0tbmqML4NqK+e3aDiO02vsf/WgbsdpcPoZE=
+go.opentelemetry.io/otel/metric v1.37.0/go.mod h1:04wGrZurHYKOc+RKeye86GwKiTb9FKm1WHtO+4EVr2E=
+go.opentelemetry.io/otel/sdk v1.37.0 h1:ItB0QUqnjesGRvNcmAcU0LyvkVyGJ2xftD29bWdDvKI=
+go.opentelemetry.io/otel/sdk v1.37.0/go.mod h1:VredYzxUvuo2q3WRcDnKDjbdvmO0sCzOvVAiY+yUkAg=
+go.opentelemetry.io/otel/sdk/metric v1.37.0 h1:90lI228XrB9jCMuSdA0673aubgRobVZFhbjxHHspCPc=
+go.opentelemetry.io/otel/sdk/metric v1.37.0/go.mod h1:cNen4ZWfiD37l5NhS+Keb5RXVWZWpRE+9WyVCpbo5ps=
+go.opentelemetry.io/otel/trace v1.37.0 h1:HLdcFNbRQBE2imdSEgm/kwqmQj1Or1l/7bW6mxVK7z4=
+go.opentelemetry.io/otel/trace v1.37.0/go.mod h1:TlgrlQ+PtQO5XFerSPUYG0JSgGyryXewPGyayAWSBS0=
+golang.org/x/net v0.46.0 h1:giFlY12I07fugqwPuWJi68oOnpfqFnJIJzaIIm2JVV4=
+golang.org/x/net v0.46.0/go.mod h1:Q9BGdFy1y4nkUwiLvT5qtyhAnEHgnQ/zd8PfU6nc210=
+golang.org/x/sys v0.37.0 h1:fdNQudmxPjkdUTPnLn5mdQv7Zwvbvpaxqs831goi9kQ=
+golang.org/x/sys v0.37.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
+golang.org/x/text v0.30.0 h1:yznKA/E9zq54KzlzBEAWn1NXSQ8DIp/NYMy88xJjl4k=
+golang.org/x/text v0.30.0/go.mod h1:yDdHFIX9t+tORqspjENWgzaCVXgk0yYnYuSZ8UzzBVM=
+gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk=
+gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20251014184007-4626949a642f h1:1FTH6cpXFsENbPR5Bu8NQddPSaUUE6NA2XdZdDSAJK4=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20251014184007-4626949a642f/go.mod h1:7i2o+ce6H/6BluujYR+kqX3GKH+dChPTQU19wjRPiGk=
+google.golang.org/grpc v1.76.0 h1:UnVkv1+uMLYXoIz6o7chp59WfQUYA2ex/BXQ9rHZu7A=
+google.golang.org/grpc v1.76.0/go.mod h1:Ju12QI8M6iQJtbcsV+awF5a4hfJMLi4X0JLo94ULZ6c=
+google.golang.org/protobuf v1.36.10 h1:AYd7cD/uASjIL6Q9LiTjz8JLcrh/88q5UObnmY3aOOE=
+google.golang.org/protobuf v1.36.10/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

metadata-collector/main.go

@@ -15,84 +15,92 @@
 package main
 
 import (
+	"context"
+	"flag"
 	"fmt"
 	"log/slog"
 	"os"
+	"os/signal"
+	"syscall"
 
-	"github.com/NVIDIA/go-nvml/pkg/nvml"
 	"github.com/nvidia/nvsentinel/commons/pkg/logger"
+	"github.com/nvidia/nvsentinel/metadata-collector/pkg/collector"
+	"github.com/nvidia/nvsentinel/metadata-collector/pkg/nvml"
+	"github.com/nvidia/nvsentinel/metadata-collector/pkg/writer"
 )
 
 const (
-	defaultAgentName = "metadata-collector"
+	defaultAgentName  = "metadata-collector"
+	defaultOutputPath = "/var/lib/nvsentinel/gpu_metadata.json"
 )
 
 var (
 	version = "dev"
 	commit  = "none"
 	date    = "unknown"
+
+	outputPath = flag.String("output-path", defaultOutputPath, "Path to write the GPU metadata JSON file")
 )
 
 func main() {
+	flag.Parse()
+
 	logger.SetDefaultStructuredLogger(defaultAgentName, version)
 	slog.Info("Starting metadata-collector", "version", version, "commit", commit, "date", date)
 
-	if err := run(); err != nil {
+	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+
+	if err := run(ctx); err != nil {
 		slog.Error("Metadata collector failed", "error", err)
+		cancel()
 		os.Exit(1)
 	}
+
+	cancel()
+	slog.Info("Metadata collector completed successfully")
 }
 
-func run() error {
-	ret := nvml.Init()
-	if ret != nvml.SUCCESS {
-		return fmt.Errorf("failed to initialize NVML: %v", nvml.ErrorString(ret))
+func run(ctx context.Context) error {
+	slog.Info("Initializing NVML")
+
+	nvmlWrapper := &nvml.NVMLWrapper{}
+	if err := nvmlWrapper.Init(); err != nil {
+		return fmt.Errorf("failed to initialize NVML: %w", err)
 	}
 
 	defer func() {
-		ret := nvml.Shutdown()
-		if ret != nvml.SUCCESS {
-			slog.Error("Failed to shutdown NVML", "error", nvml.ErrorString(ret))
+		if err := nvmlWrapper.Shutdown(); err != nil {
+			slog.Error("Failed to shutdown NVML", "error", err)
 		}
 	}()
 
-	count, ret := nvml.DeviceGetCount()
-	if ret != nvml.SUCCESS {
-		return fmt.Errorf("failed to get device count: %v", nvml.ErrorString(ret))
-	}
+	slog.Info("Collecting GPU metadata")
 
-	hostname, _ := os.Hostname()
+	metadataCollector := collector.NewCollector(nvmlWrapper)
 
-	slog.Info("GPU metadata collection started", "node", hostname, "gpu_count", count)
-
-	if nvmlVersion, ret := nvml.SystemGetNVMLVersion(); ret == nvml.SUCCESS {
-		slog.Info("NVML version", "version", nvmlVersion)
+	metadata, err := metadataCollector.Collect(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to collect GPU metadata: %w", err)
 	}
 
-	fmt.Printf("\n=== GPU Metadata Collector ===\n")
-	fmt.Printf("Node: %s\n", hostname)
-	fmt.Printf("GPUs Found: %d\n", count)
-
-	fmt.Println("\n=== GPU Details ===")
+	slog.Info("GPU metadata collected",
+		"node", metadata.NodeName,
+		"gpu_count", len(metadata.GPUs),
+		"nvswitch_count", len(metadata.NVSwitches),
+	)
 
-	for i := range count {
-		device, ret := nvml.DeviceGetHandleByIndex(i)
-		if ret != nvml.SUCCESS {
-			slog.Warn("Failed to get device", "gpu_id", i, "error", nvml.ErrorString(ret))
-			continue
-		}
+	slog.Info("Writing metadata to file", "output_path", *outputPath)
 
-		name, _ := device.GetName()
-		uuid, _ := device.GetUUID()
-
-		slog.Info("GPU discovered", "gpu_id", i, "name", name, "uuid", uuid)
+	metadataWriter, err := writer.NewWriter(*outputPath)
+	if err != nil {
+		return fmt.Errorf("failed to create metadata writer: %w", err)
+	}
 
-		fmt.Printf("\nGPU %d:\n", i)
-		fmt.Printf("  Name: %s\n", name)
-		fmt.Printf("  UUID: %s\n", uuid)
+	if err := metadataWriter.Write(metadata); err != nil {
+		return fmt.Errorf("failed to write metadata: %w", err)
 	}
 
-	slog.Info("Metadata collector hello world completed successfully")
+	slog.Info("Successfully wrote GPU metadata", "output_path", *outputPath)
 
 	return nil
 }

metadata-collector/pkg/collector/collector.go

@@ -0,0 +1,131 @@
+// Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package collector
+
+import (
+	"context"
+	"fmt"
+	"log/slog"
+	"os"
+	"time"
+
+	gonvml "github.com/NVIDIA/go-nvml/pkg/nvml"
+	"github.com/nvidia/nvsentinel/data-models/pkg/model"
+	"github.com/nvidia/nvsentinel/metadata-collector/pkg/nvml"
+)
+
+type Collector struct {
+	nvml *nvml.NVMLWrapper
+}
+
+func NewCollector(nvmlWrapper *nvml.NVMLWrapper) *Collector {
+	return &Collector{
+		nvml: nvmlWrapper,
+	}
+}
+
+func (c *Collector) Collect(ctx context.Context) (*model.GPUMetadata, error) {
+	count, err := c.nvml.GetDeviceCount()
+	if err != nil {
+		return nil, fmt.Errorf("failed to get GPU device count: %w", err)
+	}
+
+	nodeName, err := os.Hostname()
+	if err != nil {
+		return nil, fmt.Errorf("failed to get hostname: %w", err)
+	}
+
+	deviceMap, parsedTopology, err := c.prepareTopologyData(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to prepare topology data: %w", err)
+	}
+
+	metadata := &model.GPUMetadata{
+		Version:   "1.0",
+		Timestamp: time.Now().UTC().Format(time.RFC3339),
+		NodeName:  nodeName,
+		GPUs:      make([]model.GPUInfo, 0, count),
+	}
+
+	if err := c.collectGPUData(count, metadata, deviceMap, parsedTopology); err != nil {
+		return nil, fmt.Errorf("failed to collect GPU data: %w", err)
+	}
+
+	return metadata, nil
+}
+
+func (c *Collector) prepareTopologyData(
+	ctx context.Context,
+) (map[string]gonvml.Device, map[int]nvml.GPUNVLinkTopology, error) {
+	slog.Info("Building device map for NVLink topology discovery")
+
+	deviceMap, err := c.nvml.BuildDeviceMap()
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to build device map: %w", err)
+	}
+
+	slog.Info("Parsing NVLink topology from nvidia-smi")
+
+	parsedTopology, err := c.nvml.ParseNVLinkTopologyWithContext(ctx)
+	if err != nil {
+		slog.Warn("Failed to parse nvidia-smi topology, remote link IDs will be -1", "error", err)
+
+		parsedTopology = make(map[int]nvml.GPUNVLinkTopology)
+	}
+
+	return deviceMap, parsedTopology, nil
+}
+
+func (c *Collector) collectGPUData(
+	count int,
+	metadata *model.GPUMetadata,
+	deviceMap map[string]gonvml.Device,
+	parsedTopology map[int]nvml.GPUNVLinkTopology,
+) error {
+	nvswitchSet := make(map[string]bool)
+
+	var chassisSerial *string
+
+	for i := range count {
+		nvmlGPUInfo, err := c.nvml.GetGPUInfo(i)
+		if err != nil {
+			return fmt.Errorf("failed to get GPU info for GPU %d: %w", i, err)
+		}
+
+		if i == 0 {
+			chassisSerial = c.nvml.GetChassisSerial(i)
+		}
+
+		nvswitches, err := c.nvml.CollectNVLinkTopology(nvmlGPUInfo, i, deviceMap, parsedTopology)
+		if err != nil {
+			slog.Warn("Failed to collect NVLink topology for GPU", "gpu_id", i, "error", err)
+		} else {
+			for pci := range nvswitches {
+				nvswitchSet[pci] = true
+			}
+		}
+
+		metadata.GPUs = append(metadata.GPUs, *nvmlGPUInfo)
+	}
+
+	metadata.ChassisSerial = chassisSerial
+
+	metadata.NVSwitches = make([]string, 0, len(nvswitchSet))
+	for pci := range nvswitchSet {
+		metadata.NVSwitches = append(metadata.NVSwitches, pci)
+	}
+
+	return nil
+}