Wait for MPS daemon .ready before advertising shared resources

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) <noreply@anthropic.com>
Co-Authored-By: runatom-ai <258621014+runatom-ai@users.noreply.github.com>
Signed-off-by: Jonathan Meiri <33288957+Meiri28@users.noreply.github.com>

Author: 3 people

cmd/mps-control-daemon/main.go

@@ -213,7 +213,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 +223,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.")

cmd/mps-control-daemon/mps/daemon.go

@@ -206,6 +206,16 @@ func (d *Daemon) AssertHealthy() error {
 	return err
 }
 
+// Ready returns true once the MPS daemons have signalled that initialization
+// has completed by creating the node-global .ready file under the MPS root.
+// AssertHealthy only proves the control pipe is responsive, which happens
+// before per-device memory limits and thread percentages are applied; Ready
+// gates on the full configuration being in place.
+func (d *Daemon) Ready() bool {
+	_, err := os.Stat(d.root.ReadyFilePath())
+	return err == nil
+}
+
 // EchoPipeToControl sends the specified command to the MPS control daemon.
 func (d *Daemon) EchoPipeToControl(command string) (string, error) {
 	var out bytes.Buffer

cmd/mps-control-daemon/mps/daemon_test.go

@@ -0,0 +1,49 @@
+/**
+# Copyright 2026 NVIDIA CORPORATION
+#
+# 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)}
+
+	// No .ready file yet: the daemon must not report ready. This is the window
+	// in which AssertHealthy can already succeed (control pipe responsive)
+	// while per-device memory/thread configuration is not yet applied.
+	require.False(t, d.Ready(), "daemon must not be ready before the .ready file exists")
+
+	// Once the MPS control daemon has finished initialization it creates the
+	// .ready file; the daemon must then report ready.
+	require.NoError(t, os.WriteFile(filepath.Join(root, ".ready"), nil, 0o644))
+	require.True(t, d.Ready(), "daemon must be ready once the .ready file exists")
+
+	// Removing the file (e.g. on daemon stop) flips readiness back to false.
+	require.NoError(t, os.Remove(filepath.Join(root, ".ready")))
+	require.False(t, d.Ready(), "daemon must not be ready after the .ready file is removed")
+}

cmd/mps-control-daemon/mps/root.go

@@ -52,6 +52,16 @@ func (r Root) startedFile(resourceName spec.ResourceName) string {
 	return r.Path(string(resourceName), ".started")
 }
 
+// ReadyFilePath returns the path to the node-global .ready file. Unlike the
+// per-resource .started file, this single marker is created only after all MPS
+// daemons have completed initialization (compute mode, pinned memory limits,
+// and active thread percentage). Consumers such as the device plugin use it to
+// avoid advertising MPS-shared resources before the daemons are fully
+// configured.
+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...)

internal/plugin/mps.go

@@ -19,6 +19,7 @@ package plugin
 import (
 	"errors"
 	"fmt"
+	"time"
 
 	"k8s.io/klog/v2"
 	pluginapi "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1"
@@ -28,6 +29,15 @@ import (
 	"github.com/NVIDIA/k8s-device-plugin/internal/rm"
 )
 
+const (
+	// mpsReadyCheckInterval is how often we poll for MPS daemon readiness.
+	mpsReadyCheckInterval = 5 * time.Second
+	// mpsReadyCheckTimeout bounds how long we wait for the MPS daemon to become
+	// ready before giving up. On timeout the caller (plugin startup) fails and
+	// is retried by the plugin manager, so this is an upper bound per attempt.
+	mpsReadyCheckTimeout = 5 * time.Minute
+)
+
 type mpsOptions struct {
 	enabled      bool
 	resourceName spec.ResourceName
@@ -62,12 +72,36 @@ func (m *mpsOptions) waitForDaemon() error {
 	if m == nil || !m.enabled {
 		return nil
 	}
-	// TODO: Check the .ready file here.
-	// TODO: Have some retry strategy here.
+
+	deadline := time.Now().Add(mpsReadyCheckTimeout)
+	for {
+		err := m.checkDaemonReady()
+		if err == nil {
+			klog.InfoS("MPS daemon is ready", "resource", m.resourceName)
+			return nil
+		}
+		if time.Now().After(deadline) {
+			return fmt.Errorf("timed out waiting for MPS daemon for %v: %w", m.resourceName, err)
+		}
+		klog.InfoS("Waiting for MPS daemon to be ready", "resource", m.resourceName, "reason", err)
+		time.Sleep(mpsReadyCheckInterval)
+	}
+}
+
+// checkDaemonReady reports whether the MPS daemon has finished initialization.
+// It requires both that the daemon has signalled readiness via its .ready file
+// (created only after compute mode, pinned memory limits, and thread
+// percentages are applied) and that the control pipe is responsive. Checking
+// only AssertHealthy is insufficient: the pipe becomes responsive before the
+// per-device configuration is in place, so a pod scheduled in that window
+// could run without the configured MPS memory/thread limits.
+func (m *mpsOptions) checkDaemonReady() error {
+	if !m.daemon.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
 }