GPU requirement and auto-detect NVIDIA extensions (#173)

src/spec-configuration/configuration.ts

@@ -21,10 +21,16 @@ export type UserEnvProbe = 'none' | 'loginInteractiveShell' | 'interactiveShell'
 
 export type DevContainerConfigCommand = 'initializeCommand' | 'onCreateCommand' | 'updateContentCommand' | 'postCreateCommand' | 'postStartCommand' | 'postAttachCommand';
 
+export interface HostGPURequirements {
+	cores?: number;
+	memory?: string;
+}
+
 export interface HostRequirements {
 	cpus?: number;
 	memory?: string;
 	storage?: string;
+	gpu?: boolean | 'optional' | HostGPURequirements;
 }
 
 export interface DevContainerFeature {

src/spec-node/dockerCompose.ts

@@ -6,7 +6,7 @@
 import * as yaml from 'js-yaml';
 import * as shellQuote from 'shell-quote';
 
-import { createContainerProperties, startEventSeen, ResolverResult, getTunnelInformation, DockerResolverParameters, inspectDockerImage, getEmptyContextFolder, getFolderImageName, SubstitutedConfig } from './utils';
+import { createContainerProperties, startEventSeen, ResolverResult, getTunnelInformation, DockerResolverParameters, inspectDockerImage, getEmptyContextFolder, getFolderImageName, SubstitutedConfig, checkDockerSupportForGPU } from './utils';
 import { ContainerProperties, setupInContainer, ResolverProgress } from '../spec-common/injectHeadless';
 import { ContainerError } from '../spec-common/errors';
 import { Workspace } from '../spec-utils/workspaces';
@@ -392,7 +392,8 @@ async function startContainer(params: DockerResolverParameters, buildParams: Doc
 		// Persisted folder is a path that will be maintained between sessions
 		// Note: As a fallback, persistedFolder is set to the build's tmpDir() directory
 		const additionalLabels = labels ? idLabels.concat(Object.keys(labels).map(key => `${key}=${labels[key]}`)) : idLabels;
-		const overrideFilePath = await writeFeaturesComposeOverrideFile(updatedImageName, currentImageName, mergedConfig, config, versionPrefix, imageDetails, service, additionalLabels, params.additionalMounts, persistedFolder, featuresStartOverrideFilePrefix, buildCLIHost, output);
+		const overrideFilePath = await writeFeaturesComposeOverrideFile(updatedImageName, currentImageName, mergedConfig, config, versionPrefix, imageDetails, service, additionalLabels, params.additionalMounts, persistedFolder, featuresStartOverrideFilePrefix, buildCLIHost, params, output);
+
 		if (overrideFilePath) {
 			// Add file path to override file as parameter
 			composeGlobalArgs.push('-f', overrideFilePath);
@@ -455,9 +456,10 @@ async function writeFeaturesComposeOverrideFile(
 	overrideFilePath: string,
 	overrideFilePrefix: string,
 	buildCLIHost: CLIHost,
+	params: DockerResolverParameters,
 	output: Log,
 ) {
-	const composeOverrideContent = await generateFeaturesComposeOverrideContent(updatedImageName, originalImageName, mergedConfig, config, versionPrefix, imageDetails, service, additionalLabels, additionalMounts);
+	const composeOverrideContent = await generateFeaturesComposeOverrideContent(updatedImageName, originalImageName, mergedConfig, config, versionPrefix, imageDetails, service, additionalLabels, additionalMounts, params);
 	const overrideFileHasContents = !!composeOverrideContent && composeOverrideContent.length > 0 && composeOverrideContent.trim() !== '';
 	if (overrideFileHasContents) {
 		output.write(`Docker Compose override file for creating container:\n${composeOverrideContent}`);
@@ -486,6 +488,7 @@ async function generateFeaturesComposeOverrideContent(
 	service: any,
 	additionalLabels: string[],
 	additionalMounts: Mount[],
+	params: DockerResolverParameters,
 ) {
 	const overrideImage = updatedImageName !== originalImageName;
 
@@ -507,6 +510,18 @@ async function generateFeaturesComposeOverrideContent(
 	const userCommand = overrideCommand ? [] : composeCommand /* $ already escaped. */
 		|| (composeEntrypoint ? [/* Ignore image CMD per docker-compose.yml spec. */] : ((await imageDetails()).Config.Cmd || []).map(c => c.replace(/\$/g, '$$$$'))); // $ > $$ to escape docker-compose.yml's interpolation.
 
+	const hasGpuRequirement = config.hostRequirements?.gpu;
+	const addGpuCapability = hasGpuRequirement && await checkDockerSupportForGPU(params);
+	if (hasGpuRequirement && hasGpuRequirement !== 'optional' && !addGpuCapability) {
+		params.common.output.write('No GPU support found yet a GPU was required - consider marking it as "optional"', LogLevel.Warning);
+	}
+	const gpuResources = addGpuCapability ? `
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - capabilities: [gpu]` : '';
+
 	return `${versionPrefix}services:
   '${config.service}':${overrideImage ? `
     image: ${updatedImageName}` : ''}
@@ -528,7 +543,7 @@ while sleep 1 & wait $$!; do :; done", "-"${userEntrypoint.map(a => `, ${JSON.st
     labels:${additionalLabels.map(label => `
       - ${label.replace(/\$/g, '$$$$')}`).join('')}` : ''}${mounts.length ? `
     volumes:${mounts.map(m => `
-      - ${m.source}:${m.target}`).join('')}` : ''}${volumeMounts.length ? `
+      - ${m.source}:${m.target}`).join('')}` : ''}${gpuResources}${volumeMounts.length ? `
 volumes:${volumeMounts.map(m => `
   ${m.source}:${m.external ? '\n    external: true' : ''}`).join('')}` : ''}
 `;

src/spec-node/imageMetadata.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { ContainerError } from '../spec-common/errors';
-import { DevContainerConfig, DevContainerConfigCommand, DevContainerFromDockerComposeConfig, DevContainerFromDockerfileConfig, DevContainerFromImageConfig, getDockerComposeFilePaths, getDockerfilePath, HostRequirements, isDockerFileConfig, PortAttributes, UserEnvProbe } from '../spec-configuration/configuration';
+import { DevContainerConfig, DevContainerConfigCommand, DevContainerFromDockerComposeConfig, DevContainerFromDockerfileConfig, DevContainerFromImageConfig, getDockerComposeFilePaths, getDockerfilePath, HostGPURequirements, HostRequirements, isDockerFileConfig, PortAttributes, UserEnvProbe } from '../spec-configuration/configuration';
 import { Feature, FeaturesConfig, Mount, parseMount } from '../spec-configuration/containerFeaturesConfiguration';
 import { ContainerDetails, DockerCLIParameters, ImageDetails } from '../spec-shutdown/dockerUtils';
 import { Log } from '../spec-utils/log';
@@ -169,13 +169,43 @@ function mergeHostRequirements(imageMetadata: ImageMetadataEntry[]) {
 	const cpus = Math.max(...imageMetadata.map(m => m.hostRequirements?.cpus || 0));
 	const memory = Math.max(...imageMetadata.map(m => parseBytes(m.hostRequirements?.memory || '0')));
 	const storage = Math.max(...imageMetadata.map(m => parseBytes(m.hostRequirements?.storage || '0')));
-	return cpus || memory || storage ? {
+	const gpu = imageMetadata.map(m => m.hostRequirements?.gpu).reduce(mergeGpuRequirements, undefined);
+	return cpus || memory || storage || gpu ? {
 		cpus,
 		memory: memory ? `${memory}` : undefined,
 		storage: storage ? `${storage}` : undefined,
+		gpu: gpu,
 	} : undefined;
 }
 
+function mergeGpuRequirements(a: undefined | boolean | 'optional' | HostGPURequirements, b: undefined | boolean | 'optional' | HostGPURequirements): undefined | boolean | 'optional' | HostGPURequirements {
+	// simple cases if either are undefined/false we use the other one
+	if (a === undefined || a === false) {
+		return b;
+	} else if (b === undefined || b === false) {
+		return a;
+	} else if (a === 'optional' && b === 'optional ') {
+		return 'optional';
+	} else {
+		const aObject = asHostGPURequirements(a);
+		const bObject = asHostGPURequirements(b);
+		const cores = Math.max(aObject.cores || 0, bObject.cores || 0);
+		const memory = Math.max(parseBytes(aObject.memory || '0'), parseBytes(bObject.memory || '0'));
+		return {
+			cores: cores ? cores : undefined,
+			memory: memory ? `${memory}` : undefined,
+		};
+	}
+}
+
+function asHostGPURequirements(a: undefined | boolean | 'optional' | HostGPURequirements): HostGPURequirements {
+	if (typeof a !== 'object') {
+		return {};
+	} else {
+		return a as HostGPURequirements;
+	}
+}
+
 function parseBytes(str: string) {
 	const m = /^(\d+)([tgmk]b)?$/.exec(str);
 	if (m) {

src/spec-node/singleContainer.ts

@@ -4,8 +4,8 @@
  *--------------------------------------------------------------------------------------------*/
 
 
-import { createContainerProperties, startEventSeen, ResolverResult, getTunnelInformation, getDockerfilePath, getDockerContextPath, DockerResolverParameters, isDockerFileConfig, uriToWSLFsPath, WorkspaceConfiguration, getFolderImageName, inspectDockerImage, logUMask, SubstitutedConfig } from './utils';
-import { ContainerProperties, setupInContainer, ResolverProgress } from '../spec-common/injectHeadless';
+import { createContainerProperties, startEventSeen, ResolverResult, getTunnelInformation, getDockerfilePath, getDockerContextPath, DockerResolverParameters, isDockerFileConfig, uriToWSLFsPath, WorkspaceConfiguration, getFolderImageName, inspectDockerImage, logUMask, SubstitutedConfig, checkDockerSupportForGPU } from './utils';
+import { ContainerProperties, setupInContainer, ResolverProgress, ResolverParameters } from '../spec-common/injectHeadless';
 import { ContainerError, toErrorText } from '../spec-common/errors';
 import { ContainerDetails, listContainers, DockerCLIParameters, inspectContainers, dockerCLI, dockerPtyCLI, toPtyExecParameters, ImageDetails, toExecParameters } from '../spec-shutdown/dockerUtils';
 import { DevContainerConfig, DevContainerFromDockerfileConfig, DevContainerFromImageConfig } from '../spec-configuration/configuration';
@@ -305,6 +305,20 @@ 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) {
+	const extraArguments: string[] = [];
+	if (config.hostRequirements?.gpu) {
+		if (await checkDockerSupportForGPU(params)) {
+			common.output.write(`GPU support found, add GPU flags to docker call.`);
+			extraArguments.push('--gpus', 'all');
+		} else {
+			if (config.hostRequirements?.gpu !== 'optional') {
+				common.output.write('No GPU support found yet a GPU was required - consider marking it as "optional"', LogLevel.Warning);
+			}
+		}
+	}
+	return extraArguments;
+}
 
 export async function spawnDevContainer(params: DockerResolverParameters, config: DevContainerFromDockerfileConfig | DevContainerFromImageConfig, mergedConfig: MergedDevContainerConfig, imageName: string, labels: string[], workspaceMount: string | undefined, imageDetails: (() => Promise<ImageDetails>) | undefined, containerUser: string | undefined, extraLabels: Record<string, string>) {
 	const { common } = params;
@@ -372,6 +386,7 @@ while sleep 1 & wait $!; do :; done`, '-']; // `wait $!` allows for the `trap` t
 		...containerEnv,
 		...containerUserArgs,
 		...(config.runArgs || []),
+		...(await extraRunArgs(common, params, config) || []),
 		...featureArgs,
 		...entrypoint,
 		...Object.keys(extraLabels).map(key => ['-l', `${key}=${extraLabels[key]}`]).flat(),

src/spec-node/utils.ts

@@ -15,7 +15,7 @@ import { ContainerProperties, getContainerProperties, ResolverParameters } from
 import { Workspace } from '../spec-utils/workspaces';
 import { URI } from 'vscode-uri';
 import { ShellServer } from '../spec-common/shellServer';
-import { inspectContainer, inspectImage, getEvents, ContainerDetails, DockerCLIParameters, dockerExecFunction, dockerPtyCLI, dockerPtyExecFunction, toDockerImageName, DockerComposeCLI, ImageDetails } from '../spec-shutdown/dockerUtils';
+import { inspectContainer, inspectImage, getEvents, ContainerDetails, DockerCLIParameters, dockerExecFunction, dockerPtyCLI, dockerPtyExecFunction, toDockerImageName, DockerComposeCLI, ImageDetails, dockerCLI } from '../spec-shutdown/dockerUtils';
 import { getRemoteWorkspaceFolder } from './dockerCompose';
 import { findGitRootFolder } from '../spec-common/git';
 import { parentURI, uriToFsPath } from '../spec-configuration/configurationCommonUtils';
@@ -178,6 +178,12 @@ async function hasLabels(params: DockerResolverParameters, info: any, expectedLa
 		.every(name => actualLabels[name] === expectedLabels[name]);
 }
 
+export async function checkDockerSupportForGPU(params: DockerCLIParameters | DockerResolverParameters): Promise<Boolean> {
+	const result = await dockerCLI(params, 'info', '-f', '{{.Runtimes.nvidia}}');
+	const runtimeFound = result.stdout.includes('nvidia-container-runtime');
+	return runtimeFound;
+}
+
 export async function inspectDockerImage(params: DockerResolverParameters | DockerCLIParameters, imageName: string, pullImageOnError: boolean) {
 	try {
 		return await inspectImage(params, imageName);

src/test/imageMetadata.test.ts

@@ -6,6 +6,7 @@
 import * as assert from 'assert';
 import * as path from 'path';
 import { URI } from 'vscode-uri';
+import { HostGPURequirements } from '../spec-configuration/configuration';
 import { Feature, FeaturesConfig, FeatureSet, Mount } from '../spec-configuration/containerFeaturesConfiguration';
 import { experimentalImageMetadataDefault } from '../spec-node/devContainers';
 import { getDevcontainerMetadata, getDevcontainerMetadataLabel, getImageMetadata, getImageMetadataFromContainer, imageMetadataLabel, internalGetImageMetadata0, mergeConfiguration } from '../spec-node/imageMetadata';
@@ -395,6 +396,39 @@ describe('Image Metadata', function () {
 			assert.strictEqual((merged.mounts?.[2] as Mount).source, 'source5');
 		});
 	});
+
+	it('should merge gpu requirements from devcontainer.json and features', () => {
+		const merged = mergeConfiguration({
+			configFilePath: URI.parse('file:///devcontainer.json'),
+			image: 'image',
+			hostRequirements: {
+				gpu: 'optional'
+			}
+		}, [
+			{
+				hostRequirements: {
+					gpu: true
+				}
+			},
+			{
+				hostRequirements: {
+					gpu: {
+						cores: 4
+					}
+				}
+			},
+			{
+				hostRequirements: {
+					gpu: {
+						memory: '8gb'
+					}
+				}
+			}
+		]);
+		const gpuRequirement = merged.hostRequirements?.gpu as HostGPURequirements;
+		assert.strictEqual(gpuRequirement?.cores, 4);
+		assert.strictEqual(gpuRequirement?.memory, '8589934592');
+	});
 });
 
 function getFeaturesConfig(features: Feature[]): FeaturesConfig {