From 07131cdc5de69f8380bf85386c29a987970bdb55 Mon Sep 17 00:00:00 2001 From: Jonathan Meiri <33288957+Meiri28@users.noreply.github.com> Date: Sun, 2 Aug 2026 16:21:04 +0300 Subject: [PATCH] Wait for MPS daemon .ready before advertising shared resources MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The device plugin's waitForDaemon only ran a single AssertHealthy check before serving and registering the resource with kubelet. AssertHealthy issues get_default_active_thread_percentage, which only proves the MPS control pipe is responsive — and the pipe becomes responsive at Daemon.Start (mpsControlBin -d) before the per-device pinned memory limits and active thread percentage are applied. A pod scheduled in that window starts against MPS with the daemon defaults (no pinned memory limit, 100% threads) rather than the configured limits, silently bypassing the intended isolation. The MPS control daemon already creates a node-global .ready file, but only after every daemon's full initialization completes. Nothing consumed it (the two TODOs in waitForDaemon noted exactly this), so the readiness signal was unused. Gate readiness on that file: - Add Root.ReadyFilePath so the marker path has a single definition, and use it in the MPS control daemon for both create and remove instead of the hardcoded "/mps/.ready". - Add Daemon.Ready, which reports whether the .ready file exists. - Rewrite waitForDaemon to poll checkDaemonReady (Ready AND AssertHealthy) every 5s up to a 5m bound, replacing the single unconditional AssertHealthy. On timeout the caller fails and is retried by the plugin manager, so the bound is per-attempt. This closes both TODOs and ensures shared MPS resources are not advertised until the daemon is fully configured. Co-Authored-By: Claude Opus 4.7 (1M context) Co-Authored-By: runatom-ai <258621014+runatom-ai@users.noreply.github.com> Signed-off-by: Jonathan Meiri <33288957+Meiri28@users.noreply.github.com> --- cmd/mps-control-daemon/main.go | 10 +++- cmd/mps-control-daemon/mps/daemon.go | 13 +++++ cmd/mps-control-daemon/mps/daemon_test.go | 62 +++++++++++++++++++++++ cmd/mps-control-daemon/mps/root.go | 6 +++ cmd/nvidia-device-plugin/main.go | 2 +- internal/plugin/api.go | 8 ++- internal/plugin/mps.go | 38 ++++++++++++-- internal/plugin/mps_test.go | 41 +++++++++++++++ internal/plugin/server.go | 4 +- 9 files changed, 172 insertions(+), 12 deletions(-) create mode 100644 cmd/mps-control-daemon/mps/daemon_test.go create mode 100644 internal/plugin/mps_test.go diff --git a/cmd/mps-control-daemon/main.go b/cmd/mps-control-daemon/main.go index 28b491ac8..eac149b32 100644 --- a/cmd/mps-control-daemon/main.go +++ b/cmd/mps-control-daemon/main.go @@ -205,6 +205,12 @@ func startDaemons(c *cli.Context, cfg *Config) ([]*mps.Daemon, bool, error) { klog.Info("No devices are configured for MPS sharing; Waiting indefinitely.") } + // Clear any stale .ready before starting: /mps is a hostPath, so a marker + // from a previous run can survive a restart and signal readiness too early. + if err := os.Remove(mps.ContainerRoot.ReadyFilePath()); err != nil && !os.IsNotExist(err) { + return mpsDaemons, true, fmt.Errorf("failed to remove stale .ready file: %w", err) + } + // Loop through all MPS daemons and start them. // If any daemon fails to start, all daemons are started again. for _, mpsDaemon := range mpsDaemons { @@ -213,7 +219,7 @@ func startDaemons(c *cli.Context, cfg *Config) ([]*mps.Daemon, bool, error) { return mpsDaemons, true, nil } } - readyFile, err := os.Create("/mps/.ready") + readyFile, err := os.Create(mps.ContainerRoot.ReadyFilePath()) if err != nil { return mpsDaemons, true, fmt.Errorf("failed to create .ready file") } @@ -223,7 +229,7 @@ func startDaemons(c *cli.Context, cfg *Config) ([]*mps.Daemon, bool, error) { } func stopDaemons(mpsDaemons ...*mps.Daemon) error { - if err := os.Remove("/mps/.ready"); err != nil { + if err := os.Remove(mps.ContainerRoot.ReadyFilePath()); err != nil { klog.Warningf("Failed to remove .ready file: %v", err) } klog.Info("Stopping MPS daemons.") diff --git a/cmd/mps-control-daemon/mps/daemon.go b/cmd/mps-control-daemon/mps/daemon.go index 0351289cc..b9ed0d779 100644 --- a/cmd/mps-control-daemon/mps/daemon.go +++ b/cmd/mps-control-daemon/mps/daemon.go @@ -206,6 +206,19 @@ func (d *Daemon) AssertHealthy() error { return err } +// Ready reports whether the .ready file exists, i.e. the MPS daemons have +// finished initialization. A stat error other than not-exist is returned. +func (d *Daemon) Ready() (bool, error) { + _, err := os.Stat(d.root.ReadyFilePath()) + if os.IsNotExist(err) { + return false, nil + } + if err != nil { + return false, err + } + return true, nil +} + // EchoPipeToControl sends the specified command to the MPS control daemon. func (d *Daemon) EchoPipeToControl(command string) (string, error) { var out bytes.Buffer diff --git a/cmd/mps-control-daemon/mps/daemon_test.go b/cmd/mps-control-daemon/mps/daemon_test.go new file mode 100644 index 000000000..46e0890ae --- /dev/null +++ b/cmd/mps-control-daemon/mps/daemon_test.go @@ -0,0 +1,62 @@ +/** +# Copyright (c) NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +**/ + +package mps + +import ( + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestReadyFilePath(t *testing.T) { + require.Equal(t, "/mps/.ready", ContainerRoot.ReadyFilePath()) + require.Equal(t, "/custom/root/.ready", Root("/custom/root").ReadyFilePath()) +} + +func TestDaemonReady(t *testing.T) { + root := t.TempDir() + d := &Daemon{root: Root(root)} + + ready, err := d.Ready() + require.NoError(t, err) + require.False(t, ready, "not ready before the .ready file exists") + + require.NoError(t, os.WriteFile(filepath.Join(root, ".ready"), nil, 0o644)) + ready, err = d.Ready() + require.NoError(t, err) + require.True(t, ready, "ready once the .ready file exists") + + require.NoError(t, os.Remove(filepath.Join(root, ".ready"))) + ready, err = d.Ready() + require.NoError(t, err) + require.False(t, ready, "not ready after the .ready file is removed") +} + +// TestDaemonReadyStatError verifies a stat error other than not-exist is +// surfaced rather than reported as "not ready". The root is a regular file, so +// stat-ing a path beneath it fails with ENOTDIR. +func TestDaemonReadyStatError(t *testing.T) { + f := filepath.Join(t.TempDir(), "not-a-dir") + require.NoError(t, os.WriteFile(f, nil, 0o644)) + d := &Daemon{root: Root(f)} + + ready, err := d.Ready() + require.Error(t, err) + require.False(t, ready) +} diff --git a/cmd/mps-control-daemon/mps/root.go b/cmd/mps-control-daemon/mps/root.go index 90655d12e..a1c58cffd 100644 --- a/cmd/mps-control-daemon/mps/root.go +++ b/cmd/mps-control-daemon/mps/root.go @@ -52,6 +52,12 @@ func (r Root) startedFile(resourceName spec.ResourceName) string { return r.Path(string(resourceName), ".started") } +// ReadyFilePath returns the node-global .ready file, created only after all MPS +// daemons finish initialization. +func (r Root) ReadyFilePath() string { + return r.Path(".ready") +} + // Path returns a path relative to the MPS root. func (r Root) Path(parts ...string) string { pathparts := append([]string{string(r)}, parts...) diff --git a/cmd/nvidia-device-plugin/main.go b/cmd/nvidia-device-plugin/main.go index fbe9ddffa..d1975ff7a 100644 --- a/cmd/nvidia-device-plugin/main.go +++ b/cmd/nvidia-device-plugin/main.go @@ -406,7 +406,7 @@ func startPlugins(c *cli.Context, o *options) ([]plugin.Interface, bool, error) } // 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(c.Context, o.kubeletSocket); err != nil { klog.Errorf("Failed to start plugin: %v", err) return plugins, true, nil } diff --git a/internal/plugin/api.go b/internal/plugin/api.go index 92cfa2ecb..b8fbed21d 100644 --- a/internal/plugin/api.go +++ b/internal/plugin/api.go @@ -16,11 +16,15 @@ package plugin -import "github.com/NVIDIA/k8s-device-plugin/internal/rm" +import ( + "context" + + "github.com/NVIDIA/k8s-device-plugin/internal/rm" +) // Interface defines the API for the plugin package type Interface interface { Devices() rm.Devices - Start(string) error + Start(context.Context, string) error Stop() error } diff --git a/internal/plugin/mps.go b/internal/plugin/mps.go index 763e94367..23ee4b91a 100644 --- a/internal/plugin/mps.go +++ b/internal/plugin/mps.go @@ -17,9 +17,12 @@ package plugin import ( + "context" "errors" "fmt" + "time" + "k8s.io/apimachinery/pkg/util/wait" "k8s.io/klog/v2" pluginapi "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1" @@ -28,6 +31,13 @@ import ( "github.com/NVIDIA/k8s-device-plugin/internal/rm" ) +const ( + mpsReadyCheckInterval = 5 * time.Second + // mpsReadyCheckTimeout bounds a single wait attempt; on timeout the plugin + // manager retries. + mpsReadyCheckTimeout = 5 * time.Minute +) + type mpsOptions struct { enabled bool resourceName spec.ResourceName @@ -58,16 +68,34 @@ func (o *options) getMPSOptions(resourceManager rm.ResourceManager) (mpsOptions, return m, nil } -func (m *mpsOptions) waitForDaemon() error { +func (m *mpsOptions) waitForDaemon(ctx context.Context) error { if m == nil || !m.enabled { return nil } - // TODO: Check the .ready file here. - // TODO: Have some retry strategy here. + + return wait.PollUntilContextTimeout(ctx, mpsReadyCheckInterval, mpsReadyCheckTimeout, true, func(context.Context) (bool, error) { + if err := m.checkDaemonReady(); err != nil { + klog.InfoS("Waiting for MPS daemon to be ready", "resource", m.resourceName, "reason", err) + return false, nil + } + klog.InfoS("MPS daemon is ready", "resource", m.resourceName) + return true, nil + }) +} + +// checkDaemonReady requires the .ready file (written after full configuration) +// and a responsive pipe; AssertHealthy alone responds before config is applied. +func (m *mpsOptions) checkDaemonReady() error { + ready, err := m.daemon.Ready() + if err != nil { + return fmt.Errorf("checking MPS daemon readiness: %w", err) + } + if !ready { + return fmt.Errorf("MPS daemon has not signalled readiness") + } if err := m.daemon.AssertHealthy(); err != nil { - return fmt.Errorf("error checking MPS daemon health: %w", err) + return fmt.Errorf("MPS daemon is not healthy: %w", err) } - klog.InfoS("MPS daemon is healthy", "resource", m.resourceName) return nil } diff --git a/internal/plugin/mps_test.go b/internal/plugin/mps_test.go new file mode 100644 index 000000000..0c5fbe897 --- /dev/null +++ b/internal/plugin/mps_test.go @@ -0,0 +1,41 @@ +/** +# Copyright (c) NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +**/ + +package plugin + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "github.com/NVIDIA/k8s-device-plugin/cmd/mps-control-daemon/mps" +) + +// TestCheckDaemonReadyRequiresReadyFile is the regression for this PR's race: +// even when the control pipe would be healthy, readiness must be withheld until +// the .ready file exists. With no .ready file, checkDaemonReady returns +// not-ready before ever consulting the pipe (AssertHealthy). +func TestCheckDaemonReadyRequiresReadyFile(t *testing.T) { + root := t.TempDir() // no .ready file + m := &mpsOptions{ + enabled: true, + daemon: mps.NewDaemon(nil, mps.Root(root)), + } + + err := m.checkDaemonReady() + require.Error(t, err) + require.Contains(t, err.Error(), "has not signalled readiness") +} diff --git a/internal/plugin/server.go b/internal/plugin/server.go index 9281b48ef..afb6bd448 100644 --- a/internal/plugin/server.go +++ b/internal/plugin/server.go @@ -126,10 +126,10 @@ func (plugin *nvidiaDevicePlugin) Devices() rm.Devices { // Start starts the gRPC server, registers the device plugin with the Kubelet, // and starts the device healthchecks. -func (plugin *nvidiaDevicePlugin) Start(kubeletSocket string) error { +func (plugin *nvidiaDevicePlugin) Start(ctx context.Context, kubeletSocket string) error { plugin.initialize() - if err := plugin.mps.waitForDaemon(); err != nil { + if err := plugin.mps.waitForDaemon(ctx); err != nil { return fmt.Errorf("error waiting for MPS daemon: %w", err) }