From bc2da4a0c3c565cf077f0f16518fcfa5b108e2a5 Mon Sep 17 00:00:00 2001
From: Zach Koopmans <zkoopmans@google.com>
Date: Thu, 12 Oct 2023 11:15:11 -0700
Subject: [PATCH] Add internal annotation for nvproxy.

Add an internal annotation so we can set NVProxy in k8s on a per
sandbox (pod) basis.

PiperOrigin-RevId: 572960928
---
 runsc/boot/loader.go      |  6 +++---
 runsc/specutils/nvidia.go | 18 +++++++++++++++++-
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/runsc/boot/loader.go b/runsc/boot/loader.go
index 51d3e4374..851c6dc16 100644
--- a/runsc/boot/loader.go
+++ b/runsc/boot/loader.go
@@ -305,7 +305,7 @@ func New(args Args) (*Loader, error) {
 		return nil, fmt.Errorf("setting up memory usage: %w", err)
 	}
 
-	if args.Conf.NVProxy {
+	if specutils.NVProxyEnabled(args.Spec, args.Conf) {
 		nvproxy.Init()
 	}
 
@@ -362,7 +362,7 @@ func New(args Args) (*Loader, error) {
 	if err != nil {
 		return nil, fmt.Errorf("creating platform: %w", err)
 	}
-	if args.Conf.NVProxy && p.OwnsPageTables() {
+	if specutils.NVProxyEnabled(args.Spec, args.Conf) && p.OwnsPageTables() {
 		return nil, fmt.Errorf("--nvproxy is incompatible with platform %s: owns page tables", args.Conf.Platform)
 	}
 	k := &kernel.Kernel{
@@ -657,7 +657,7 @@ func (l *Loader) installSeccompFilters() error {
 			HostNetworkRawSockets: hostnet && l.root.conf.EnableRaw,
 			HostFilesystem:        l.root.conf.DirectFS,
 			ProfileEnable:         l.root.conf.ProfileEnable,
-			NVProxy:               l.root.conf.NVProxy,
+			NVProxy:               specutils.NVProxyEnabled(l.root.spec, l.root.conf),
 			TPUProxy:              l.root.conf.TPUProxy,
 			ControllerFD:          l.ctrl.srv.FD(),
 		}
diff --git a/runsc/specutils/nvidia.go b/runsc/specutils/nvidia.go
index 4f3de002c..d33b6b5c9 100644
--- a/runsc/specutils/nvidia.go
+++ b/runsc/specutils/nvidia.go
@@ -23,16 +23,32 @@ import (
 	"strings"
 
 	specs "github.com/opencontainers/runtime-spec/specs-go"
+	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/runsc/config"
 )
 
 const nvdEnvVar = "NVIDIA_VISIBLE_DEVICES"
 
+// annotationNVProxy enables nvproxy.
+const annotationNVProxy = "dev.gvisor.spec.nvproxy"
+
+// NVProxyEnabled checks both the nvproxy annotation and conf.NVProxy to see if nvproxy is enabled.
+func NVProxyEnabled(spec *specs.Spec, conf *config.Config) bool {
+	if conf.NVProxy {
+		return true
+	}
+	val, ok := spec.Annotations[annotationNVProxy]
+	if ok && val != "true" {
+		log.Warningf("nvproxy annotation is set to invalid value %q. Ignoring.", val)
+	}
+	return ok && val == "true"
+}
+
 // GPUFunctionalityRequested returns true if the user intends for the sandbox
 // to have access to GPU functionality (e.g. access to /dev/nvidiactl),
 // irrespective of whether or not they want access to any specific GPU.
 func GPUFunctionalityRequested(spec *specs.Spec, conf *config.Config) bool {
-	if !conf.NVProxy {
+	if !NVProxyEnabled(spec, conf) {
 		// nvproxy disabled.
 		return false
 	}