initial idea - adds new GPU Requirements, detects nvidia extension and updated the arguments if it's requested and exists

Author: chris-major-improbable

src/spec-configuration/configuration.ts

@@ -20,10 +20,17 @@ 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 | HostGPURequirements;
 }
 
 export interface DevContainerFeature {

src/spec-node/singleContainer.ts

@@ -291,6 +291,17 @@ export async function findDevContainer(params: DockerCLIParameters | DockerResol
 	return details.filter(container => container.State.Status !== 'removing')[0];
 }
 
+export async function extraRunArgs(params: DockerCLIParameters | DockerResolverParameters, config: DevContainerFromDockerfileConfig | DevContainerFromImageConfig) {
+	const extraArguments: string[] = [];
+	if (config.hostRequirements?.gpu) {
+		const result = await dockerCLI(params, 'info', '-f', '{{.Runtimes.nvidia}}');
+		const runtimeFound = result.stdout.includes('nvidia-container-runtime');
+		if (runtimeFound) {
+			extraArguments.push('--gpus', 'all');
+		}
+	}
+	return extraArguments;
+}
 
 export async function spawnDevContainer(params: DockerResolverParameters, config: DevContainerFromDockerfileConfig | DevContainerFromImageConfig, collapsedFeaturesConfig: CollapsedFeaturesConfig | undefined, imageName: string, labels: string[], workspaceMount: string | undefined, imageDetails: (() => Promise<ImageDetails>) | undefined) {
 	const { common } = params;
@@ -370,6 +381,7 @@ while sleep 1 & wait $!; do :; done`, '-']; // `wait $!` allows for the `trap` t
 		...containerEnv,
 		...containerUser,
 		...(config.runArgs || []),
+		...(await extraRunArgs(params, config) || []),
 		...featureArgs,
 		...entrypoint,
 		imageName,