visor/pkg/sentry/devices/nvproxy/save_restore.go

// Copyright 2023 The gVisor Authors.
//
// 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 nvproxy

import (
	goContext "context"
	"fmt"
	"path/filepath"

	"golang.org/x/sys/unix"
	"gvisor.dev/gvisor/pkg/abi/nvgpu"
	"gvisor.dev/gvisor/pkg/context"
	"gvisor.dev/gvisor/pkg/devutil"
	"gvisor.dev/gvisor/pkg/fdnotifier"
	"gvisor.dev/gvisor/pkg/log"
	"gvisor.dev/gvisor/pkg/sentry/vfs"
)

type saveRestoreGoogleContextID int

const (
	// CtxDeviceRemapping is a Context.Value key for a DeviceRemapping.
	CtxDeviceRemapping saveRestoreGoogleContextID = iota
)

// DeviceRemappingFromContext returns the DeviceRemapping associated with ctx,
// or nil if no such DeviceRemapping exists.
func DeviceRemappingFromContext(ctx goContext.Context) *DeviceRemapping {
	if r := ctx.Value(CtxDeviceRemapping); r != nil {
		return r.(*DeviceRemapping)
	}
	return nil
}

// AnyFrontendDevicesOpen returns true if at least one "frontend device"
// (/dev/nvidia# or /dev/nvidiactl) has at least one open file descriptor.
func AnyFrontendDevicesOpen(vfsObj *vfs.VirtualFilesystem) bool {
	nvp := nvproxyFromVFS(vfsObj)
	if nvp == nil {
		return false
	}
	nvp.fdsMu.Lock()
	defer nvp.fdsMu.Unlock()
	return len(nvp.frontendFDs) != 0
}

// DeviceRemapID identifies a device for remapping across save/restore.
type DeviceRemapID struct {
	// Minor is the minor device number. This is assigned by
	// kernel-open/nvidia/nv-pci.c:nv_pci_probe() =>
	// kernel-open/nvidia/nv.c:nv_linux_add_device_locked().
	Minor uint32 `json:"minor"`

	// PCIVendorID and PCIDeviceID are PCI vendor and device IDs respectively.
	// These are properties of the underlying hardware.
	PCIVendorID uint16 `json:"pci_vendor_id"`
	PCIDeviceID uint16 `json:"pci_device_id"`

	// GPUID is nv_state_t::gpu_id. This is generated by
	// src/nvidia/arch/nvalloc/unix/src/osinit.c:RmInitPrivateState() =>
	// nv_generate_id_from_pci_info().
	GPUID uint32 `json:"gpu_id"`

	// DeviceInstance is the "broadcast device instance number", or
	// equivalently the value of NV0080_ALLOC_PARAMETERS::deviceId, or
	// equivalently the index of this device in OBJGPUMGR::pGpuGrpTable. This
	// is assigned by
	// src/nvidia/src/kernel/gpu_mgr/gpu_mgr.c:gpumgrCreateDevice() =>
	// gpumgrAllocDeviceInstance().
	DeviceInstance uint32 `json:"device_instance"`

	// SubDeviceInstance is the "unicast subdevice instance number", or
	// equivalently the value of NV2080_ALLOC_PARAMETERS::subDeviceId, or
	// equivalently the value of OBJGPU::subdeviceInstance. This is set by an
	// invocation of NV2080_CTRL_CMD_GPU_SET_SDM.
	SubDeviceInstance uint32 `json:"subdevice_instance"`

	// UUID is the "UUID", including the "GPU-" prefix, as printed by
	// `nvidia-smi -L`. This is provided by the GSP, and retrieved and
	// stringified by src/nvidia/src/kernel/gpu/gpu.c:gpuGetGidInfo_IMPL().
	UUID string `json:"uuid"`

	// If DevGoferClient is not nil, it provides access to the represented
	// device.
	DevGoferClient *devutil.GoferClient `json:"-"`
}

// String implements fmt.Stringer.String.
func (id *DeviceRemapID) String() string {
	return fmt.Sprintf("{Minor:%d PCIVendorID:0x%04x PCIDeviceID:0x%04x GPUID:%#x DeviceInstance:%d SubDeviceInstance:%d UUID:%s}", id.Minor, id.PCIVendorID, id.PCIDeviceID, id.GPUID, id.DeviceInstance, id.SubDeviceInstance, id.UUID)
}

// CheckDevicesRemappable checks that the set of devices in ids can be saved.
func CheckDevicesRemappable(ids []DeviceRemapID) error {
	if len(ids) == 0 {
		return nil
	}
	// Check that all hardware is identical. MakeDeviceRemapping currently
	// relies on this property during both save and restore to avoid having to
	// remap based on hardware identity. This property is satisfied on all
	// known GCE GPU instance types.
	for _, id := range ids[1:] {
		if ids[0].PCIVendorID != id.PCIVendorID || ids[0].PCIDeviceID != id.PCIDeviceID {
			return fmt.Errorf("can't remap with heterogeneous GPUs: devices %v and %v have different PCI vendor/device IDs", ids[0], id)
		}
	}
	// Check that all subdevice instances are 0, which implies that each device
	// has a distinct device instance. This is currently relied upon to avoid
	// having to preserve device-subdevice relationships in
	// MakeDeviceRemapping(), or change them in nvproxy.afterLoadImpl(). This
	// should be the case when using the OSS kernel driver, since
	// src/nvidia/src/kernel/gpu_mgr/gpu_mgr.c:gpumgrCreateDevice() returns
	// NV_ERR_INVALID_REQUEST if !gpumgrIsSubDeviceCountOne() (based on the
	// code structure, there is likely redacted handling for the
	// !gpumgrIsSubDeviceCountOne() case that is only present in the
	// proprietary driver).
	for _, id := range ids {
		if id.SubDeviceInstance != 0 {
			return fmt.Errorf("device %v has non-zero subdevice instance; all devices: %v", id, ids)
		}
	}
	return nil
}

// DeviceRemapping specifies the remapping from a set of saved devices to a set
// of devices available during restore.
type DeviceRemapping struct {
	// XByY maps Y to X.
	NewDeviceByOld            map[*DeviceRemapID]*DeviceRemapID
	OldDeviceByMinor          map[uint32]*DeviceRemapID
	OldDeviceByDeviceInstance map[uint32]*DeviceRemapID
}

// MakeDeviceRemapping produces a DeviceRemapping. It takes ownership of oldIDs
// and newIDs.
//
// Preconditions: CheckDevicesRemappable(oldIDs) == nil.
func MakeDeviceRemapping(oldIDs, newIDs []DeviceRemapID) (*DeviceRemapping, error) {
	if len(oldIDs) > len(newIDs) {
		return nil, fmt.Errorf("can't remap %d saved devices to %d restored devices", len(oldIDs), len(newIDs))
	}
	if len(oldIDs) == 0 {
		return nil, nil
	}
	if err := CheckDevicesRemappable(newIDs); err != nil {
		return nil, err
	}
	dr := &DeviceRemapping{
		NewDeviceByOld:            make(map[*DeviceRemapID]*DeviceRemapID),
		OldDeviceByMinor:          make(map[uint32]*DeviceRemapID),
		OldDeviceByDeviceInstance: make(map[uint32]*DeviceRemapID),
	}
	for i := range oldIDs {
		oldID := &oldIDs[i]
		newID := &newIDs[i]
		dr.NewDeviceByOld[oldID] = newID
		dr.OldDeviceByMinor[oldID.Minor] = oldID
		dr.OldDeviceByDeviceInstance[oldID.DeviceInstance] = oldID
	}
	return dr, nil
}

// beforeSave is invoked by stateify.
func (nvp *nvproxy) beforeSave() {
	// No need to hold any locks since the sandbox is paused.
	// This avoids lock inversion between nvp.clientsMu and rootClient.objsMu.
	for _, client := range nvp.clients {
		for _, obj := range client.resources {
			if _, ok := obj.impl.(restorableObjectImpl); !ok {
				panic(fmt.Sprintf("nvproxy: object handle %v:%v (class %v, type %T) is not restorable", client.handle, obj.handle, obj.class, obj.impl))
			}
		}
	}
}

// afterLoad is invoked by stateify.
func (nvp *nvproxy) afterLoad(ctx goContext.Context) {
	Init()
	abiEntry, ok := abis[nvp.version]
	if !ok {
		panic(fmt.Sprintf("driver version %q not found in abis map", nvp.version))
	}
	nvp.abi = abiEntry.cons()
	if dr := DeviceRemappingFromContext(ctx); dr != nil {
		// Remap frontendFDs to their new frontendDevices.
		// frontendFD.vfsfd.VirtualDentry() will still point to the original
		// dentry, so e.g. /proc/[pid]/fd/ will be incorrect, but we don't have
		// a way to get the correct dentry, which may not even exist
		// (frontendFDs might have been opened from application-created device
		// special files in any filesystem not mounted nodev).
		func() {
			nvp.fdsMu.Lock()
			defer nvp.fdsMu.Unlock()
			for fd := range nvp.frontendFDs {
				oldMinor := fd.dev.minor
				if oldMinor > nvgpu.NV_MINOR_DEVICE_NUMBER_REGULAR_MAX {
					continue
				}
				oldDev, ok := dr.OldDeviceByMinor[oldMinor]
				if !ok {
					panic(fmt.Sprintf("nvproxy: remapping does not contain a saved device with minor number %d: %+v", oldMinor, dr))
				}
				fd.dev = nvp.regularDevs[dr.NewDeviceByOld[oldDev].Minor]
			}
		}()
		// Remap objects to their new devices.
		// Calls to CheckDevicesRemappable() already checked that all subdevice
		// instances are 0, so we don't need to update any NV20_SUBDEVICE_0
		// objects.
		for _, client := range nvp.clients {
			for _, obj := range client.resources {
				if obj.class == nvgpu.NV01_DEVICE_0 {
					obj := obj.impl.(*rmAllocObject)
					var allocParams nvgpu.NV0080_ALLOC_PARAMETERS
					// src/nvidia/src/kernel/gpu/device.c:deviceConstruct_IMPL()
					// treats NULL allocParams with class NV01_DEVICE_0 as
					// equivalent to (what Go calls) the zero value.
					if len(obj.params.allocParams) != 0 {
						allocParams.UnmarshalBytes(obj.params.allocParams)
					}
					oldDev, ok := dr.OldDeviceByDeviceInstance[allocParams.DeviceID]
					if !ok {
						panic(fmt.Sprintf("nvproxy: remapping does not contain a saved device with device instance %d: %+v", allocParams.DeviceID, dr))
					}
					allocParams.DeviceID = dr.NewDeviceByOld[oldDev].DeviceInstance
					if size := allocParams.SizeBytes(); len(obj.params.allocParams) < size {
						obj.params.allocParams = make([]byte, size)
					}
					allocParams.MarshalBytes(obj.params.allocParams)
				}
			}
		}
	}

	// Ensure that frontendFDs have host FDs before restoring objects that may
	// use them.
	func() {
		nvp.fdsMu.Lock()
		defer nvp.fdsMu.Unlock()
		for fd := range nvp.frontendFDs {
			fd.load(ctx)
		}
	}()

	// Produce a topological sort of all live objects so that we can restore
	// each object before its dependents. Algorithm adapted from Cormen et al.,
	// "Introduction to Algorithms", 3rd ed., 2009, Sec. 22.4 "Topological
	// sort".
	objsVisited := make(map[*object]bool) // false => discovered, true => finished
	var objsSorted []*object
	var dfsConsider func(obj *object)
	dfsConsider = func(obj *object) {
		if finished, ok := objsVisited[obj]; ok {
			if !finished {
				panic(fmt.Sprintf("nvproxy: circular dependency through object handle %v:%v (class %v)", obj.client.handle, obj.handle, obj.class))
			}
			return
		}
		objsVisited[obj] = false
		for dep := range obj.deps {
			dfsConsider(dep)
		}
		objsVisited[obj] = true
		objsSorted = append(objsSorted, obj)
	}
	// No need to hold any locks since the sandbox is paused.
	// This avoids lock inversion between nvp.clientsMu and rootClient.objsMu.
	for _, client := range nvp.clients {
		for _, obj := range client.resources {
			dfsConsider(obj)
		}
	}

	// Restore objects in sorted order.
	var depHs []nvgpu.Handle
	for _, obj := range objsSorted {
		for dep := range obj.deps {
			depHs = append(depHs, dep.handle)
		}
		// This type assertion was checked in nvp.beforeSave().
		if err := obj.impl.(restorableObjectImpl).Restore(ctx); err != nil {
			panic(fmt.Sprintf("nvproxy: failed to restore object handle %v:%v (class %v, type %T, dependencies %v): %v", obj.client.handle, obj.handle, obj.class, obj.impl, depHs, err))
		}
		if log.IsLogging(log.Debug) {
			log.Debugf("nvproxy: restored object handle %v:%v (class %v, type %T, dependencies %v)", obj.client.handle, obj.handle, obj.class, obj.impl, depHs)
		}
		// Reuse slice across iterations.
		depHs = depHs[:0]
	}
}

func (fd *frontendFD) load(ctx goContext.Context) {
	basename := fd.dev.basename()
	fd.hostFD = openHostDevFileForRestore(ctx, basename, fd.dev.nvp.useDevGofer, fd.containerName, fd.vfsfd.StatusFlags())
	if err := fdnotifier.AddFD(fd.hostFD, &fd.internalQueue); err != nil {
		panic(fmt.Sprintf("fdnotifier.AddFD(%d) failed for %s device: %v", fd.hostFD, basename, err))
	}
	fd.memmapFile.SetFD(int(fd.hostFD))
}

// beforeSave is invoked by stateify.
func (fd *frontendFD) beforeSave() {
	// no-op
}

// afterLoad is invoked by stateify.
func (fd *frontendFD) afterLoad(ctx goContext.Context) {
	// no-op: nvproxy.afterLoad() => frontendFD.load() restores fd.hostFD
}

// beforeSave is invoked by stateify.
func (fd *uvmFD) beforeSave() {
	// no-op
}

// afterLoad is invoked by stateify.
func (fd *uvmFD) afterLoad(ctx goContext.Context) {
	fd.hostFD = openHostDevFileForRestore(ctx, "nvidia-uvm", fd.dev.nvp.useDevGofer, fd.containerName, fd.vfsfd.StatusFlags())
	if err := fdnotifier.AddFD(fd.hostFD, &fd.queue); err != nil {
		panic(fmt.Sprintf("fdnotifier.AddFD(%d) failed for UVM device: %v", fd.hostFD, err))
	}
	fd.memmapFile.SetFD(int(fd.hostFD))
}

// openHostDevFileForRestore differs from openHostDevFile in the following ways:
//   - The context passed to restore is not bound to a specific container, so
//     devutil.CtxDevGoferClient is not set. Instead, a container name is passed by the
//     caller.
//   - For the same reason, openHostDevFileForRestore does not return a container name.
//   - Restore is unable to propagate errors except via panics, so
//     openHostDevFileForRestore panics on failure rather than restoring an error.
func openHostDevFileForRestore(ctx goContext.Context, relpath string, useDevGofer bool, containerName string, openFlags uint32) int32 {
	if useDevGofer {
		clientProvider := devutil.GoferClientProviderFromContext(ctx)
		if clientProvider == nil {
			panic("devutil.CtxDevGoferClientProvider is not set")
		}
		devClient := clientProvider.GetDevGoferClient(containerName)
		hostFD, err := devClient.OpenAt(ctx.(context.Context), relpath, openFlags)
		if err != nil {
			panic(fmt.Sprintf("failed to open device gofer %s: %v", relpath, err))
		}
		return int32(hostFD)
	}
	abspath := filepath.Join("/dev", relpath)
	hostFD, err := unix.Openat(-1, abspath, int(openFlags&unix.O_ACCMODE|unix.O_NOFOLLOW), 0)
	if err != nil {
		panic(fmt.Sprintf("failed to open host %s: %v", abspath, err))
	}
	return int32(hostFD)
}

type restorableObjectImpl interface {
	objectImpl

	// Restore recreates the represented object in the given driver FD.
	Restore(ctx goContext.Context) error
}

// Restore implements restorableObjectImpl.Restore.
func (o *rmAllocObject) Restore(ctx goContext.Context) error {
	return o.params.restore()
}

// Restore implements restorableObjectImpl.Restore.
func (c *rootClient) Restore(ctx goContext.Context) error {
	// For root clients, the requested handle need to be passed as hRoot rather
	// than hObjectNew; see
	// src/nvidia/src/kernel/rmapi/entry_points.c:_nv04AllocWithAccessSecInfo()
	// => src/nvidia/src/kernel/rmapi/alloc_free.c:rmapiAllocWithSecInfo()
	// (renamed to hClient and *phObject respectively) => _rmAlloc() =>
	// src/nvidia/src/libraries/resserv/src/rs_server.c:serverAllocResource()
	// => serverAllocClient() => _serverCreateEntryAndLockForNewClient().
	// RS_COMPATABILITY_MODE is defined to 1 in src/nvidia/Makefile.
	temp := c.params
	temp.ioctlParams.HRoot = c.handle
	return temp.restore()
}