Add option for GPU availability (microsoft/vscode-remote-release#9385)

Author: chrmarti

src/spec-node/devContainers.ts

@@ -8,7 +8,7 @@ import * as crypto from 'crypto';
 import * as os from 'os';
 
 import { mapNodeOSToGOOS, mapNodeArchitectureToGOARCH } from '../spec-configuration/containerCollectionsOCI';
-import { DockerResolverParameters, DevContainerAuthority, UpdateRemoteUserUIDDefault, BindMountConsistency, getCacheFolder } from './utils';
+import { DockerResolverParameters, DevContainerAuthority, UpdateRemoteUserUIDDefault, BindMountConsistency, getCacheFolder, GPUAvailability } from './utils';
 import { createNullLifecycleHook, finishBackgroundTasks, ResolverParameters, UserEnvProbe } from '../spec-common/injectHeadless';
 import { GoARCH, GoOS, getCLIHost, loadNativeModule } from '../spec-common/commonUtils';
 import { resolve } from './configContainer';
@@ -28,6 +28,7 @@ export interface ProvisionOptions {
 	containerSystemDataFolder: string | undefined;
 	workspaceFolder: string | undefined;
 	workspaceMountConsistency?: BindMountConsistency;
+	gpuAvailability?: GPUAvailability;
 	mountWorkspaceGitRoot: boolean;
 	configFile: URI | undefined;
 	overrideConfigFile: URI | undefined;
@@ -101,7 +102,7 @@ export async function launch(options: ProvisionOptions, providedIdLabels: string
 }
 
 export async function createDockerParams(options: ProvisionOptions, disposables: (() => Promise<unknown> | undefined)[]): Promise<DockerResolverParameters> {
-	const { persistedFolder, additionalMounts, updateRemoteUserUIDDefault, containerDataFolder, containerSystemDataFolder, workspaceMountConsistency, mountWorkspaceGitRoot, remoteEnv, experimentalLockfile, experimentalFrozenLockfile, omitLoggerHeader, secretsP } = options;
+	const { persistedFolder, additionalMounts, updateRemoteUserUIDDefault, containerDataFolder, containerSystemDataFolder, workspaceMountConsistency, gpuAvailability, mountWorkspaceGitRoot, remoteEnv, experimentalLockfile, experimentalFrozenLockfile, omitLoggerHeader, secretsP } = options;
 	let parsedAuthority: DevContainerAuthority | undefined;
 	if (options.workspaceFolder) {
 		parsedAuthority = { hostPath: options.workspaceFolder } as DevContainerAuthority;
@@ -212,6 +213,7 @@ export async function createDockerParams(options: ProvisionOptions, disposables:
 		dockerComposeCLI: dockerComposeCLI,
 		dockerEnv: cliHost.env,
 		workspaceMountConsistencyDefault: workspaceMountConsistency,
+		gpuAvailability: gpuAvailability || 'detect',
 		mountWorkspaceGitRoot,
 		updateRemoteUserUIDOnMacOS: false,
 		cacheMount: 'bind',

src/spec-node/devContainersSpecCLI.ts

@@ -105,6 +105,7 @@ function provisionOptions(y: Argv) {
 		'container-system-data-folder': { type: 'string', description: 'Container system data folder where system data inside the container will be stored.' },
 		'workspace-folder': { type: 'string', description: 'Workspace folder path. The devcontainer.json will be looked up relative to this path.' },
 		'workspace-mount-consistency': { choices: ['consistent' as 'consistent', 'cached' as 'cached', 'delegated' as 'delegated'], default: 'cached' as 'cached', description: 'Workspace mount consistency.' },
+		'gpu-availability': { choices: ['all' as 'all', 'detect' as 'detect', 'none' as 'none'], default: 'detect' as 'detect', description: 'Availability of GPUs in case the dev container requires any. `all` expects a GPU to be available.' },
 		'mount-workspace-git-root': { type: 'boolean', default: true, description: 'Mount the workspace using its Git root.' },
 		'id-label': { type: 'string', description: 'Id label(s) of the format name=value. These will be set on the container and used to query for an existing container. If no --id-label is given, one will be inferred from the --workspace-folder path.' },
 		'config': { type: 'string', description: 'devcontainer.json path. The default is to use .devcontainer/devcontainer.json or, if that does not exist, .devcontainer.json in the workspace folder.' },
@@ -179,6 +180,7 @@ async function provision({
 	'container-system-data-folder': containerSystemDataFolder,
 	'workspace-folder': workspaceFolderArg,
 	'workspace-mount-consistency': workspaceMountConsistency,
+	'gpu-availability': gpuAvailability,
 	'mount-workspace-git-root': mountWorkspaceGitRoot,
 	'id-label': idLabel,
 	config,
@@ -233,6 +235,7 @@ async function provision({
 		containerSystemDataFolder,
 		workspaceFolder,
 		workspaceMountConsistency,
+		gpuAvailability,
 		mountWorkspaceGitRoot,
 		configFile: config ? URI.file(path.resolve(process.cwd(), config)) : undefined,
 		overrideConfigFile: overrideConfig ? URI.file(path.resolve(process.cwd(), overrideConfig)) : undefined,

src/spec-node/singleContainer.ts

@@ -329,7 +329,7 @@ export async function findDevContainer(params: DockerCLIParameters | DockerResol
 	return details.filter(container => container.State.Status !== 'removing')[0];
 }
 
-export async function extraRunArgs(common: ResolverParameters, params: DockerCLIParameters | DockerResolverParameters, config: DevContainerFromDockerfileConfig | DevContainerFromImageConfig) {
+export async function extraRunArgs(common: ResolverParameters, params: DockerResolverParameters, config: DevContainerFromDockerfileConfig | DevContainerFromImageConfig) {
 	const extraArguments: string[] = [];
 	if (config.hostRequirements?.gpu) {
 		if (await checkDockerSupportForGPU(params)) {

src/spec-node/utils.ts

@@ -35,6 +35,8 @@ export { uriToFsPath, parentURI } from '../spec-configuration/configurationCommo
 
 export type BindMountConsistency = 'consistent' | 'cached' | 'delegated' | undefined;
 
+export type GPUAvailability = 'all' | 'detect' | 'none';
+
 // Generic retry function
 export async function retry<T>(fn: () => Promise<T>, options: { retryIntervalMilliseconds: number; maxRetries: number; output: Log }): Promise<T> {
 	const { retryIntervalMilliseconds, maxRetries, output } = options;
@@ -100,6 +102,7 @@ export interface DockerResolverParameters {
 	dockerComposeCLI: () => Promise<DockerComposeCLI>;
 	dockerEnv: NodeJS.ProcessEnv;
 	workspaceMountConsistencyDefault: BindMountConsistency;
+	gpuAvailability: GPUAvailability;
 	mountWorkspaceGitRoot: boolean;
 	updateRemoteUserUIDOnMacOS: boolean;
 	cacheMount: 'volume' | 'bind' | 'none';
@@ -202,7 +205,13 @@ async function hasLabels(params: DockerResolverParameters, info: any, expectedLa
 		.every(name => actualLabels[name] === expectedLabels[name]);
 }
 
-export async function checkDockerSupportForGPU(params: DockerCLIParameters | DockerResolverParameters): Promise<Boolean> {
+export async function checkDockerSupportForGPU(params: DockerResolverParameters): Promise<Boolean> {
+	if (params.gpuAvailability === 'all') {
+		return true;
+	}
+	if (params.gpuAvailability === 'none') {
+		return false;
+	}
 	const result = await dockerCLI(params, 'info', '-f', '{{.Runtimes.nvidia}}');
 	const runtimeFound = result.stdout.includes('nvidia-container-runtime');
 	return runtimeFound;