Wait for MPS daemon readiness before advertising shared resources

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.")

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

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)
+}

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...)

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
 		}

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
 }

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
 }
 

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")
+}

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)
 	}