|
| 1 | +// This Source Code Form is subject to the terms of the Mozilla Public |
| 2 | +// License, v. 2.0. If a copy of the MPL was not distributed with this |
| 3 | +// file, You can obtain one at http://mozilla.org/MPL/2.0/. |
| 4 | + |
| 5 | +package main |
| 6 | + |
| 7 | +import ( |
| 8 | + "bufio" |
| 9 | + "context" |
| 10 | + "fmt" |
| 11 | + "log" |
| 12 | + "os" |
| 13 | + "os/exec" |
| 14 | + "os/signal" |
| 15 | + "path/filepath" |
| 16 | + "strings" |
| 17 | + "sync" |
| 18 | + "syscall" |
| 19 | + "time" |
| 20 | + |
| 21 | + "github.com/goaux/decowriter" |
| 22 | +) |
| 23 | + |
| 24 | +const ( |
| 25 | + // FabricManager |
| 26 | + fmCmdFile = "/usr/bin/nv-fabricmanager" |
| 27 | + fmConfigFile = "/usr/share/nvidia/nvswitch/fabricmanager.cfg" |
| 28 | + fmStopTimeout = 5 * time.Second |
| 29 | + |
| 30 | + // NVLSM |
| 31 | + smCmdFile = "/opt/nvidia/nvlsm/sbin/nvlsm" |
| 32 | + smConfigFile = "/usr/share/nvidia/nvlsm/nvlsm.conf" |
| 33 | + smPidFile = "/var/run/nvlsm.pid" |
| 34 | + smSocket = "/var/run/nvidia-fabricmanager/fm_sm_ipc.socket" |
| 35 | + smStopTimeout = 5 * time.Second |
| 36 | + smSocketWait = 15 * time.Second |
| 37 | +) |
| 38 | + |
| 39 | +func runCommand(ctx context.Context, wg *sync.WaitGroup, doneCb func(), waitDelay time.Duration, path string, arg ...string) { |
| 40 | + wg.Add(1) |
| 41 | + |
| 42 | + cmd := exec.CommandContext(ctx, path, arg...) |
| 43 | + cmd.WaitDelay = waitDelay |
| 44 | + cmd.Cancel = func() error { |
| 45 | + return cmd.Process.Signal(os.Interrupt) |
| 46 | + } |
| 47 | + |
| 48 | + // TODO line writer to log module |
| 49 | + name := filepath.Base(path) |
| 50 | + cmd.Stdout = decowriter.New(bufio.NewWriter(os.Stdout), []byte(name+": "), []byte{}) |
| 51 | + cmd.Stderr = decowriter.New(bufio.NewWriter(os.Stderr), []byte(name+": "), []byte{}) |
| 52 | + |
| 53 | + go func() { |
| 54 | + log.Printf("nvidia-fabricmanager-wrapper: running command: %s %s\n", path, strings.Join(arg, " ")) |
| 55 | + |
| 56 | + err := cmd.Run() |
| 57 | + if err == nil { |
| 58 | + log.Printf("nvidia-fabricmanager-wrapper: command %s [%d] completed successfully\n", path, cmd.Process.Pid) |
| 59 | + } else if exitErr, ok := err.(*exec.ExitError); ok { |
| 60 | + if exitErr.Exited() { |
| 61 | + log.Printf("nvidia-fabricmanager-wrapper: command %s [%d] exited with code %d\n", path, exitErr.Pid(), |
| 62 | + exitErr.ExitCode()) |
| 63 | + } else { |
| 64 | + log.Printf("nvidia-fabricmanager-wrapper: command %s [%d] was terminated\n", path, exitErr.Pid()) |
| 65 | + } |
| 66 | + } else { |
| 67 | + log.Printf("nvidia-fabricmanager-wrapper: failed to run command %s: %v\n", path, err) |
| 68 | + } |
| 69 | + |
| 70 | + wg.Done() |
| 71 | + doneCb() |
| 72 | + }() |
| 73 | +} |
| 74 | + |
| 75 | +func waitForFile(ctx context.Context, filepath string, timeout time.Duration) error { |
| 76 | + timer := time.NewTimer(timeout) |
| 77 | + defer timer.Stop() |
| 78 | + |
| 79 | + for { |
| 80 | + select { |
| 81 | + case <-ctx.Done(): |
| 82 | + return fmt.Errorf("parent context canceled: %w", ctx.Err()) |
| 83 | + case <-timer.C: |
| 84 | + return fmt.Errorf("timeout waiting for file") |
| 85 | + default: |
| 86 | + if _, err := os.Stat(filepath); err == nil { |
| 87 | + return nil |
| 88 | + } |
| 89 | + time.Sleep(100 * time.Millisecond) |
| 90 | + } |
| 91 | + } |
| 92 | +} |
| 93 | + |
| 94 | +func main() { |
| 95 | + var cmdWg sync.WaitGroup |
| 96 | + |
| 97 | + signal.Ignore(syscall.SIGHUP) |
| 98 | + |
| 99 | + runCtx, gracefulShutdown := context.WithCancel(context.Background()) |
| 100 | + |
| 101 | + signalsChan := make(chan os.Signal, 1) |
| 102 | + signal.Notify(signalsChan, os.Interrupt) |
| 103 | + signal.Notify(signalsChan, syscall.SIGTERM) |
| 104 | + |
| 105 | + go func() { |
| 106 | + received := <-signalsChan |
| 107 | + signal.Stop(signalsChan) |
| 108 | + log.Printf("nvidia-fabricmanager-wrapper: received signal '%s', initiating a graceful shutdown\n", received.String()) |
| 109 | + gracefulShutdown() |
| 110 | + }() |
| 111 | + |
| 112 | + nvswitchPorts := findNvswitchMgmtPorts() |
| 113 | + for _, port := range nvswitchPorts { |
| 114 | + log.Printf("nvidia-fabricmanager-wrapper: found NVSwitch LPF: device=%s guid=0x%x\n", port.IBDevice, port.PortGUID) |
| 115 | + } |
| 116 | + |
| 117 | + fmSmMgmtPortGUID := "" |
| 118 | + if len(nvswitchPorts) > 0 { |
| 119 | + fmSmMgmtPortGUID = fmt.Sprintf("0x%x", nvswitchPorts[0].PortGUID) |
| 120 | + log.Printf("nvidia-fabricmanager-wrapper: using NVSwitch management port GUID: %s\n", fmSmMgmtPortGUID) |
| 121 | + } else { |
| 122 | + log.Println("nvidia-fabricmanager-wrapper: No InfiniBand NVSwitch detected. On Blackwell HGX baseboards and newer", |
| 123 | + "with NVLink 5.0+, please load kernel module 'ib_umad' for NVLSM to run along FabricManager. Otherwise it will", |
| 124 | + "fail to start with error NV_WARN_NOTHING_TO_DO, and GPU workloads will report CUDA_ERROR_SYSTEM_NOT_READY.") |
| 125 | + } |
| 126 | + |
| 127 | + if fmSmMgmtPortGUID != "" { |
| 128 | + if err := os.Mkdir(filepath.Dir(smSocket), 0755); err != nil { |
| 129 | + log.Printf("nvidia-fabricmanager-wrapper: error creating socket directory: %v\n", err) |
| 130 | + } |
| 131 | + |
| 132 | + runCommand(runCtx, &cmdWg, gracefulShutdown, smStopTimeout, smCmdFile, "--config", smConfigFile, |
| 133 | + "--guid", fmSmMgmtPortGUID, "--pid_file", smPidFile, "--log_file", "stdout") |
| 134 | + |
| 135 | + // vendor startup script waits for 5 seconds for NVLSM socket to be available before starting FM |
| 136 | + // let's wait for the actual GRPC socket to be created by the plugin |
| 137 | + log.Println("nvidia-fabricmanager-wrapper: waiting for socket creation at", smSocket) |
| 138 | + err := waitForFile(runCtx, smSocket, smSocketWait) |
| 139 | + if err != nil { |
| 140 | + log.Printf("nvidia-fabricmanager-wrapper: error waiting for socket: %v\n", err) |
| 141 | + } else { |
| 142 | + log.Println("nvidia-fabricmanager-wrapper: socket found at", smSocket) |
| 143 | + } |
| 144 | + // for safety |
| 145 | + time.Sleep(time.Second) |
| 146 | + } |
| 147 | + |
| 148 | + fmCmdArgs := []string{"--config", fmConfigFile} |
| 149 | + if fmSmMgmtPortGUID != "" { |
| 150 | + fmCmdArgs = append(fmCmdArgs, "--fm-sm-mgmt-port-guid", fmSmMgmtPortGUID) |
| 151 | + } |
| 152 | + runCommand(runCtx, &cmdWg, gracefulShutdown, fmStopTimeout, fmCmdFile, fmCmdArgs...) |
| 153 | + |
| 154 | + log.Println("nvidia-fabricmanager-wrapper: initialization completed") |
| 155 | + cmdWg.Wait() |
| 156 | +} |
0 commit comments