From d65975c9df77fb92baa17c4ace64cfe449f055d5 Mon Sep 17 00:00:00 2001
From: Gyuho Lee <6799218+gyuho@users.noreply.github.com>
Date: Sat, 30 Nov 2024 20:45:46 +0800
Subject: [PATCH] fix(cmd/gpud): handle "run
 --expected-port-states-nvidia-infiniband" flag (#212)

Signed-off-by: Gyuho Lee <gyuhox@gmail.com>
---
 cmd/gpud/command/run.go | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/cmd/gpud/command/run.go b/cmd/gpud/command/run.go
index f7aa67cb..9751f789 100644
--- a/cmd/gpud/command/run.go
+++ b/cmd/gpud/command/run.go
@@ -7,6 +7,7 @@ import (
 	"os/signal"
 	"time"
 
+	"github.com/leptonai/gpud/components/accelerator/nvidia/infiniband"
 	"github.com/leptonai/gpud/config"
 	lepServer "github.com/leptonai/gpud/internal/server"
 	"github.com/leptonai/gpud/log"
@@ -45,6 +46,14 @@ func cmdRun(cliContext *cli.Context) error {
 		config.WithKubeletIgnoreConnectionErrors(kubeletIgnoreConnectionErrors),
 	}
 
+	if expectedPortStates != "" {
+		portStates := &infiniband.ExpectedPortStates{}
+		if err := json.Unmarshal([]byte(expectedPortStates), portStates); err != nil {
+			return err
+		}
+		configOpts = append(configOpts, config.WithExpectedPortStates(*portStates))
+	}
+
 	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
 	cfg, err := config.DefaultConfig(ctx, configOpts...)
 	cancel()