diff --git a/charts/hami/README.md b/charts/hami/README.md index 6229688e5c..be21bb6ede 100644 --- a/charts/hami/README.md +++ b/charts/hami/README.md @@ -192,6 +192,8 @@ This document provides detailed descriptions of all configurable values paramete |-----------|-------------|---------------| | `devicePlugin.pluginPath` | Plugin path | `/var/lib/kubelet/device-plugins` | | `devicePlugin.libPath` | Library path | `/usr/local/vgpu` | +| `devicePlugin.hostPID` | Use the host PID namespace for the device plugin | `true` | +| `devicePlugin.hostPIDBroker.enabled` | Let HAMi core ask the device plugin for its host PID. This requires `devicePlugin.hostPID`. See [Host PID broker](../../docs/develop/hostpid-broker.md) | `false` | | `devicePlugin.nvidiaNodeSelector` | NVIDIA node selector | `{"gpu": "on"}` | | `devicePlugin.updateStrategy.type` | Update strategy type | `RollingUpdate` | | `devicePlugin.updateStrategy.rollingUpdate.maxUnavailable` | Maximum unavailable count | `1` | diff --git a/charts/hami/templates/device-plugin/daemonsetnvidia.yaml b/charts/hami/templates/device-plugin/daemonsetnvidia.yaml index d3baa4cdd9..8891ffa0a7 100644 --- a/charts/hami/templates/device-plugin/daemonsetnvidia.yaml +++ b/charts/hami/templates/device-plugin/daemonsetnvidia.yaml @@ -1,3 +1,6 @@ +{{- if and .Values.devicePlugin.enabled .Values.devicePlugin.hostPIDBroker.enabled (not .Values.devicePlugin.hostPID) }} +{{- fail "devicePlugin.hostPIDBroker requires devicePlugin.hostPID" }} +{{- end }} {{- if .Values.devicePlugin.enabled }} apiVersion: apps/v1 kind: DaemonSet @@ -94,6 +97,10 @@ spec: value: {{ .Values.devicePlugin.deviceListStrategy }} - name: HOOK_PATH value: {{ .Values.global.gpuHookPath }} + {{- if .Values.devicePlugin.hostPIDBroker.enabled }} + - name: LIBVGPU_HOSTPID_BROKER + value: "1" + {{- end }} {{- if typeIs "bool" .Values.devicePlugin.passDeviceSpecsEnabled }} - name: PASS_DEVICE_SPECS value: {{ .Values.devicePlugin.passDeviceSpecsEnabled | quote }} @@ -169,6 +176,10 @@ spec: mountPath: /etc/hami/numa-refit-ca readOnly: true {{- end }} + {{- if .Values.devicePlugin.hostPIDBroker.enabled }} + - name: hostpid-broker + mountPath: /var/run/hami/hostpid + {{- end }} {{- if typeIs "string" .Values.devicePlugin.nvidiaDriverRoot }} # We always mount the driver root at /driver-root in the container. # This is required for CDI detection to work correctly. @@ -266,6 +277,12 @@ spec: hostPath: path: /var/run/cdi type: DirectoryOrCreate + {{- if .Values.devicePlugin.hostPIDBroker.enabled }} + - name: hostpid-broker + hostPath: + path: /var/run/hami/hostpid + type: DirectoryOrCreate + {{- end }} - name: usrbin hostPath: path: /usr/bin diff --git a/charts/hami/values.yaml b/charts/hami/values.yaml index 9d031ee346..e4c5a85988 100644 --- a/charts/hami/values.yaml +++ b/charts/hami/values.yaml @@ -426,6 +426,10 @@ devicePlugin: podAnnotations: {} hostPID: true + # Let HAMi core ask the device plugin for its host PID. + # This requires hostPID to be true. + hostPIDBroker: + enabled: false hostNetwork: false securityContext: privileged: true diff --git a/cmd/device-plugin/nvidia/hostpid_broker.go b/cmd/device-plugin/nvidia/hostpid_broker.go new file mode 100644 index 0000000000..8f446aeecb --- /dev/null +++ b/cmd/device-plugin/nvidia/hostpid_broker.go @@ -0,0 +1,78 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * Copyright (c) 2026 The HAMi Authors. + */ + +package main + +import ( + "errors" + "fmt" + "os" + + "k8s.io/klog/v2" + + "github.com/Project-HAMi/HAMi/pkg/device-plugin/nvidiadevice/nvinternal/hostpid" +) + +type runningHostPIDBroker struct { + broker hostPIDBroker + done chan struct{} + serveErr error +} + +type hostPIDBroker interface { + Serve() error + Close() error +} + +type hostPIDBrokerListener func() (hostPIDBroker, error) + +func startHostPIDBroker() (*runningHostPIDBroker, error) { + return startHostPIDBrokerWithListener(func() (hostPIDBroker, error) { + return hostpid.ListenDefault() + }) +} + +func startHostPIDBrokerWithListener( + listen hostPIDBrokerListener) (*runningHostPIDBroker, error) { + if !hostpid.Enabled(os.Getenv(hostpid.EnvironmentVariable)) { + return nil, nil + } + broker, err := listen() + if err != nil { + return nil, err + } + running := &runningHostPIDBroker{ + broker: broker, + done: make(chan struct{}), + } + go func() { + running.serveErr = broker.Serve() + close(running.done) + }() + klog.Infof("Host PID broker is listening on %s", hostpid.ServerSocketPath) + return running, nil +} + +func (running *runningHostPIDBroker) stop() error { + closeErr := running.broker.Close() + <-running.done + return closeErr +} + +func (running *runningHostPIDBroker) failure() error { + if running == nil { + return nil + } + select { + case <-running.done: + if running.serveErr != nil { + return fmt.Errorf("host PID broker stopped: %w", running.serveErr) + } + return errors.New("host PID broker stopped unexpectedly") + default: + return nil + } +} diff --git a/cmd/device-plugin/nvidia/hostpid_broker_lifecycle_test.go b/cmd/device-plugin/nvidia/hostpid_broker_lifecycle_test.go new file mode 100644 index 0000000000..60b999d6c8 --- /dev/null +++ b/cmd/device-plugin/nvidia/hostpid_broker_lifecycle_test.go @@ -0,0 +1,130 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * Copyright (c) 2026 The HAMi Authors. + */ + +package main + +import ( + "errors" + "testing" + + "github.com/Project-HAMi/HAMi/pkg/device-plugin/nvidiadevice/nvinternal/plugin" + "github.com/Project-HAMi/HAMi/pkg/device-plugin/nvidiadevice/nvinternal/rm" +) + +type fakeDevicePlugin struct { + devices rm.Devices + start func(string) error + stopErr error + startCalls int + stopCalls int +} + +func (p *fakeDevicePlugin) Devices() rm.Devices { + return p.devices +} + +func (p *fakeDevicePlugin) Start(socket string) error { + p.startCalls++ + if p.start != nil { + return p.start(socket) + } + return nil +} + +func (p *fakeDevicePlugin) Stop() error { + p.stopCalls++ + return p.stopErr +} + +func devicePluginWithDevice() *fakeDevicePlugin { + return &fakeDevicePlugin{ + devices: rm.Devices{"GPU-0": &rm.Device{}}, + } +} + +func TestStartPluginServersDetectsBrokerFailureBeforeStart(t *testing.T) { + wantErr := errors.New("broker failed") + done := make(chan struct{}) + close(done) + running := &runningHostPIDBroker{ + done: done, + serveErr: wantErr, + } + p := devicePluginWithDevice() + + started, restart, err := startPluginServers( + []plugin.Interface{p}, "/tmp/kubelet.sock", running) + if started != 0 || restart || !errors.Is(err, wantErr) { + t.Fatalf("started=%d restart=%v err=%v, want broker failure", + started, restart, err) + } + if p.startCalls != 0 || p.stopCalls != 0 { + t.Fatalf("start calls=%d stop calls=%d, want 0 and 0", + p.startCalls, p.stopCalls) + } +} + +func TestStartPluginServersCleansUpAfterBrokerFailure(t *testing.T) { + wantServeErr := errors.New("broker failed") + wantStopErr := errors.New("plugin stop failed") + done := make(chan struct{}) + running := &runningHostPIDBroker{done: done} + p := devicePluginWithDevice() + p.stopErr = wantStopErr + p.start = func(string) error { + running.serveErr = wantServeErr + close(done) + return nil + } + + started, restart, err := startPluginServers( + []plugin.Interface{p}, "/tmp/kubelet.sock", running) + if started != 0 || restart || !errors.Is(err, wantServeErr) || + !errors.Is(err, wantStopErr) { + t.Fatalf("started=%d restart=%v err=%v, want joined failures", + started, restart, err) + } + if p.startCalls != 1 || p.stopCalls != 1 { + t.Fatalf("start calls=%d stop calls=%d, want 1 and 1", + p.startCalls, p.stopCalls) + } +} + +func TestStartPluginServersRequestsRestartAfterStartFailure(t *testing.T) { + wantErr := errors.New("plugin start failed") + p := devicePluginWithDevice() + p.start = func(socket string) error { + if socket != "/tmp/kubelet.sock" { + t.Fatalf("socket=%q", socket) + } + return wantErr + } + + started, restart, err := startPluginServers( + []plugin.Interface{p}, "/tmp/kubelet.sock", nil) + if started != 0 || !restart || err != nil { + t.Fatalf("started=%d restart=%v err=%v", started, restart, err) + } + if p.startCalls != 1 || p.stopCalls != 0 { + t.Fatalf("start calls=%d stop calls=%d, want 1 and 0", + p.startCalls, p.stopCalls) + } +} + +func TestStartPluginServersSkipsEmptyPlugins(t *testing.T) { + empty := &fakeDevicePlugin{} + ready := devicePluginWithDevice() + + started, restart, err := startPluginServers( + []plugin.Interface{empty, ready}, "/tmp/kubelet.sock", nil) + if started != 1 || restart || err != nil { + t.Fatalf("started=%d restart=%v err=%v", started, restart, err) + } + if empty.startCalls != 0 || ready.startCalls != 1 { + t.Fatalf("empty starts=%d ready starts=%d, want 0 and 1", + empty.startCalls, ready.startCalls) + } +} diff --git a/cmd/device-plugin/nvidia/hostpid_broker_test.go b/cmd/device-plugin/nvidia/hostpid_broker_test.go new file mode 100644 index 0000000000..a425bc9282 --- /dev/null +++ b/cmd/device-plugin/nvidia/hostpid_broker_test.go @@ -0,0 +1,167 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * Copyright (c) 2026 The HAMi Authors. + */ + +package main + +import ( + "errors" + "sync" + "testing" + + "github.com/Project-HAMi/HAMi/pkg/device-plugin/nvidiadevice/nvinternal/hostpid" +) + +type fakeHostPIDBroker struct { + serveStarted chan struct{} + serveRelease chan struct{} + closeOnce sync.Once + serveErr error + closeErr error +} + +func newFakeHostPIDBroker() *fakeHostPIDBroker { + return &fakeHostPIDBroker{ + serveStarted: make(chan struct{}), + serveRelease: make(chan struct{}), + } +} + +func (broker *fakeHostPIDBroker) Serve() error { + close(broker.serveStarted) + <-broker.serveRelease + return broker.serveErr +} + +func (broker *fakeHostPIDBroker) Close() error { + broker.closeOnce.Do(func() { + close(broker.serveRelease) + }) + return broker.closeErr +} + +func TestStartHostPIDBrokerDisabled(t *testing.T) { + for _, value := range []string{"", "0", "true", "false", "01", " 1"} { + t.Run(value, func(t *testing.T) { + t.Setenv(hostpid.EnvironmentVariable, value) + listenerCalled := false + running, err := startHostPIDBrokerWithListener( + func() (hostPIDBroker, error) { + listenerCalled = true + return nil, nil + }) + if err != nil || running != nil { + t.Fatalf("running=%v err=%v", running, err) + } + if listenerCalled { + t.Fatal("listener was called while broker was disabled") + } + }) + } +} + +func TestStartHostPIDBrokerDefaultDisabled(t *testing.T) { + t.Setenv(hostpid.EnvironmentVariable, "") + running, err := startHostPIDBroker() + if err != nil || running != nil { + t.Fatalf("running=%v err=%v", running, err) + } +} + +func TestStartHostPIDBrokerListenFailure(t *testing.T) { + t.Setenv(hostpid.EnvironmentVariable, "1") + wantErr := errors.New("listen failed") + running, err := startHostPIDBrokerWithListener( + func() (hostPIDBroker, error) { + return nil, wantErr + }) + if running != nil || !errors.Is(err, wantErr) { + t.Fatalf("running=%v err=%v, want %v", running, err, wantErr) + } +} + +func TestStartAndStopHostPIDBroker(t *testing.T) { + t.Setenv(hostpid.EnvironmentVariable, "1") + broker := newFakeHostPIDBroker() + running, err := startHostPIDBrokerWithListener( + func() (hostPIDBroker, error) { + return broker, nil + }) + if err != nil { + t.Fatal(err) + } + <-broker.serveStarted + if err := running.failure(); err != nil { + t.Fatalf("running broker failure=%v", err) + } + if err := running.stop(); err != nil { + t.Fatalf("stop broker: %v", err) + } + if err := running.failure(); err == nil { + t.Fatal("stopped broker was not reported") + } +} + +func TestHostPIDBrokerServeFailure(t *testing.T) { + t.Setenv(hostpid.EnvironmentVariable, "1") + wantErr := errors.New("serve failed") + broker := newFakeHostPIDBroker() + broker.serveErr = wantErr + close(broker.serveRelease) + running, err := startHostPIDBrokerWithListener( + func() (hostPIDBroker, error) { + return broker, nil + }) + if err != nil { + t.Fatal(err) + } + <-running.done + if err := running.failure(); !errors.Is(err, wantErr) { + t.Fatalf("failure=%v, want %v", err, wantErr) + } +} + +func TestHostPIDBrokerCloseFailure(t *testing.T) { + t.Setenv(hostpid.EnvironmentVariable, "1") + wantErr := errors.New("close failed") + broker := newFakeHostPIDBroker() + broker.closeErr = wantErr + running, err := startHostPIDBrokerWithListener( + func() (hostPIDBroker, error) { + return broker, nil + }) + if err != nil { + t.Fatal(err) + } + <-broker.serveStarted + if err := running.stop(); !errors.Is(err, wantErr) { + t.Fatalf("stop=%v, want %v", err, wantErr) + } +} + +func TestRunningHostPIDBrokerFailure(t *testing.T) { + var disabled *runningHostPIDBroker + if err := disabled.failure(); err != nil { + t.Fatalf("disabled broker failure=%v", err) + } + + running := &runningHostPIDBroker{done: make(chan struct{})} + if err := running.failure(); err != nil { + t.Fatalf("running broker failure=%v", err) + } + + serveErr := errors.New("accept failed") + running.serveErr = serveErr + close(running.done) + if err := running.failure(); !errors.Is(err, serveErr) { + t.Fatalf("stopped broker failure=%v, want %v", err, serveErr) + } + + stopped := &runningHostPIDBroker{done: make(chan struct{})} + close(stopped.done) + if err := stopped.failure(); err == nil { + t.Fatal("clean broker stop was not reported") + } +} diff --git a/cmd/device-plugin/nvidia/main.go b/cmd/device-plugin/nvidia/main.go index c91454bf8a..d61490098e 100644 --- a/cmd/device-plugin/nvidia/main.go +++ b/cmd/device-plugin/nvidia/main.go @@ -18,6 +18,7 @@ package main import ( "encoding/json" + "errors" "flag" "fmt" "os" @@ -259,7 +260,7 @@ func loadConfig(c *cli.Context, flags []cli.Flag) (*spec.Config, error) { return config, nil } -func start(c *cli.Context, o *options) error { +func start(c *cli.Context, o *options) (resultErr error) { util.NodeName = os.Getenv(util.NodeNameEnvName) client.InitGlobalClient() @@ -271,6 +272,23 @@ func start(c *cli.Context, o *options) error { } defer watcher.Close() + hostPIDBroker, err := startHostPIDBroker() + if err != nil { + return fmt.Errorf("failed to start host PID broker: %w", err) + } + var hostPIDBrokerDone <-chan struct{} + hostPIDBrokerFailureReported := false + if hostPIDBroker != nil { + hostPIDBrokerDone = hostPIDBroker.done + defer func() { + stopErr := hostPIDBroker.stop() + if !hostPIDBrokerFailureReported { + stopErr = errors.Join(stopErr, hostPIDBroker.serveErr) + } + resultErr = errors.Join(resultErr, stopErr) + }() + } + /*Loading config files*/ klog.Infof("Start working on node %s", util.NodeName) klog.Info("Starting OS watcher.") @@ -289,7 +307,7 @@ restart: } klog.Info("Starting Plugins.") - plugins, restartPlugins, err := startPlugins(c, o) + plugins, restartPlugins, err := startPlugins(c, o, hostPIDBroker) if err != nil { return fmt.Errorf("error starting plugins: %v", err) } @@ -304,6 +322,11 @@ restart: // some messages, trigger a restart of the plugins, or exit the program. for { select { + case <-hostPIDBrokerDone: + hostPIDBrokerFailureReported = true + resultErr = hostPIDBroker.failure() + goto exit + // If the restart timeout has expired, then restart the plugins case <-restartTimeout: goto restart @@ -336,14 +359,15 @@ restart: } } exit: - err = stopPlugins(plugins) - if err != nil { - return fmt.Errorf("error stopping plugins: %v", err) + if err := stopPlugins(plugins); err != nil { + resultErr = errors.Join(resultErr, + fmt.Errorf("error stopping plugins: %v", err)) } - return nil + return resultErr } -func startPlugins(c *cli.Context, o *options) ([]plugin.Interface, bool, error) { +func startPlugins(c *cli.Context, o *options, + hostPIDBroker *runningHostPIDBroker) ([]plugin.Interface, bool, error) { // Load the configuration file klog.Info("Loading configuration.") config, err := loadConfig(c, o.flags) @@ -398,29 +422,50 @@ func startPlugins(c *cli.Context, o *options) ([]plugin.Interface, bool, error) return nil, false, fmt.Errorf("error getting plugins: %v", err) } - // Loop through all plugins, starting them if they have any devices - // to serve. If even one plugin fails to start properly, try - // starting them all again. + started, restartPlugins, err := startPluginServers(plugins, + o.kubeletSocket, hostPIDBroker) + if err != nil { + return nil, false, err + } + if restartPlugins { + return plugins, true, nil + } + + if started == 0 { + klog.Info("No devices found. Waiting indefinitely.") + } + + return plugins, false, nil +} + +func startPluginServers(plugins []plugin.Interface, kubeletSocket string, + hostPIDBroker *runningHostPIDBroker) (int, bool, error) { started := 0 + startedPlugins := make([]plugin.Interface, 0, len(plugins)) for _, p := range plugins { // Just continue if there are no devices to serve for plugin p. if len(p.Devices()) == 0 { continue } + if err := hostPIDBroker.failure(); err != nil { + return started, false, + errors.Join(err, stopPlugins(startedPlugins)) + } + // Start the gRPC server for plugin p and connect it with the kubelet. - if err := p.Start(o.kubeletSocket); err != nil { + if err := p.Start(kubeletSocket); err != nil { klog.Errorf("Failed to start plugin: %v", err) - return plugins, true, nil + return started, true, nil + } + startedPlugins = append(startedPlugins, p) + if err := hostPIDBroker.failure(); err != nil { + return started, false, + errors.Join(err, stopPlugins(startedPlugins)) } started++ } - - if started == 0 { - klog.Info("No devices found. Waiting indefinitely.") - } - - return plugins, false, nil + return started, false, nil } func stopPlugins(plugins []plugin.Interface) error { diff --git a/docs/develop/hostpid-broker.md b/docs/develop/hostpid-broker.md new file mode 100644 index 0000000000..6f300c7695 --- /dev/null +++ b/docs/develop/hostpid-broker.md @@ -0,0 +1,154 @@ +# Host PID broker + +The feature is disabled by default. + +## Purpose + +HAMi-core needs the host PID of each CUDA process because NVML reports processes in the host PID namespace. The current fallback discovers that PID by creating a CUDA primary context while holding the post init lock. That work becomes serial when many processes call `cuInit()` together. + +The host PID broker returns the caller's own host PID from Linux `SO_PEERCRED`. It runs inside the NVIDIA device plugin, whose pod already uses the host PID namespace. It does not read host procfs and does not accept a PID supplied by the client. + +## Requirements + +1. The NVIDIA device plugin must run on Linux as root. + +2. `devicePlugin.hostPID` must remain `true`. + +3. The workload must use a HAMi-core build that supports protocol version 1 and the `LIBVGPU_HOSTPID_BROKER` gate. + +4. The container runtime must support the read-only nested bind mount used for `/tmp/vgpulock/hostpid`. + +5. The shared `/tmp/vgpulock` parent must be owned by root and use mode `01777`. The sticky bit allows legacy lock creation while preventing an ordinary workload user from replacing an entry owned by another user. + +## Enablement + +Set the chart value below: + +```yaml +devicePlugin: + hostPID: true + hostPIDBroker: + enabled: true +``` + +The chart rejects a configuration that enables the broker while disabling the device plugin host PID namespace. + +When enabled, the chart does four things: + +1. It sets `LIBVGPU_HOSTPID_BROKER=1` in the device plugin. + +2. It mounts the host directory `/var/run/hami/hostpid` into the device plugin. + +3. The device plugin creates `/var/run/hami/hostpid/broker.sock` and serves protocol version 1. + +4. Each allocation that receives HAMi-core also receives `LIBVGPU_HOSTPID_BROKER=1` and a read-only mount from `/var/run/hami/hostpid` to `/tmp/vgpulock/hostpid`. + +The device plugin prepares `/tmp/vgpulock` with mode `01777` before returning an allocation. This also applies when the broker is disabled and HAMi-core uses the existing fallback. Allocation fails if the directory cannot be prepared safely. + +Preparation opens `/tmp` without following a symlink, verifies its owner and sticky rule, then creates and opens `vgpulock` relative to that descriptor. The first `mkdirat()` requests mode `01777`. A descriptor-based `chmod` restores bits removed by the process umask. A final identity and mode check rejects replacement during preparation. + +The allocation response contains one writable parent mount at `/tmp/vgpulock`. When the broker is enabled, it also contains one read-only broker mount at `/tmp/vgpulock/hostpid`. The integration replaces duplicate or path-equivalent entries with these canonical mounts and preserves unrelated mounts. It lists the parent before the nested broker mount so the parent does not hide the broker mount when the runtime applies the response. + +Before applying the current gate, the allocation helper clears the reserved broker environment key. When the broker is disabled, it also removes stale broker mounts while preserving the writable parent mount and unrelated mounts. + +No value other than the exact string `1` enables the server or client. + +## Protocol + +The protocol uses one request and one response on a Unix stream connection. Every integer is unsigned and encoded in network byte order. + +| Field | Request bytes | Response bytes | +| --- | ---: | ---: | +| Magic `HPID` | 4 | 4 | +| Version | 2 | 2 | +| Command or status | 2 | 2 | +| Host PID | 0 | 4 | + +Protocol version 1 supports command 1, which means get the caller's host PID. Status 0 is success. Status 1 means the request was invalid. + +The server reads the peer credentials from the connected socket after validating the request. The client never sends a PID. + +## Security boundary + +1. The server requires effective UID 0 for its default path. + +2. The server directory is owned by root with mode `0711`. + +3. The shared lock parent is owned by root with mode `01777`. Allocation rejects an unsafe owner, object type, symlink, or final mode. + +4. A root-owned `0600` lock file prevents two brokers from replacing each other. + +5. The server rejects symlink directories, symlink lock files, regular file collisions, sockets owned by another UID, and active sockets. + +6. The server removes only a stale socket owned by the expected UID. During shutdown it removes the path only if its device and inode still match the socket it created. + +7. The workload sees the broker directory through a read-only mount. The HAMi-core client checks the directory owner, directory write bits, socket type, socket owner, read-only mount flag, and connected peer UID before trusting a response. + +8. A caller can request only its own PID. The kernel supplies that identity through `SO_PEERCRED` in the broker's host PID namespace. + +The socket is available only to workloads that receive the allocation mount. A workload can still create connection pressure. The server limits active handlers, applies one transaction deadline, and closes excess connections. A client that cannot complete the transaction uses the existing bounded HAMi-core fallback. + +## Failure behavior + +| Condition | Server behavior | HAMi-core behavior | +| --- | --- | --- | +| Feature disabled | No socket is created | Existing NVML discovery path | +| Server cannot start | Device plugin startup fails | No new allocation is served by that plugin instance | +| Lock parent cannot be prepared safely | The allocation fails | No unsafe parent mount is returned | +| Socket missing or stale in a workload | No broker response | Existing NVML discovery path | +| Unsafe owner, mode, mount, or peer UID | Request is not trusted | Existing NVML discovery path | +| Malformed protocol reply | Request fails | Existing NVML discovery path | +| Slow or unresponsive broker | Server and client deadlines close the request | Existing NVML discovery path | +| Broker exits after plugin startup | Device plugin exits with the broker error | Kubernetes restarts the device plugin | + +The broker never returns a guessed PID. A successful reply contains the PID supplied by the kernel. Every other outcome is a failure that leaves PID discovery to the existing path. + +## Rollout + +1. Install the server capable HAMi release with `hostPIDBroker.enabled=false`. + +2. Install the compatible HAMi-core library. With no broker mount it uses the existing fallback. + +3. Enable `hostPIDBroker.enabled` and wait for every NVIDIA device plugin pod to become ready. + +4. Restart or recreate selected workloads so their allocation responses include the broker mount and environment gate. + +5. Validate the selected workloads before widening the rollout. Check correct host PID assignment, CUDA context accounting, broker use, and fallback behavior. + +Existing workloads do not gain a new mount when the device plugin changes. They continue through their existing path until they are recreated. + +## Rollback + +1. Set `hostPIDBroker.enabled=false`. + +2. Wait for the NVIDIA device plugin rollout. + +3. Recreate workloads when the broker mount should be removed. Workloads that still have a compatible broker mount can continue until they exit. + +4. A newer HAMi-core with no broker available uses the existing fallback, so the library does not need to be rolled back first. + +This rollout contract applies to the broker feature. The separate PR 248 lock migration still requires its own mixed binary policy. + +## Validation required before release + +1. Go race tests for the broker, lifecycle, and allocation integration. + +2. The actual C client to Go server contract test. + +3. Missing, stale, unsafe, malformed, slow, saturated, restarting, and dying broker cases. + +4. Linux and CUDA builds for HAMi-core plus the client and context accounting tests. + +5. Kubernetes validation with the feature disabled, enabled, rolled forward, and rolled back. + +6. Concurrent `cuInit()` and first primary context benchmarks with raw output, environment details, source revisions, and checksums. + +## Known limits + +1. The design depends on Linux `SO_PEERCRED` and a device plugin in the host PID namespace. + +2. Sandboxed runtimes must be tested in real execution. A mounted Unix socket may be blocked or may not preserve the peer identity needed by this design. + +3. Existing workloads need recreation to receive or remove the allocation mount. + +4. Additional NVIDIA architectures, driver versions, CRI-O, rootless runtimes, gVisor, and Kata remain separate compatibility cells until each is tested. diff --git a/go.mod b/go.mod index f929d60442..384478bf78 100644 --- a/go.mod +++ b/go.mod @@ -23,6 +23,7 @@ require ( github.com/stretchr/testify v1.12.1 github.com/urfave/cli/v2 v2.27.7 golang.org/x/net v0.58.0 + golang.org/x/sys v0.47.0 golang.org/x/term v0.45.0 golang.org/x/tools v0.49.0 google.golang.org/grpc v1.83.1 @@ -81,7 +82,6 @@ require ( golang.org/x/mod v0.40.0 // indirect golang.org/x/oauth2 v0.36.0 // indirect golang.org/x/sync v0.22.0 // indirect - golang.org/x/sys v0.47.0 // indirect golang.org/x/telemetry v0.0.0-20260811182544-a038080d80e5 // indirect golang.org/x/text v0.41.0 // indirect golang.org/x/time v0.14.0 // indirect diff --git a/go.sum b/go.sum index b5120e9cf8..55eba36daa 100644 --- a/go.sum +++ b/go.sum @@ -143,8 +143,6 @@ github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0t github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc= github.com/russross/blackfriday/v2 v2.1.0 h1:JIOH55/0cWyOuilr9/qlrm0BSXldqnqwMsf35Ld67mk= github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= -github.com/sirupsen/logrus v1.10.0 h1:T8MxJJXVZkfcC5zSRMRAg2F8+lxjmUCGGWPzFxO+Msc= -github.com/sirupsen/logrus v1.10.0/go.mod h1:FXZFonkDAnFozmO+5hGAFvB0Yg9/j2SIhA/QuIkP180= github.com/sirupsen/logrus v1.10.1 h1:xi4336Zh11WpU14fXR6I67V3yaTPQYwRx2WEtHbRg4Q= github.com/sirupsen/logrus v1.10.1/go.mod h1:vsQHnG7xzNsxk3NrwboUiWPnIC3dmbjcGPykD7+tiHk= github.com/spf13/cobra v1.10.2 h1:DMTTonx5m65Ic0GOoRY2c16WCbHxOOw6xxezuLaBpcU= @@ -227,8 +225,6 @@ gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E= google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa h1:mZHHdPZl0dbGHCflZgAq/Q468DWVFcU2whhB2KAo8fk= google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= -google.golang.org/grpc v1.83.0 h1:JeNZEKJFbQxArAMl+hiytHauacDNqJUllNfmIMmpqnQ= -google.golang.org/grpc v1.83.0/go.mod h1:kDyl6SKsiHKt0uylY5gtn5cEjkrIOhQOGDgIc4JGwzQ= google.golang.org/grpc v1.83.1 h1:HIO0+BEtBP6soyqvqC8sNUjZ7bTs+0hFQuFF+RAy++Y= google.golang.org/grpc v1.83.1/go.mod h1:kDyl6SKsiHKt0uylY5gtn5cEjkrIOhQOGDgIc4JGwzQ= google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af h1:+5/Sw3GsDNlEmu7TfklWKPdQ0Ykja5VEmq2i817+jbI= @@ -246,30 +242,20 @@ gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gotest.tools/v3 v3.5.2 h1:7koQfIKdy+I8UTetycgUqXWSDwpgv193Ka+qRsmBY8Q= gotest.tools/v3 v3.5.2/go.mod h1:LtdLGcnqToBH83WByAAi/wiwSFCArdFIUV/xxN4pcjA= -k8s.io/api v0.36.3 h1:NxB+05W2UGqXWFXcLO0RB5cnqnUPP5v5sVlaOH0Iz4w= -k8s.io/api v0.36.3/go.mod h1:JzLQKqRHC5+I8RVj/lS3lCg0mg6nWI9Fo/Sk3ElxHzg= k8s.io/api v0.36.4 h1:RxrvqCL6vgH5/+UnTeu1IIFqYmGfy0hnyrod1rn35Oo= k8s.io/api v0.36.4/go.mod h1:S2B3orCFBDhrgyWbLeuKcT2QdHIpQesBkCYSlWtwUOw= k8s.io/apiextensions-apiserver v0.36.0 h1:Wt7E8J+VBCbj4FjiBfDTK/neXDDjyJVJc7xfuOHImZ0= k8s.io/apiextensions-apiserver v0.36.0/go.mod h1:kGDjH0msuiIB3tgsYRV0kS9GqpMYMUsQ3GHv7TApyug= -k8s.io/apimachinery v0.36.3 h1:PkzMRBRG8joFD8EhCuQAtNPvJlxb82FwplP26HIzvAM= -k8s.io/apimachinery v0.36.3/go.mod h1:cTSjBWgPe/6CQyBKzY/hDIRWCQQQeK0mfLbml0UYFHE= k8s.io/apimachinery v0.36.4 h1:PT2UzkupGuAx/+xT5XjiMJ1WGpY3fn9/hdAvjweRet4= k8s.io/apimachinery v0.36.4/go.mod h1:p2I2dipt7JHG+quVwQ1d02d28O4GdDi77RByQ13MTpk= -k8s.io/client-go v0.36.3 h1:M4JdVzXxYcZk4fGpfDdYnxSwhLKWCFoQsHW6t+z8Hfg= -k8s.io/client-go v0.36.3/go.mod h1:gcPwr0c87vjjG6HB6pWEqOeuYVoXSsREjzux2j6GF30= k8s.io/client-go v0.36.4 h1:MDvfDNvMSt0Br94SK8neviVlwL9qifw9B26hJCpD1K0= k8s.io/client-go v0.36.4/go.mod h1:pNK4WKELbwlEDvtbE8l22lEZL5THYF61H5EealokZmA= k8s.io/klog/v2 v2.140.0 h1:Tf+J3AH7xnUzZyVVXhTgGhEKnFqye14aadWv7bzXdzc= k8s.io/klog/v2 v2.140.0/go.mod h1:o+/RWfJ6PwpnFn7OyAG3QnO47BFsymfEfrz6XyYSSp0= k8s.io/kube-openapi v0.0.0-20260317180543-43fb72c5454a h1:xCeOEAOoGYl2jnJoHkC3hkbPJgdATINPMAxaynU2Ovg= k8s.io/kube-openapi v0.0.0-20260317180543-43fb72c5454a/go.mod h1:uGBT7iTA6c6MvqUvSXIaYZo9ukscABYi2btjhvgKGZ0= -k8s.io/kube-scheduler v0.36.3 h1:sc18quI2CgvH23oU1eJIQ21ivJENzjH4cxATqLSSubM= -k8s.io/kube-scheduler v0.36.3/go.mod h1:M7zaLPp1Q3S6ddqZYYLiQiGw21t5bLD0FaPRXyomP7c= k8s.io/kube-scheduler v0.36.4 h1:TC9tXGuUlPbZGeQEkTiP+3VfgaKHM8k2Y5xDNmbQ7U8= k8s.io/kube-scheduler v0.36.4/go.mod h1:flvrzp0ZHjLGKsHRTWiD8MUNS0cZDMsNpDKCHI9v9CE= -k8s.io/kubelet v0.36.3 h1:dRzEnhHk35Opy6wjWR4YBcN5RI9lB2npUY37TghFuPU= -k8s.io/kubelet v0.36.3/go.mod h1:4USFGr21Ioka+b964Beq0NvV5b5aca3RWJ1/kfq+RLw= k8s.io/kubelet v0.36.4 h1:mlmXnkrq3H02r/r0H/8M2jdPY7f4I4u4cA0tHnsPzY0= k8s.io/kubelet v0.36.4/go.mod h1:jcOhk4E8cdUBn7WswW67WH9waQTe37G057ttnYdcaKY= k8s.io/utils v0.0.0-20260210185600-b8788abfbbc2 h1:AZYQSJemyQB5eRxqcPky+/7EdBj0xi3g0ZcxxJ7vbWU= diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/hostpid/broker_linux.go b/pkg/device-plugin/nvidiadevice/nvinternal/hostpid/broker_linux.go new file mode 100644 index 0000000000..1b8e81b0da --- /dev/null +++ b/pkg/device-plugin/nvidiadevice/nvinternal/hostpid/broker_linux.go @@ -0,0 +1,396 @@ +//go:build linux + +/* + * SPDX-License-Identifier: Apache-2.0 + * + * Copyright (c) 2026 The HAMi Authors. + */ + +package hostpid + +import ( + "errors" + "fmt" + "io" + "net" + "os" + "path/filepath" + "sync" + "sync/atomic" + "syscall" + "time" + + "golang.org/x/sys/unix" + "k8s.io/klog/v2" +) + +const ( + transactionTimeout = 500 * time.Millisecond + activeProbeTimeout = 50 * time.Millisecond + acceptRetryInitial = 5 * time.Millisecond + acceptRetryMaximum = time.Second + maxHandlers = 512 + serverDirectoryMode = 0o711 + serverSocketMode = 0o666 + serverLockMode = 0o600 +) + +type socketIdentity struct { + device uint64 + inode uint64 +} + +type Broker struct { + listener *net.UnixListener + socketPath string + socket socketIdentity + lockFile *os.File + handlerSlots chan struct{} + handlerWG sync.WaitGroup + handlerMu sync.Mutex + closing atomic.Bool + dropped atomic.Uint64 + closeOnce sync.Once + closeErr error +} + +func ListenDefault() (*Broker, error) { + if os.Geteuid() != 0 { + return nil, errors.New("the host PID broker must run as root") + } + return listen(ServerSocketPath, 0) +} + +func listen(socketPath string, ownerUID int) (*Broker, error) { + return listenWithHandlerLimit(socketPath, ownerUID, maxHandlers) +} + +func listenWithHandlerLimit(socketPath string, ownerUID int, + handlerLimit int) (*Broker, error) { + if handlerLimit <= 0 { + return nil, errors.New("host PID broker handler limit must be positive") + } + directory := filepath.Dir(socketPath) + if err := prepareDirectory(directory, ownerUID); err != nil { + return nil, err + } + + lockFile, err := acquireLock(directory, ownerUID) + if err != nil { + return nil, err + } + releaseLock := true + defer func() { + if releaseLock { + _ = unix.Flock(int(lockFile.Fd()), unix.LOCK_UN) + _ = lockFile.Close() + } + }() + + if err := removeStaleSocket(socketPath, ownerUID); err != nil { + return nil, err + } + address := &net.UnixAddr{Name: socketPath, Net: "unix"} + listener, err := net.ListenUnix("unix", address) + if err != nil { + return nil, fmt.Errorf("listen on host PID broker socket: %w", err) + } + listener.SetUnlinkOnClose(false) + cleanupListener := true + defer func() { + if cleanupListener { + _ = listener.Close() + _ = os.Remove(socketPath) + } + }() + + if err := os.Chmod(socketPath, serverSocketMode); err != nil { + return nil, fmt.Errorf("set host PID broker socket mode: %w", err) + } + identity, err := readSocketIdentity(socketPath, ownerUID) + if err != nil { + return nil, err + } + + cleanupListener = false + releaseLock = false + return &Broker{ + listener: listener, + socketPath: socketPath, + socket: identity, + lockFile: lockFile, + handlerSlots: make(chan struct{}, handlerLimit), + }, nil +} + +func prepareDirectory(directory string, ownerUID int) error { + if err := os.MkdirAll(directory, serverDirectoryMode); err != nil { + return fmt.Errorf("create host PID broker directory: %w", err) + } + info, err := os.Lstat(directory) + if err != nil { + return fmt.Errorf("inspect host PID broker directory: %w", err) + } + stat, ok := info.Sys().(*syscall.Stat_t) + if !ok || !info.IsDir() || info.Mode()&os.ModeSymlink != 0 { + return errors.New("host PID broker directory is not a real directory") + } + if int(stat.Uid) != ownerUID { + return fmt.Errorf("host PID broker directory owner is %d, want %d", + stat.Uid, ownerUID) + } + if err := os.Chmod(directory, serverDirectoryMode); err != nil { + return fmt.Errorf("set host PID broker directory mode: %w", err) + } + return nil +} + +func acquireLock(directory string, ownerUID int) (*os.File, error) { + lockPath := filepath.Join(directory, "broker.lock") + fd, err := unix.Open(lockPath, + unix.O_CREAT|unix.O_RDWR|unix.O_CLOEXEC|unix.O_NOFOLLOW, + serverLockMode) + if err != nil { + return nil, fmt.Errorf("open host PID broker lock: %w", err) + } + lockFile := os.NewFile(uintptr(fd), lockPath) + if lockFile == nil { + _ = unix.Close(fd) + return nil, errors.New("open host PID broker lock file") + } + if err := unix.Flock(fd, unix.LOCK_EX|unix.LOCK_NB); err != nil { + _ = lockFile.Close() + return nil, fmt.Errorf("lock host PID broker directory: %w", err) + } + + info, err := lockFile.Stat() + if err != nil { + _ = unix.Flock(fd, unix.LOCK_UN) + _ = lockFile.Close() + return nil, fmt.Errorf("inspect host PID broker lock: %w", err) + } + stat, ok := info.Sys().(*syscall.Stat_t) + if !ok || !info.Mode().IsRegular() || int(stat.Uid) != ownerUID || + info.Mode().Perm()&0o077 != 0 { + _ = unix.Flock(fd, unix.LOCK_UN) + _ = lockFile.Close() + return nil, errors.New("host PID broker lock is not trusted") + } + return lockFile, nil +} + +func removeStaleSocket(socketPath string, ownerUID int) error { + info, err := os.Lstat(socketPath) + if errors.Is(err, os.ErrNotExist) { + return nil + } + if err != nil { + return fmt.Errorf("inspect host PID broker socket: %w", err) + } + stat, ok := info.Sys().(*syscall.Stat_t) + if !ok || info.Mode()&os.ModeSocket == 0 { + return errors.New("host PID broker socket path is not a socket") + } + if int(stat.Uid) != ownerUID { + return fmt.Errorf("host PID broker socket owner is %d, want %d", + stat.Uid, ownerUID) + } + + connection, dialErr := net.DialTimeout("unix", socketPath, + activeProbeTimeout) + if dialErr == nil { + _ = connection.Close() + return errors.New("another host PID broker is already listening") + } + if !errors.Is(dialErr, syscall.ECONNREFUSED) && + !errors.Is(dialErr, os.ErrNotExist) { + return fmt.Errorf("probe existing host PID broker socket: %w", dialErr) + } + if err := os.Remove(socketPath); err != nil && + !errors.Is(err, os.ErrNotExist) { + return fmt.Errorf("remove stale host PID broker socket: %w", err) + } + return nil +} + +func readSocketIdentity(socketPath string, ownerUID int) (socketIdentity, error) { + info, err := os.Lstat(socketPath) + if err != nil { + return socketIdentity{}, fmt.Errorf( + "inspect new host PID broker socket: %w", err) + } + stat, ok := info.Sys().(*syscall.Stat_t) + if !ok || info.Mode()&os.ModeSocket == 0 || int(stat.Uid) != ownerUID { + return socketIdentity{}, errors.New("new host PID broker socket is not trusted") + } + return socketIdentity{device: uint64(stat.Dev), inode: stat.Ino}, nil +} + +func (broker *Broker) Serve() error { + var backoff time.Duration + for { + connection, err := broker.listener.AcceptUnix() + if err != nil { + if broker.closing.Load() { + return nil + } + if isTemporaryAcceptError(err) { + backoff = nextAcceptBackoff(backoff) + time.Sleep(backoff) + continue + } + return fmt.Errorf("accept host PID broker connection: %w", err) + } + backoff = 0 + if !broker.startHandler(connection) { + _ = connection.Close() + } + } +} + +func isTemporaryAcceptError(err error) bool { + return errors.Is(err, syscall.EMFILE) || + errors.Is(err, syscall.ENFILE) || + errors.Is(err, syscall.ENOBUFS) || + errors.Is(err, syscall.ENOMEM) +} + +func nextAcceptBackoff(current time.Duration) time.Duration { + if current == 0 { + return acceptRetryInitial + } + next := current * 2 + if next > acceptRetryMaximum { + return acceptRetryMaximum + } + return next +} + +func (broker *Broker) startHandler(connection *net.UnixConn) bool { + broker.handlerMu.Lock() + defer broker.handlerMu.Unlock() + if broker.closing.Load() { + return false + } + select { + case broker.handlerSlots <- struct{}{}: + broker.handlerWG.Add(1) + go broker.handle(connection) + return true + default: + return false + } +} + +func (broker *Broker) handle(connection *net.UnixConn) { + defer func() { + _ = connection.Close() + <-broker.handlerSlots + broker.handlerWG.Done() + }() + _ = connection.SetDeadline(time.Now().Add(transactionTimeout)) + + request := make([]byte, requestSize) + if _, err := io.ReadFull(connection, request); err != nil { + broker.logDroppedTransaction("request read", err) + return + } + if !validRequest(request) { + response := makeResponse(statusInvalidRequest, 0) + writeResponse(connection, response) + return + } + + pid, err := peerPID(connection) + if err != nil { + broker.logDroppedTransaction("peer credentials", err) + return + } + if pid <= 0 { + return + } + response := makeResponse(statusOK, uint32(pid)) + writeResponse(connection, response) +} + +func (broker *Broker) logDroppedTransaction(operation string, err error) { + count := broker.dropped.Add(1) + if count&(count-1) != 0 { + return + } + klog.V(4).Infof( + "Dropped host PID broker transaction during %s (count=%d): %v", + operation, count, err) +} + +func writeResponse(connection *net.UnixConn, + response [responseSize]byte) { + written := 0 + for written < len(response) { + count, err := connection.Write(response[written:]) + if err != nil || count == 0 { + return + } + written += count + } +} + +func peerPID(connection *net.UnixConn) (int32, error) { + rawConnection, err := connection.SyscallConn() + if err != nil { + return 0, err + } + var credentials *unix.Ucred + var credentialErr error + if err := rawConnection.Control(func(fd uintptr) { + credentials, credentialErr = unix.GetsockoptUcred( + int(fd), unix.SOL_SOCKET, unix.SO_PEERCRED) + }); err != nil { + return 0, err + } + if credentialErr != nil { + return 0, credentialErr + } + if credentials == nil || credentials.Pid <= 0 { + return 0, errors.New("host PID broker received invalid peer credentials") + } + return credentials.Pid, nil +} + +func (broker *Broker) Close() error { + broker.closeOnce.Do(func() { + broker.handlerMu.Lock() + broker.closing.Store(true) + broker.handlerMu.Unlock() + + listenerErr := broker.listener.Close() + broker.handlerWG.Wait() + removeErr := broker.removeOwnedSocket() + unlockErr := unix.Flock(int(broker.lockFile.Fd()), unix.LOCK_UN) + lockCloseErr := broker.lockFile.Close() + broker.closeErr = errors.Join(listenerErr, removeErr, unlockErr, + lockCloseErr) + }) + return broker.closeErr +} + +func (broker *Broker) removeOwnedSocket() error { + var stat unix.Stat_t + if err := unix.Lstat(broker.socketPath, &stat); err != nil { + if errors.Is(err, os.ErrNotExist) { + return nil + } + return fmt.Errorf("inspect host PID broker socket during cleanup: %w", + err) + } + if uint64(stat.Dev) != broker.socket.device || + stat.Ino != broker.socket.inode || + stat.Mode&unix.S_IFMT != unix.S_IFSOCK { + return nil + } + if err := unix.Unlink(broker.socketPath); err != nil && + !errors.Is(err, os.ErrNotExist) { + return fmt.Errorf("remove host PID broker socket: %w", err) + } + return nil +} diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/hostpid/broker_linux_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/hostpid/broker_linux_test.go new file mode 100644 index 0000000000..e9cdf62038 --- /dev/null +++ b/pkg/device-plugin/nvidiadevice/nvinternal/hostpid/broker_linux_test.go @@ -0,0 +1,625 @@ +//go:build linux + +/* + * SPDX-License-Identifier: Apache-2.0 + * + * Copyright (c) 2026 The HAMi Authors. + */ + +package hostpid + +import ( + "encoding/binary" + "errors" + "fmt" + "io" + "net" + "os" + "os/exec" + "path/filepath" + "strconv" + "strings" + "sync" + "syscall" + "testing" + "time" +) + +const subprocessHelperEnvironment = "HAMI_HOSTPID_BROKER_HELPER" + +const externalCClientEnvironment = "HAMI_HOSTPID_C_CLIENT" + +const externalCStormEnvironment = "HAMI_HOSTPID_C_STORM" + +func startTestBroker(t *testing.T) (*Broker, string) { + t.Helper() + directory := t.TempDir() + socketPath := filepath.Join(directory, "broker.sock") + broker, err := listen(socketPath, os.Geteuid()) + if err != nil { + t.Fatalf("listen: %v", err) + } + serveResult := make(chan error, 1) + go func() { + serveResult <- broker.Serve() + }() + t.Cleanup(func() { + if err := broker.Close(); err != nil { + t.Errorf("close broker: %v", err) + } + if err := <-serveResult; err != nil { + t.Errorf("serve broker: %v", err) + } + }) + return broker, socketPath +} + +func queryBroker(socketPath string) (uint16, uint32, error) { + connection, err := net.DialTimeout("unix", socketPath, time.Second) + if err != nil { + return 0, 0, err + } + defer connection.Close() + if err := connection.SetDeadline(time.Now().Add(time.Second)); err != nil { + return 0, 0, err + } + request := []byte{'H', 'P', 'I', 'D', 0, 1, 0, 1} + if _, err := connection.Write(request); err != nil { + return 0, 0, err + } + response := make([]byte, responseSize) + if _, err := io.ReadFull(connection, response); err != nil { + return 0, 0, err + } + if string(response[:4]) != "HPID" || + binary.BigEndian.Uint16(response[4:6]) != protocolVersion { + return 0, 0, errors.New("invalid broker response") + } + return binary.BigEndian.Uint16(response[6:8]), + binary.BigEndian.Uint32(response[8:12]), nil +} + +func TestBrokerReturnsPeerPID(t *testing.T) { + _, socketPath := startTestBroker(t) + status, pid, err := queryBroker(socketPath) + if err != nil { + t.Fatal(err) + } + if status != statusOK || pid != uint32(os.Getpid()) { + t.Fatalf("got status %d PID %d, want status 0 PID %d", + status, pid, os.Getpid()) + } +} + +func TestBrokerReturnsSubprocessPID(t *testing.T) { + _, socketPath := startTestBroker(t) + command := exec.Command(os.Args[0], + "-test.run=^TestBrokerSubprocessHelper$") + command.Env = append(os.Environ(), + subprocessHelperEnvironment+"="+socketPath) + output, err := command.CombinedOutput() + if err != nil { + t.Fatalf("helper failed: %v\n%s", err, output) + } +} + +func TestBrokerExternalCClient(t *testing.T) { + clientPath := os.Getenv(externalCClientEnvironment) + if clientPath == "" { + t.Skip("HAMI_HOSTPID_C_CLIENT is not set") + } + _, socketPath := startTestBroker(t) + command := exec.Command(clientPath, socketPath) + output, err := command.CombinedOutput() + if err != nil { + t.Fatalf("external C client failed: %v\n%s", err, output) + } +} + +func TestBrokerExternalCClientStorm(t *testing.T) { + clientPath := os.Getenv(externalCStormEnvironment) + if clientPath == "" { + t.Skip("HAMI_HOSTPID_C_STORM is not set") + } + _, socketPath := startTestBroker(t) + command := exec.Command(clientPath, socketPath, "300") + output, err := command.CombinedOutput() + if err != nil { + t.Fatalf("external C client storm failed: %v\n%s", err, output) + } + if !strings.Contains(string(output), + "clients=300 successful=300 failed=0") { + t.Fatalf("unexpected C client storm output: %s", output) + } +} + +func TestBrokerSubprocessHelper(t *testing.T) { + socketPath := os.Getenv(subprocessHelperEnvironment) + if socketPath == "" { + return + } + status, pid, err := queryBroker(socketPath) + if err != nil { + t.Fatal(err) + } + if status != statusOK || pid != uint32(os.Getpid()) { + t.Fatalf("got status %d PID %d, want status 0 PID %d", + status, pid, os.Getpid()) + } +} + +func TestBrokerRejectsInvalidRequest(t *testing.T) { + _, socketPath := startTestBroker(t) + tests := map[string][]byte{ + "magic": {'B', 'A', 'D', '!', 0, 1, 0, 1}, + "version": {'H', 'P', 'I', 'D', 0, 2, 0, 1}, + "command": {'H', 'P', 'I', 'D', 0, 1, 0, 2}, + } + for name, request := range tests { + t.Run(name, func(t *testing.T) { + connection, err := net.Dial("unix", socketPath) + if err != nil { + t.Fatal(err) + } + defer connection.Close() + if _, err := connection.Write(request); err != nil { + t.Fatal(err) + } + response := make([]byte, responseSize) + if _, err := io.ReadFull(connection, response); err != nil { + t.Fatal(err) + } + if status := binary.BigEndian.Uint16(response[6:8]); status != statusInvalidRequest { + t.Fatalf("got status %d", status) + } + if pid := binary.BigEndian.Uint32(response[8:12]); pid != 0 { + t.Fatalf("got PID %d", pid) + } + }) + } +} + +func TestBrokerRecoversFromEarlyClose(t *testing.T) { + _, socketPath := startTestBroker(t) + connection, err := net.Dial("unix", socketPath) + if err != nil { + t.Fatal(err) + } + if err := connection.Close(); err != nil { + t.Fatal(err) + } + + status, pid, err := queryBroker(socketPath) + if err != nil || status != statusOK || pid != uint32(os.Getpid()) { + t.Fatalf("broker did not recover: status=%d pid=%d err=%v", + status, pid, err) + } +} + +func TestTemporaryAcceptErrors(t *testing.T) { + for _, acceptErr := range []error{ + syscall.EMFILE, + syscall.ENFILE, + syscall.ENOBUFS, + syscall.ENOMEM, + fmt.Errorf("wrapped: %w", syscall.EMFILE), + } { + if !isTemporaryAcceptError(acceptErr) { + t.Errorf("expected temporary accept error: %v", acceptErr) + } + } + if isTemporaryAcceptError(syscall.EINVAL) { + t.Fatal("EINVAL must remain a permanent accept error") + } +} + +func TestAcceptBackoffIsBounded(t *testing.T) { + backoff := time.Duration(0) + for range 32 { + backoff = nextAcceptBackoff(backoff) + if backoff > acceptRetryMaximum { + t.Fatalf("backoff %v exceeds maximum %v", backoff, + acceptRetryMaximum) + } + } + if backoff != acceptRetryMaximum { + t.Fatalf("backoff=%v, want %v", backoff, acceptRetryMaximum) + } +} + +func TestBrokerTimesOutPartialRequest(t *testing.T) { + _, socketPath := startTestBroker(t) + connection, err := net.Dial("unix", socketPath) + if err != nil { + t.Fatal(err) + } + if _, err := connection.Write([]byte{'H', 'P', 'I', 'D'}); err != nil { + t.Fatal(err) + } + time.Sleep(transactionTimeout + 100*time.Millisecond) + if err := connection.SetReadDeadline(time.Now().Add(time.Second)); err != nil { + t.Fatal(err) + } + buffer := make([]byte, 1) + if count, err := connection.Read(buffer); count != 0 || err == nil { + t.Fatalf("partial request connection stayed open: n=%d err=%v", + count, err) + } + _ = connection.Close() + + status, pid, err := queryBroker(socketPath) + if err != nil || status != statusOK || pid != uint32(os.Getpid()) { + t.Fatalf("broker did not recover: status=%d pid=%d err=%v", + status, pid, err) + } +} + +func TestBrokerHandlesConcurrentClients(t *testing.T) { + _, socketPath := startTestBroker(t) + const clients = 300 + errorsChannel := make(chan error, clients) + var waitGroup sync.WaitGroup + + for range clients { + waitGroup.Add(1) + go func() { + defer waitGroup.Done() + status, pid, err := queryBroker(socketPath) + if err != nil { + errorsChannel <- err + return + } + if status != statusOK || pid != uint32(os.Getpid()) { + errorsChannel <- fmt.Errorf("status=%d pid=%d", status, pid) + } + }() + } + waitGroup.Wait() + close(errorsChannel) + for err := range errorsChannel { + t.Error(err) + } +} + +func TestBrokerBoundsSlowClientsAndRecovers(t *testing.T) { + directory := t.TempDir() + socketPath := filepath.Join(directory, "broker.sock") + const handlerLimit = 4 + broker, err := listenWithHandlerLimit(socketPath, os.Geteuid(), + handlerLimit) + if err != nil { + t.Fatal(err) + } + serveResult := make(chan error, 1) + go func() { serveResult <- broker.Serve() }() + t.Cleanup(func() { + if err := broker.Close(); err != nil { + t.Errorf("close broker: %v", err) + } + if err := <-serveResult; err != nil { + t.Errorf("serve broker: %v", err) + } + }) + + connections := make([]net.Conn, 0, handlerLimit) + for range handlerLimit { + connection, err := net.Dial("unix", socketPath) + if err != nil { + t.Fatal(err) + } + connections = append(connections, connection) + } + deadline := time.Now().Add(time.Second) + for len(broker.handlerSlots) != handlerLimit && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if got := len(broker.handlerSlots); got != handlerLimit { + t.Fatalf("active handlers=%d, want %d", got, handlerLimit) + } + + overflow, err := net.Dial("unix", socketPath) + if err != nil { + t.Fatal(err) + } + if err := overflow.SetDeadline(time.Now().Add(transactionTimeout)); err != nil { + t.Fatal(err) + } + request := []byte{'H', 'P', 'I', 'D', 0, 1, 0, 1} + _, _ = overflow.Write(request) + response := make([]byte, responseSize) + if _, err := io.ReadFull(overflow, response); err == nil { + t.Fatal("overflow client received a response") + } + _ = overflow.Close() + + for _, connection := range connections { + _ = connection.Close() + } + deadline = time.Now().Add(2 * transactionTimeout) + for { + status, pid, queryErr := queryBroker(socketPath) + if queryErr == nil && status == statusOK && pid == uint32(os.Getpid()) { + break + } + if time.Now().After(deadline) { + t.Fatalf("broker did not recover: status=%d pid=%d err=%v", + status, pid, queryErr) + } + time.Sleep(time.Millisecond) + } +} + +func TestBrokerCloseBoundsPartialClient(t *testing.T) { + directory := t.TempDir() + socketPath := filepath.Join(directory, "broker.sock") + broker, err := listenWithHandlerLimit(socketPath, os.Geteuid(), 1) + if err != nil { + t.Fatal(err) + } + serveResult := make(chan error, 1) + go func() { serveResult <- broker.Serve() }() + + connection, err := net.Dial("unix", socketPath) + if err != nil { + t.Fatal(err) + } + if _, err := connection.Write([]byte{'H'}); err != nil { + t.Fatal(err) + } + deadline := time.Now().Add(time.Second) + for len(broker.handlerSlots) != 1 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + + begin := time.Now() + if err := broker.Close(); err != nil { + t.Fatal(err) + } + if elapsed := time.Since(begin); elapsed > 2*transactionTimeout { + t.Fatalf("close took %s", elapsed) + } + if err := <-serveResult; err != nil { + t.Fatal(err) + } + _ = connection.Close() +} + +func TestBrokerRestartsAfterClose(t *testing.T) { + directory := t.TempDir() + socketPath := filepath.Join(directory, "broker.sock") + + for attempt := range 2 { + broker, err := listen(socketPath, os.Geteuid()) + if err != nil { + t.Fatalf("listen attempt %d: %v", attempt, err) + } + serveResult := make(chan error, 1) + go func() { serveResult <- broker.Serve() }() + status, pid, err := queryBroker(socketPath) + if err != nil || status != statusOK || pid != uint32(os.Getpid()) { + t.Fatalf("query attempt %d: status=%d pid=%d err=%v", + attempt, status, pid, err) + } + if err := broker.Close(); err != nil { + t.Fatalf("close attempt %d: %v", attempt, err) + } + if err := <-serveResult; err != nil { + t.Fatalf("serve attempt %d: %v", attempt, err) + } + } +} + +func TestBrokerCreatesTrustedModes(t *testing.T) { + _, socketPath := startTestBroker(t) + directoryInfo, err := os.Stat(filepath.Dir(socketPath)) + if err != nil { + t.Fatal(err) + } + socketInfo, err := os.Stat(socketPath) + if err != nil { + t.Fatal(err) + } + if got := directoryInfo.Mode().Perm(); got != serverDirectoryMode { + t.Fatalf("directory mode is %#o", got) + } + if got := socketInfo.Mode().Perm(); got != serverSocketMode { + t.Fatalf("socket mode is %#o", got) + } +} + +func TestBrokerRejectsPathCollision(t *testing.T) { + directory := t.TempDir() + socketPath := filepath.Join(directory, "broker.sock") + if err := os.WriteFile(socketPath, []byte("keep"), 0o600); err != nil { + t.Fatal(err) + } + if broker, err := listen(socketPath, os.Geteuid()); err == nil { + _ = broker.Close() + t.Fatal("broker accepted a regular file collision") + } + contents, err := os.ReadFile(socketPath) + if err != nil || string(contents) != "keep" { + t.Fatalf("collision was changed: contents=%q err=%v", contents, err) + } +} + +func TestBrokerRemovesStaleSocket(t *testing.T) { + directory := t.TempDir() + socketPath := filepath.Join(directory, "broker.sock") + listener, err := net.ListenUnix("unix", &net.UnixAddr{ + Name: socketPath, + Net: "unix", + }) + if err != nil { + t.Fatal(err) + } + listener.SetUnlinkOnClose(false) + if err := listener.Close(); err != nil { + t.Fatal(err) + } + + broker, err := listen(socketPath, os.Geteuid()) + if err != nil { + t.Fatal(err) + } + if err := broker.Close(); err != nil { + t.Fatal(err) + } +} + +func TestBrokerDoesNotRemoveActiveSocket(t *testing.T) { + directory := t.TempDir() + socketPath := filepath.Join(directory, "broker.sock") + listener, err := net.ListenUnix("unix", &net.UnixAddr{ + Name: socketPath, + Net: "unix", + }) + if err != nil { + t.Fatal(err) + } + defer listener.Close() + + broker, err := listen(socketPath, os.Geteuid()) + if err == nil { + _ = broker.Close() + t.Fatal("broker replaced an active socket") + } + if _, err := os.Lstat(socketPath); err != nil { + t.Fatalf("active socket was removed: %v", err) + } +} + +func TestBrokerRejectsSymlinkDirectory(t *testing.T) { + parent := t.TempDir() + target := filepath.Join(parent, "target") + link := filepath.Join(parent, "link") + if err := os.Mkdir(target, 0o700); err != nil { + t.Fatal(err) + } + if err := os.Symlink(target, link); err != nil { + t.Fatal(err) + } + broker, err := listen(filepath.Join(link, "broker.sock"), os.Geteuid()) + if err == nil { + _ = broker.Close() + t.Fatal("broker accepted a symlink directory") + } +} + +func TestBrokerRejectsSymlinkLock(t *testing.T) { + directory := t.TempDir() + target := filepath.Join(directory, "target") + lockPath := filepath.Join(directory, "broker.lock") + if err := os.WriteFile(target, []byte("keep"), 0o600); err != nil { + t.Fatal(err) + } + if err := os.Symlink(target, lockPath); err != nil { + t.Fatal(err) + } + broker, err := listen(filepath.Join(directory, "broker.sock"), + os.Geteuid()) + if err == nil { + _ = broker.Close() + t.Fatal("broker accepted a symlink lock") + } + contents, err := os.ReadFile(target) + if err != nil || string(contents) != "keep" { + t.Fatalf("lock target was changed: contents=%q err=%v", contents, err) + } +} + +func TestBrokerRejectsSecondListener(t *testing.T) { + broker, socketPath := startTestBroker(t) + second, err := listen(socketPath, os.Geteuid()) + if err == nil { + _ = second.Close() + t.Fatal("second broker acquired the socket") + } + status, pid, queryErr := queryBroker(socketPath) + if queryErr != nil || status != statusOK || pid != uint32(os.Getpid()) { + t.Fatalf("first broker stopped: status=%d pid=%d err=%v", + status, pid, queryErr) + } + _ = broker +} + +func TestBrokerLeavesReplacementDuringClose(t *testing.T) { + directory := t.TempDir() + socketPath := filepath.Join(directory, "broker.sock") + oldSocketPath := filepath.Join(directory, "old.sock") + broker, err := listen(socketPath, os.Geteuid()) + if err != nil { + t.Fatal(err) + } + serveResult := make(chan error, 1) + go func() { serveResult <- broker.Serve() }() + + if err := os.Rename(socketPath, oldSocketPath); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(socketPath, []byte("replacement"), 0o600); err != nil { + t.Fatal(err) + } + if err := broker.Close(); err != nil { + t.Fatal(err) + } + if err := <-serveResult; err != nil { + t.Fatal(err) + } + contents, err := os.ReadFile(socketPath) + if err != nil || string(contents) != "replacement" { + t.Fatalf("replacement was changed: contents=%q err=%v", contents, err) + } + if err := os.Remove(socketPath); err != nil { + t.Fatal(err) + } + if err := os.Remove(oldSocketPath); err != nil { + t.Fatal(err) + } +} + +func TestEnabled(t *testing.T) { + tests := map[string]bool{ + "": false, + "0": false, + "1": true, + "true": false, + "false": false, + "01": false, + " 1": false, + } + for value, expected := range tests { + t.Run(strconv.Quote(value), func(t *testing.T) { + if got := Enabled(value); got != expected { + t.Fatalf("Enabled(%q)=%v, want %v", value, got, expected) + } + }) + } +} + +func TestListenDefaultRequiresRoot(t *testing.T) { + if os.Geteuid() == 0 { + t.Skip("test requires a non-root process") + } + broker, err := ListenDefault() + if broker != nil || err == nil { + t.Fatalf("broker=%v err=%v", broker, err) + } +} + +func TestBrokerSocketIdentityUsesDeviceAndInode(t *testing.T) { + broker, socketPath := startTestBroker(t) + info, err := os.Lstat(socketPath) + if err != nil { + t.Fatal(err) + } + stat, ok := info.Sys().(*syscall.Stat_t) + if !ok || stat.Ino == 0 { + t.Fatalf("invalid socket stat: %#v", info.Sys()) + } + if uint64(stat.Dev) != broker.socket.device || + stat.Ino != broker.socket.inode { + t.Fatalf("recorded identity dev=%d ino=%d, want dev=%d ino=%d", + broker.socket.device, broker.socket.inode, stat.Dev, stat.Ino) + } +} diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/hostpid/broker_unsupported.go b/pkg/device-plugin/nvidiadevice/nvinternal/hostpid/broker_unsupported.go new file mode 100644 index 0000000000..b6694d8595 --- /dev/null +++ b/pkg/device-plugin/nvidiadevice/nvinternal/hostpid/broker_unsupported.go @@ -0,0 +1,27 @@ +//go:build !linux + +/* + * SPDX-License-Identifier: Apache-2.0 + * + * Copyright (c) 2026 The HAMi Authors. + */ + +package hostpid + +import "errors" + +var errUnsupported = errors.New("the host PID broker requires Linux") + +type Broker struct{} + +func ListenDefault() (*Broker, error) { + return nil, errUnsupported +} + +func (broker *Broker) Serve() error { + return errUnsupported +} + +func (broker *Broker) Close() error { + return nil +} diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/hostpid/config.go b/pkg/device-plugin/nvidiadevice/nvinternal/hostpid/config.go new file mode 100644 index 0000000000..6421b59f33 --- /dev/null +++ b/pkg/device-plugin/nvidiadevice/nvinternal/hostpid/config.go @@ -0,0 +1,21 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * Copyright (c) 2026 The HAMi Authors. + */ + +package hostpid + +const ( + EnvironmentVariable = "LIBVGPU_HOSTPID_BROKER" + + ServerDirectory = "/var/run/hami/hostpid" + ServerSocketPath = ServerDirectory + "/broker.sock" + + ContainerDirectory = "/tmp/vgpulock/hostpid" + ContainerSocketPath = ContainerDirectory + "/broker.sock" +) + +func Enabled(value string) bool { + return value == "1" +} diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/hostpid/protocol.go b/pkg/device-plugin/nvidiadevice/nvinternal/hostpid/protocol.go new file mode 100644 index 0000000000..33e5bc3313 --- /dev/null +++ b/pkg/device-plugin/nvidiadevice/nvinternal/hostpid/protocol.go @@ -0,0 +1,39 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * Copyright (c) 2026 The HAMi Authors. + */ + +package hostpid + +import "encoding/binary" + +const ( + protocolVersion uint16 = 1 + commandGetPID uint16 = 1 + + statusOK uint16 = 0 + statusInvalidRequest uint16 = 1 + + requestSize = 8 + responseSize = 12 +) + +var protocolMagic = [4]byte{'H', 'P', 'I', 'D'} + +func validRequest(request []byte) bool { + return len(request) == requestSize && + string(request[:4]) == string(protocolMagic[:]) && + binary.BigEndian.Uint16(request[4:6]) == protocolVersion && + binary.BigEndian.Uint16(request[6:8]) == commandGetPID +} + +func makeResponse(status uint16, pid uint32) [responseSize]byte { + var response [responseSize]byte + + copy(response[:4], protocolMagic[:]) + binary.BigEndian.PutUint16(response[4:6], protocolVersion) + binary.BigEndian.PutUint16(response[6:8], status) + binary.BigEndian.PutUint32(response[8:12], pid) + return response +} diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/hostpid/protocol_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/hostpid/protocol_test.go new file mode 100644 index 0000000000..96c8866e5c --- /dev/null +++ b/pkg/device-plugin/nvidiadevice/nvinternal/hostpid/protocol_test.go @@ -0,0 +1,62 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * Copyright (c) 2026 The HAMi Authors. + */ + +package hostpid + +import ( + "encoding/binary" + "testing" +) + +func TestValidRequest(t *testing.T) { + valid := []byte{'H', 'P', 'I', 'D', 0, 1, 0, 1} + if !validRequest(valid) { + t.Fatal("valid request was rejected") + } + + tests := map[string][]byte{ + "short": valid[:7], + "long": append(append([]byte{}, valid...), 0), + "magic": {'B', 'A', 'D', '!', 0, 1, 0, 1}, + "version": {'H', 'P', 'I', 'D', 0, 2, 0, 1}, + "command": {'H', 'P', 'I', 'D', 0, 1, 0, 2}, + "zero command": {'H', 'P', 'I', 'D', 0, 1, 0, 0}, + } + for name, request := range tests { + t.Run(name, func(t *testing.T) { + if validRequest(request) { + t.Fatal("invalid request was accepted") + } + }) + } +} + +func TestMakeResponse(t *testing.T) { + response := makeResponse(statusOK, 0x01020304) + + if string(response[:4]) != "HPID" { + t.Fatalf("unexpected magic %q", response[:4]) + } + if got := binary.BigEndian.Uint16(response[4:6]); got != protocolVersion { + t.Fatalf("unexpected version %d", got) + } + if got := binary.BigEndian.Uint16(response[6:8]); got != statusOK { + t.Fatalf("unexpected status %d", got) + } + if got := binary.BigEndian.Uint32(response[8:12]); got != 0x01020304 { + t.Fatalf("unexpected PID %#x", got) + } +} + +func TestMakeErrorResponse(t *testing.T) { + response := makeResponse(statusInvalidRequest, 0) + if got := binary.BigEndian.Uint16(response[6:8]); got != statusInvalidRequest { + t.Fatalf("unexpected status %d", got) + } + if got := binary.BigEndian.Uint32(response[8:12]); got != 0 { + t.Fatalf("unexpected PID %d", got) + } +} diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/hostpid_broker.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/hostpid_broker.go new file mode 100644 index 0000000000..5de2ba684f --- /dev/null +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/hostpid_broker.go @@ -0,0 +1,243 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * Copyright (c) 2026 The HAMi Authors. + */ + +package plugin + +import ( + "fmt" + "os" + "path/filepath" + "syscall" + + "golang.org/x/sys/unix" + kubeletdevicepluginv1beta1 "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1" + + "github.com/Project-HAMi/HAMi/pkg/device-plugin/nvidiadevice/nvinternal/hostpid" +) + +const hostPIDLockParentDirectory = "/tmp/vgpulock" + +const hostPIDLockParentMode = os.FileMode(0o777) | os.ModeSticky + +const hostPIDLockParentCreateMode = uint32(0o1777) + +const hostPIDLockTrustedOwner = uint32(0) + +var prepareHostPIDLockParentForAllocation = prepareDefaultHostPIDLockParent + +func prepareDefaultHostPIDLockParent() error { + return prepareHostPIDLockParent(hostPIDLockParentDirectory, + hostPIDLockTrustedOwner) +} + +func createHostPIDLockParent(parentFD int, baseName string) error { + return createHostPIDLockParentWith(unix.Mkdirat, parentFD, baseName) +} + +func createHostPIDLockParentWith( + mkdirat func(int, string, uint32) error, + parentFD int, baseName string) error { + err := mkdirat(parentFD, baseName, hostPIDLockParentCreateMode) + if err != nil && err != unix.EEXIST { + return err + } + return nil +} + +func prepareHostPIDLockParent(directory string, trustedOwner uint32) error { + cleanDirectory := filepath.Clean(directory) + if !filepath.IsAbs(cleanDirectory) || + cleanDirectory == string(filepath.Separator) { + return fmt.Errorf("directory must be an absolute non-root path") + } + parentDirectory := filepath.Dir(cleanDirectory) + baseName := filepath.Base(cleanDirectory) + parentFD, err := unix.Open(parentDirectory, + unix.O_RDONLY|unix.O_DIRECTORY|unix.O_CLOEXEC|unix.O_NOFOLLOW, 0) + if err != nil { + return fmt.Errorf("open parent directory without symlinks: %w", err) + } + parentFile := os.NewFile(uintptr(parentFD), parentDirectory) + if parentFile == nil { + _ = unix.Close(parentFD) + return fmt.Errorf("wrap parent directory descriptor") + } + defer parentFile.Close() + + parentInfo, err := parentFile.Stat() + if err != nil { + return fmt.Errorf("inspect parent directory: %w", err) + } + parentStat, ok := parentInfo.Sys().(*syscall.Stat_t) + parentMode := parentInfo.Mode() + if !ok || !parentInfo.IsDir() || parentStat.Uid != trustedOwner { + return fmt.Errorf("parent directory is not owned by trusted UID %d", + trustedOwner) + } + if parentMode.Perm()&0o022 != 0 && parentMode&os.ModeSticky == 0 { + return fmt.Errorf("writable parent directory is not sticky") + } + + if err := createHostPIDLockParent(parentFD, baseName); err != nil { + return fmt.Errorf("create directory: %w", err) + } + fd, err := unix.Openat(parentFD, baseName, + unix.O_RDONLY|unix.O_DIRECTORY|unix.O_CLOEXEC|unix.O_NOFOLLOW, 0) + if err != nil { + return fmt.Errorf("open directory without symlinks: %w", err) + } + file := os.NewFile(uintptr(fd), cleanDirectory) + if file == nil { + _ = unix.Close(fd) + return fmt.Errorf("wrap directory descriptor") + } + defer file.Close() + + openedInfo, err := file.Stat() + if err != nil { + return fmt.Errorf("inspect opened directory: %w", err) + } + openedStat, ok := openedInfo.Sys().(*syscall.Stat_t) + if !ok || !openedInfo.IsDir() || openedStat.Uid != trustedOwner { + return fmt.Errorf("directory is not owned by trusted UID %d", + trustedOwner) + } + if err := file.Chmod(hostPIDLockParentMode); err != nil { + return fmt.Errorf("set sticky directory mode: %w", err) + } + + verifiedInfo, err := file.Stat() + if err != nil { + return fmt.Errorf("verify opened directory: %w", err) + } + verifiedStat, ok := verifiedInfo.Sys().(*syscall.Stat_t) + var currentStat unix.Stat_t + if err := unix.Fstatat(parentFD, baseName, ¤tStat, + unix.AT_SYMLINK_NOFOLLOW); err != nil { + return fmt.Errorf("reinspect directory entry: %w", err) + } + if !ok || currentStat.Mode&unix.S_IFMT != unix.S_IFDIR || + currentStat.Dev != verifiedStat.Dev || + currentStat.Ino != verifiedStat.Ino || + currentStat.Uid != trustedOwner || + verifiedStat.Uid != trustedOwner || + currentStat.Mode&0o7777 != 0o1777 || + verifiedInfo.Mode()&(os.ModePerm|os.ModeSetuid|os.ModeSetgid| + os.ModeSticky) != + hostPIDLockParentMode { + return fmt.Errorf("directory changed while it was prepared") + } + return nil +} + +func configureHostPIDBroker( + response *kubeletdevicepluginv1beta1.ContainerAllocateResponse) { + if response.Envs != nil { + delete(response.Envs, hostpid.EnvironmentVariable) + } + if !hostpid.Enabled(os.Getenv(hostpid.EnvironmentVariable)) { + removeHostPIDMounts(response, hostpid.ContainerDirectory) + return + } + if response.Envs == nil { + response.Envs = make(map[string]string) + } + response.Envs[hostpid.EnvironmentVariable] = "1" + configureCanonicalHostPIDMount(response, hostpid.ContainerDirectory, + hostpid.ServerDirectory, true) + orderHostPIDMountPair(response) +} + +func removeHostPIDMounts( + response *kubeletdevicepluginv1beta1.ContainerAllocateResponse, + containerPath string) { + configuredMounts := make([]*kubeletdevicepluginv1beta1.Mount, 0, + len(response.Mounts)) + for _, mount := range response.Mounts { + if mountTargetsContainerPath(mount, containerPath) { + continue + } + configuredMounts = append(configuredMounts, mount) + } + response.Mounts = configuredMounts +} + +func configureHostPIDLockParentMount( + response *kubeletdevicepluginv1beta1.ContainerAllocateResponse) { + configureCanonicalHostPIDMount(response, hostPIDLockParentDirectory, + hostPIDLockParentDirectory, false) +} + +func orderHostPIDMountPair( + response *kubeletdevicepluginv1beta1.ContainerAllocateResponse) { + parentIndex := -1 + brokerIndex := -1 + for index, mount := range response.Mounts { + if mount == nil { + continue + } + if parentIndex < 0 && + mountTargetsContainerPath(mount, hostPIDLockParentDirectory) { + parentIndex = index + } + if brokerIndex < 0 && + mountTargetsContainerPath(mount, hostpid.ContainerDirectory) { + brokerIndex = index + } + } + if parentIndex < 0 || brokerIndex < 0 || parentIndex < brokerIndex { + return + } + + parentMount := response.Mounts[parentIndex] + brokerMount := response.Mounts[brokerIndex] + orderedMounts := make([]*kubeletdevicepluginv1beta1.Mount, 0, + len(response.Mounts)) + for index, mount := range response.Mounts { + switch index { + case brokerIndex: + orderedMounts = append(orderedMounts, parentMount, brokerMount) + case parentIndex: + continue + default: + orderedMounts = append(orderedMounts, mount) + } + } + response.Mounts = orderedMounts +} + +func configureCanonicalHostPIDMount( + response *kubeletdevicepluginv1beta1.ContainerAllocateResponse, + containerPath string, hostPath string, readOnly bool) { + canonicalMount := &kubeletdevicepluginv1beta1.Mount{ + ContainerPath: containerPath, + HostPath: hostPath, + ReadOnly: readOnly, + } + canonicalMountAdded := false + configuredMounts := make([]*kubeletdevicepluginv1beta1.Mount, 0, + len(response.Mounts)+1) + for _, mount := range response.Mounts { + if mountTargetsContainerPath(mount, containerPath) { + if !canonicalMountAdded { + configuredMounts = append(configuredMounts, canonicalMount) + canonicalMountAdded = true + } + continue + } + configuredMounts = append(configuredMounts, mount) + } + if !canonicalMountAdded { + configuredMounts = append(configuredMounts, canonicalMount) + } + response.Mounts = configuredMounts +} + +func mountTargetsContainerPath( + mount *kubeletdevicepluginv1beta1.Mount, containerPath string) bool { + return mount != nil && filepath.Clean(filepath.Join( + string(filepath.Separator), mount.ContainerPath)) == containerPath +} diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/hostpid_broker_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/hostpid_broker_test.go new file mode 100644 index 0000000000..0d7984aba6 --- /dev/null +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/hostpid_broker_test.go @@ -0,0 +1,419 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * Copyright (c) 2026 The HAMi Authors. + */ + +package plugin + +import ( + "os" + "path/filepath" + "runtime" + "syscall" + "testing" + + "github.com/stretchr/testify/require" + "golang.org/x/sys/unix" + kubeletdevicepluginv1beta1 "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1" + + "github.com/Project-HAMi/HAMi/pkg/device-plugin/nvidiadevice/nvinternal/hostpid" +) + +func TestMain(m *testing.M) { + prepareHostPIDLockParentForAllocation = func() error { return nil } + os.Exit(m.Run()) +} + +func TestPrepareHostPIDLockParent(t *testing.T) { + directory := filepath.Join(t.TempDir(), "vgpulock") + + require.Equal(t, uint32(0), hostPIDLockTrustedOwner) + require.NoError(t, os.MkdirAll(directory, 0o700)) + require.NoError(t, prepareHostPIDLockParent(directory, + uint32(os.Geteuid()))) + + info, err := os.Stat(directory) + require.NoError(t, err) + require.True(t, info.IsDir()) + require.Equal(t, hostPIDLockParentMode, + info.Mode()&(os.ModePerm|os.ModeSticky)) +} + +func TestPrepareHostPIDLockParentRequestsStickyCreateMode(t *testing.T) { + called := false + var requestedFD int + var requestedName string + var requestedMode uint32 + require.NoError(t, createHostPIDLockParentWith( + func(parentFD int, baseName string, mode uint32) error { + called = true + requestedFD = parentFD + requestedName = baseName + requestedMode = mode + return nil + }, 17, "vgpulock")) + require.True(t, called) + require.Equal(t, 17, requestedFD) + require.Equal(t, "vgpulock", requestedName) + require.Equal(t, uint32(0o1777), requestedMode) +} + +func TestPrepareHostPIDLockParentCreatesStickyModeOnLinux(t *testing.T) { + if runtime.GOOS != "linux" { + t.Skip("Linux mkdirat mode behavior is required") + } + + parent := t.TempDir() + parentFD, err := unix.Open(parent, + unix.O_RDONLY|unix.O_DIRECTORY|unix.O_CLOEXEC, 0) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, unix.Close(parentFD)) + }) + + var createErr error + func() { + oldUmask := unix.Umask(0) + defer unix.Umask(oldUmask) + createErr = createHostPIDLockParent(parentFD, "vgpulock") + }() + require.NoError(t, createErr) + + var createdStat unix.Stat_t + require.NoError(t, unix.Fstatat(parentFD, "vgpulock", &createdStat, + unix.AT_SYMLINK_NOFOLLOW)) + require.Equal(t, hostPIDLockParentCreateMode, + uint32(createdStat.Mode)&0o7777) +} + +func TestPrepareDefaultHostPIDLockParentAsRoot(t *testing.T) { + if os.Geteuid() != 0 || + os.Getenv("HAMI_TEST_PRODUCTION_PARENT") != "1" { + t.Skip("an isolated root mount namespace is required") + } + + require.NoError(t, prepareDefaultHostPIDLockParent()) + info, err := os.Stat(hostPIDLockParentDirectory) + require.NoError(t, err) + stat, ok := info.Sys().(*syscall.Stat_t) + require.True(t, ok) + require.Equal(t, hostPIDLockTrustedOwner, stat.Uid) + require.Equal(t, hostPIDLockParentMode, + info.Mode()&(os.ModePerm|os.ModeSticky)) +} + +func TestPrepareHostPIDLockParentRejectsUntrustedObjects(t *testing.T) { + t.Run("parent owner", func(t *testing.T) { + directory := filepath.Join(t.TempDir(), "vgpulock") + require.NoError(t, os.Mkdir(directory, 0o700)) + + err := prepareHostPIDLockParent(directory, uint32(os.Geteuid()+1)) + require.ErrorContains(t, err, + "parent directory is not owned by trusted UID") + }) + + t.Run("regular file", func(t *testing.T) { + path := filepath.Join(t.TempDir(), "vgpulock") + require.NoError(t, os.WriteFile(path, nil, 0o600)) + + require.Error(t, prepareHostPIDLockParent(path, + uint32(os.Geteuid()))) + }) + + t.Run("symlink", func(t *testing.T) { + parent := t.TempDir() + target := filepath.Join(parent, "target") + link := filepath.Join(parent, "vgpulock") + require.NoError(t, os.Mkdir(target, 0o700)) + require.NoError(t, os.Symlink(target, link)) + + require.Error(t, prepareHostPIDLockParent(link, + uint32(os.Geteuid()))) + }) +} + +func TestPrepareHostPIDLockParentRejectsUntrustedParent(t *testing.T) { + t.Run("symlink", func(t *testing.T) { + fixture := t.TempDir() + realParent := filepath.Join(fixture, "real-parent") + linkParent := filepath.Join(fixture, "link-parent") + directory := filepath.Join(linkParent, "vgpulock") + realDirectory := filepath.Join(realParent, "vgpulock") + require.NoError(t, os.Mkdir(realParent, 0o700)) + require.NoError(t, os.Symlink(realParent, linkParent)) + + require.Error(t, prepareHostPIDLockParent(directory, + uint32(os.Geteuid()))) + _, err := os.Lstat(realDirectory) + require.ErrorIs(t, err, os.ErrNotExist) + }) + + t.Run("writable without sticky bit", func(t *testing.T) { + parent := filepath.Join(t.TempDir(), "parent") + directory := filepath.Join(parent, "vgpulock") + require.NoError(t, os.Mkdir(parent, 0o700)) + require.NoError(t, os.Chmod(parent, 0o777)) + + require.Error(t, prepareHostPIDLockParent(directory, + uint32(os.Geteuid()))) + }) +} + +func TestPrepareHostPIDLockParentAllowsStickyWritableParent(t *testing.T) { + parent := filepath.Join(t.TempDir(), "parent") + directory := filepath.Join(parent, "vgpulock") + require.NoError(t, os.Mkdir(parent, 0o700)) + require.NoError(t, os.Chmod(parent, + os.FileMode(0o777)|os.ModeSticky)) + + require.NoError(t, prepareHostPIDLockParent(directory, + uint32(os.Geteuid()))) +} + +func TestConfigureHostPIDBrokerDisabled(t *testing.T) { + for _, value := range []string{"", "0", "true", "false", "01", " 1"} { + t.Run(value, func(t *testing.T) { + t.Setenv(hostpid.EnvironmentVariable, value) + response := &kubeletdevicepluginv1beta1.ContainerAllocateResponse{} + + configureHostPIDBroker(response) + + require.Empty(t, response.Envs) + require.Empty(t, response.Mounts) + }) + } +} + +func TestConfigureHostPIDBrokerDisabledRemovesStaleConfiguration(t *testing.T) { + t.Setenv(hostpid.EnvironmentVariable, "0") + before := &kubeletdevicepluginv1beta1.Mount{ + ContainerPath: "/before", + HostPath: "/before", + } + lockParent := &kubeletdevicepluginv1beta1.Mount{ + ContainerPath: hostPIDLockParentDirectory, + HostPath: hostPIDLockParentDirectory, + } + after := &kubeletdevicepluginv1beta1.Mount{ + ContainerPath: "/after", + HostPath: "/after", + } + response := &kubeletdevicepluginv1beta1.ContainerAllocateResponse{ + Envs: map[string]string{ + "KEEP": "yes", + hostpid.EnvironmentVariable: "1", + }, + Mounts: []*kubeletdevicepluginv1beta1.Mount{ + before, + { + ContainerPath: hostpid.ContainerDirectory + "/", + HostPath: "/first-stale", + }, + lockParent, + { + ContainerPath: "tmp/vgpulock/hostpid", + HostPath: "/second-stale", + }, + after, + }, + } + + configureHostPIDBroker(response) + + require.Equal(t, map[string]string{"KEEP": "yes"}, response.Envs) + require.Equal(t, []*kubeletdevicepluginv1beta1.Mount{ + before, + lockParent, + after, + }, response.Mounts) +} + +func TestConfigureHostPIDBrokerEnabled(t *testing.T) { + t.Setenv(hostpid.EnvironmentVariable, "1") + existingMount := &kubeletdevicepluginv1beta1.Mount{ + ContainerPath: "/tmp/vgpulock", + HostPath: "/tmp/vgpulock", + } + response := &kubeletdevicepluginv1beta1.ContainerAllocateResponse{ + Envs: map[string]string{"KEEP": "yes"}, + Mounts: []*kubeletdevicepluginv1beta1.Mount{existingMount}, + } + + configureHostPIDBroker(response) + + require.Equal(t, "yes", response.Envs["KEEP"]) + require.Equal(t, "1", response.Envs[hostpid.EnvironmentVariable]) + require.Len(t, response.Mounts, 2) + require.Same(t, existingMount, response.Mounts[0]) + require.Equal(t, &kubeletdevicepluginv1beta1.Mount{ + ContainerPath: hostpid.ContainerDirectory, + HostPath: hostpid.ServerDirectory, + ReadOnly: true, + }, response.Mounts[1]) +} + +func TestConfigureHostPIDBrokerIsIdempotent(t *testing.T) { + t.Setenv(hostpid.EnvironmentVariable, "1") + response := &kubeletdevicepluginv1beta1.ContainerAllocateResponse{} + + configureHostPIDBroker(response) + configureHostPIDBroker(response) + + require.Equal(t, "1", response.Envs[hostpid.EnvironmentVariable]) + require.Len(t, response.Mounts, 1) + require.Equal(t, &kubeletdevicepluginv1beta1.Mount{ + ContainerPath: hostpid.ContainerDirectory, + HostPath: hostpid.ServerDirectory, + ReadOnly: true, + }, response.Mounts[0]) +} + +func TestConfigureHostPIDBrokerReplacesConflictingMount(t *testing.T) { + t.Setenv(hostpid.EnvironmentVariable, "1") + response := &kubeletdevicepluginv1beta1.ContainerAllocateResponse{ + Mounts: []*kubeletdevicepluginv1beta1.Mount{{ + ContainerPath: hostpid.ContainerDirectory, + HostPath: "/untrusted", + ReadOnly: false, + }}, + } + + configureHostPIDBroker(response) + + require.Len(t, response.Mounts, 1) + require.Equal(t, hostpid.ServerDirectory, + response.Mounts[0].HostPath) + require.True(t, response.Mounts[0].ReadOnly) +} + +func TestConfigureHostPIDBrokerCanonicalizesDuplicateMounts(t *testing.T) { + t.Setenv(hostpid.EnvironmentVariable, "1") + before := &kubeletdevicepluginv1beta1.Mount{ + ContainerPath: "/before", + HostPath: "/before", + } + after := &kubeletdevicepluginv1beta1.Mount{ + ContainerPath: "/after", + HostPath: "/after", + } + response := &kubeletdevicepluginv1beta1.ContainerAllocateResponse{ + Mounts: []*kubeletdevicepluginv1beta1.Mount{ + before, + { + ContainerPath: hostpid.ContainerDirectory + "/", + HostPath: "/first-untrusted", + ReadOnly: false, + }, + after, + { + ContainerPath: hostPIDLockParentDirectory + + "/hostpid/../hostpid", + HostPath: "/second-untrusted", + ReadOnly: false, + }, + }, + } + + configureHostPIDBroker(response) + + require.Len(t, response.Mounts, 3) + require.Same(t, before, response.Mounts[0]) + require.Equal(t, &kubeletdevicepluginv1beta1.Mount{ + ContainerPath: hostpid.ContainerDirectory, + HostPath: hostpid.ServerDirectory, + ReadOnly: true, + }, response.Mounts[1]) + require.Same(t, after, response.Mounts[2]) +} + +func TestConfigureHostPIDBrokerOrdersParentBeforeNestedMount(t *testing.T) { + t.Setenv(hostpid.EnvironmentVariable, "1") + before := &kubeletdevicepluginv1beta1.Mount{ + ContainerPath: "/before", + HostPath: "/before", + } + middle := &kubeletdevicepluginv1beta1.Mount{ + ContainerPath: "/middle", + HostPath: "/middle", + } + after := &kubeletdevicepluginv1beta1.Mount{ + ContainerPath: "/after", + HostPath: "/after", + } + response := &kubeletdevicepluginv1beta1.ContainerAllocateResponse{ + Mounts: []*kubeletdevicepluginv1beta1.Mount{ + before, + { + ContainerPath: "tmp/vgpulock/hostpid", + HostPath: "/untrusted-broker", + ReadOnly: false, + }, + middle, + { + ContainerPath: "/tmp/./vgpulock", + HostPath: "/untrusted-parent", + ReadOnly: true, + }, + after, + }, + } + + configureHostPIDLockParentMount(response) + configureHostPIDBroker(response) + + require.Equal(t, []*kubeletdevicepluginv1beta1.Mount{ + before, + { + ContainerPath: hostPIDLockParentDirectory, + HostPath: hostPIDLockParentDirectory, + ReadOnly: false, + }, + { + ContainerPath: hostpid.ContainerDirectory, + HostPath: hostpid.ServerDirectory, + ReadOnly: true, + }, + middle, + after, + }, response.Mounts) +} + +func TestConfigureHostPIDLockParentMountCanonicalizesDuplicateMounts( + t *testing.T) { + before := &kubeletdevicepluginv1beta1.Mount{ + ContainerPath: "/before", + HostPath: "/before", + } + after := &kubeletdevicepluginv1beta1.Mount{ + ContainerPath: "/after", + HostPath: "/after", + } + response := &kubeletdevicepluginv1beta1.ContainerAllocateResponse{ + Mounts: []*kubeletdevicepluginv1beta1.Mount{ + before, + { + ContainerPath: hostPIDLockParentDirectory + "/", + HostPath: "/first-untrusted", + ReadOnly: true, + }, + after, + { + ContainerPath: "tmp/./vgpulock", + HostPath: "/second-untrusted", + ReadOnly: false, + }, + }, + } + + configureHostPIDLockParentMount(response) + + require.Len(t, response.Mounts, 3) + require.Same(t, before, response.Mounts[0]) + require.Equal(t, &kubeletdevicepluginv1beta1.Mount{ + ContainerPath: hostPIDLockParentDirectory, + HostPath: hostPIDLockParentDirectory, + ReadOnly: false, + }, response.Mounts[1]) + require.Same(t, after, response.Mounts[2]) +} diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go index bc23113f44..b92eaf25c7 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go @@ -911,8 +911,11 @@ func (plugin *NvidiaDevicePlugin) Allocate(ctx context.Context, reqs *kubeletdev os.MkdirAll(cacheFileHostDirectory, 0777) os.Chmod(cacheFileHostDirectory, 0777) - os.MkdirAll("/tmp/vgpulock", 0777) - os.Chmod("/tmp/vgpulock", 0777) + if err := prepareHostPIDLockParentForAllocation(); err != nil { + PodAllocationFailed(nodename, current, NodeLockNvidia) + return nil, fmt.Errorf( + "failed to prepare host PID lock parent: %w", err) + } response.Mounts = append(response.Mounts, &kubeletdevicepluginv1beta1.Mount{ContainerPath: fmt.Sprintf("%s/vgpu/libvgpu.so", hostHookPath), HostPath: GetLibPath(), @@ -920,10 +923,9 @@ func (plugin *NvidiaDevicePlugin) Allocate(ctx context.Context, reqs *kubeletdev &kubeletdevicepluginv1beta1.Mount{ContainerPath: fmt.Sprintf("%s/vgpu", hostHookPath), HostPath: cacheFileHostDirectory, ReadOnly: false}, - &kubeletdevicepluginv1beta1.Mount{ContainerPath: "/tmp/vgpulock", - HostPath: "/tmp/vgpulock", - ReadOnly: false}, ) + configureHostPIDLockParentMount(response) + configureHostPIDBroker(response) found := false for _, val := range currentCtr.Env { if strings.Compare(val.Name, "CUDA_DISABLE_CONTROL") == 0 { diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go index 21bda35763..0c3077c76f 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/server_test.go @@ -35,6 +35,7 @@ package plugin import ( "context" "encoding/json" + "errors" "fmt" "os" "reflect" @@ -43,6 +44,7 @@ import ( v1 "github.com/NVIDIA/k8s-device-plugin/api/config/v1" "github.com/Project-HAMi/HAMi/pkg/device" "github.com/Project-HAMi/HAMi/pkg/device-plugin/nvidiadevice/nvinternal/cdi" + "github.com/Project-HAMi/HAMi/pkg/device-plugin/nvidiadevice/nvinternal/hostpid" "github.com/Project-HAMi/HAMi/pkg/device-plugin/nvidiadevice/nvinternal/imex" "github.com/Project-HAMi/HAMi/pkg/device-plugin/nvidiadevice/nvinternal/rm" "github.com/Project-HAMi/HAMi/pkg/device/nvidia" @@ -169,7 +171,7 @@ func TestCDIAllocateResponse(t *testing.T) { } for i := range testCases { - tc := testCases[i] + tc := &testCases[i] t.Run(tc.description, func(t *testing.T) { deviceListStrategies, _ := v1.NewDeviceListStrategies(tc.deviceListStrategies) plugin := NvidiaDevicePlugin{ @@ -702,7 +704,23 @@ func TestAlignContainerDevicesWithAllocatedIDsRejectsLengthMismatch(t *testing.T require.Contains(t, err.Error(), "device number not matched") } -func TestAllocateUsesKubeletSelectedUUIDsForVGPUResponse(t *testing.T) { +func TestAllocateUsesSelectedUUIDsAndHostPIDBroker(t *testing.T) { + t.Setenv(hostpid.EnvironmentVariable, "1") + prepareCalls := 0 + previousPrepareHostPIDLockParent := prepareHostPIDLockParentForAllocation + prepareHostPIDLockParentForAllocation = func() error { + prepareCalls++ + return nil + } + defer func() { + prepareHostPIDLockParentForAllocation = + previousPrepareHostPIDLockParent + }() + previousEnableGetPreferredAllocation := enableGetPreferredAllocation + enableGetPreferredAllocation = true + defer func() { + enableGetPreferredAllocation = previousEnableGetPreferredAllocation + }() deviceListStrategies, _ := v1.NewDeviceListStrategies([]string{"envvar"}) deviceIDStrategy := v1.DeviceIDStrategyUUID memScale := 1.0 @@ -770,9 +788,62 @@ func TestAllocateUsesKubeletSelectedUUIDsForVGPUResponse(t *testing.T) { response, err := plugin.Allocate(context.Background(), request) require.NoError(t, err) + require.Equal(t, 1, prepareCalls) require.Equal(t, "GPU-03f69c50-207a-2038-9b45-23cac89cb67a", response.ContainerResponses[0].Envs[deviceListEnvVar]) require.Equal(t, "3000m", response.ContainerResponses[0].Envs["CUDA_DEVICE_MEMORY_LIMIT_0"]) require.Equal(t, "50", response.ContainerResponses[0].Envs["CUDA_DEVICE_SM_LIMIT"]) + require.Equal(t, "1", response.ContainerResponses[0].Envs[hostpid.EnvironmentVariable]) + brokerMountCount := 0 + brokerMountIndex := -1 + fallbackParentMountCount := 0 + fallbackParentMountIndex := -1 + for mountIndex, mount := range response.ContainerResponses[0].Mounts { + if mount.ContainerPath == hostpid.ContainerDirectory { + require.Equal(t, hostpid.ServerDirectory, mount.HostPath) + require.True(t, mount.ReadOnly) + brokerMountIndex = mountIndex + brokerMountCount++ + } + if mount.ContainerPath == hostPIDLockParentDirectory { + require.Equal(t, hostPIDLockParentDirectory, mount.HostPath) + require.False(t, mount.ReadOnly) + fallbackParentMountIndex = mountIndex + fallbackParentMountCount++ + } + } + require.Equal(t, 1, brokerMountCount) + require.Equal(t, 1, fallbackParentMountCount) + require.Less(t, fallbackParentMountIndex, brokerMountIndex) + + t.Setenv(hostpid.EnvironmentVariable, "") + pod.Annotations["hami.io/vgpu-devices-to-allocate"] = + "GPU-annotated-a,NVIDIA,3000,50:;" + client.KubeClient = fake.NewSimpleClientset(pod) + disabledResponse, err := plugin.Allocate(context.Background(), request) + require.NoError(t, err) + require.Equal(t, 2, prepareCalls) + require.NotContains(t, disabledResponse.ContainerResponses[0].Envs, + hostpid.EnvironmentVariable) + fallbackMountCount := 0 + for _, mount := range disabledResponse.ContainerResponses[0].Mounts { + require.NotEqual(t, hostpid.ContainerDirectory, mount.ContainerPath) + if mount.ContainerPath == hostPIDLockParentDirectory { + require.Equal(t, hostPIDLockParentDirectory, mount.HostPath) + require.False(t, mount.ReadOnly) + fallbackMountCount++ + } + } + require.Equal(t, 1, fallbackMountCount) + + prepareHostPIDLockParentForAllocation = func() error { + return errors.New("parent preparation fixture") + } + pod.Annotations["hami.io/vgpu-devices-to-allocate"] = + "GPU-annotated-a,NVIDIA,3000,50:;" + client.KubeClient = fake.NewSimpleClientset(pod) + failedResponse, err := plugin.Allocate(context.Background(), request) + require.Nil(t, failedResponse) + require.ErrorContains(t, err, "failed to prepare host PID lock parent") } func TestAllocatePreservesContainerOrderWhenOneContainerFallsBack(t *testing.T) {