added GPU requirement logic to docker compose flow using override files

Author: chris-major-improbable

src/spec-node/dockerCompose.ts

@@ -230,6 +230,25 @@ ${cacheFromOverrideContent}
 		args.push('-f', composeOverrideFile);
 	}
 
+	const hasGpuRequirement = config.hostRequirements?.gpu;
+	const addGpuCapability = hasGpuRequirement && await checkDockerSupportForGPU(params);
+	if (addGpuCapability) {
+		const composeOverrideContent = `services:
+  ${config.service}:
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - capabilities: [gpu]
+`;
+		const composeFolder = cliHost.path.join(overrideFilePath, 'docker-compose');
+		await cliHost.mkdirp(composeFolder);
+		const composeOverrideFile = cliHost.path.join(composeFolder, `${overrideFilePrefix}-gpu-${Date.now()}.yml`);
+		await cliHost.writeFile(composeOverrideFile, Buffer.from(composeOverrideContent));
+		additionalComposeOverrideFiles.push(composeOverrideFile);
+		args.push('-f', composeOverrideFile);
+	}
+
 	if (!noBuild) {
 		args.push('build');
 		if (noCache) {
@@ -291,6 +310,13 @@ async function checkForPersistedFile(cliHost: CLIHost, output: Log, files: strin
 		foundLabel: false
 	};
 }
+
+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;
+}
+
 async function startContainer(params: DockerResolverParameters, buildParams: DockerCLIParameters, config: DevContainerFromDockerComposeConfig, projectName: string, composeFiles: string[], envFile: string | undefined, container: ContainerDetails | undefined, idLabels: string[]) {
 	const { common } = params;
 	const { persistedFolder, output } = common;
@@ -300,15 +326,14 @@ async function startContainer(params: DockerResolverParameters, buildParams: Doc
 
 	common.progress(ResolverProgress.StartingContainer);
 
-	const localComposeFiles = composeFiles;
 	// If dockerComposeFile is an array, add -f <file> in order. https://docs.docker.com/compose/extends/#multiple-compose-files
-	const composeGlobalArgs = ([] as string[]).concat(...localComposeFiles.map(composeFile => ['-f', composeFile]));
+	const composeGlobalArgs = ([] as string[]).concat(...composeFiles.map(composeFile => ['-f', composeFile]));
 	if (envFile) {
 		composeGlobalArgs.push('--env-file', envFile);
 	}
 
 	const infoOutput = makeLog(buildParams.output, LogLevel.Info);
-	const composeConfig = await readDockerComposeConfig(buildParams, localComposeFiles, envFile);
+	const composeConfig = await readDockerComposeConfig(buildParams, composeFiles, envFile);
 	const services = Object.keys(composeConfig.services || {});
 	if (services.indexOf(config.service) === -1) {
 		throw new ContainerError({ description: `Service '${config.service}' configured in devcontainer.json not found in Docker Compose configuration.`, data: { fileWithError: composeFiles[0] } });
@@ -355,7 +380,7 @@ async function startContainer(params: DockerResolverParameters, buildParams: Doc
 		const noBuild = !!container; //if we have an existing container, just recreate override files but skip the build
 
 		const infoParams = { ...params, common: { ...params.common, output: infoOutput } };
-		const { collapsedFeaturesConfig, additionalComposeOverrideFiles } = await buildAndExtendDockerCompose(config, projectName, infoParams, localComposeFiles, envFile, composeGlobalArgs, config.runServices ?? [], params.buildNoCache ?? false, persistedFolder, featuresBuildOverrideFilePrefix, params.additionalCacheFroms, noBuild);
+		const { collapsedFeaturesConfig, additionalComposeOverrideFiles } = await buildAndExtendDockerCompose(config, projectName, infoParams, composeFiles, envFile, composeGlobalArgs, config.runServices ?? [], params.buildNoCache ?? false, persistedFolder, featuresBuildOverrideFilePrefix, params.additionalCacheFroms, noBuild);
 		additionalComposeOverrideFiles.forEach(overrideFilePath => composeGlobalArgs.push('-f', overrideFilePath));
 
 		let cache: Promise<ImageDetails> | undefined;
@@ -391,7 +416,7 @@ async function startContainer(params: DockerResolverParameters, buildParams: Doc
 		}
 	} catch (err) {
 		cancel!();
-		throw new ContainerError({ description: 'An error occurred starting Docker Compose up.', originalError: err, data: { fileWithError: localComposeFiles[0] } });
+		throw new ContainerError({ description: 'An error occurred starting Docker Compose up.', originalError: err, data: { fileWithError: composeFiles[0] } });
 	}
 
 	await started;

src/spec-node/singleContainer.ts

@@ -5,7 +5,7 @@
 
 
 import { createContainerProperties, startEventSeen, ResolverResult, getTunnelInformation, getDockerfilePath, getDockerContextPath, DockerResolverParameters, isDockerFileConfig, uriToWSLFsPath, WorkspaceConfiguration, getFolderImageName, inspectDockerImage, ensureDockerfileHasFinalStageName, getImageUser, logUMask } from './utils';
-import { ContainerProperties, setupInContainer, ResolverProgress } from '../spec-common/injectHeadless';
+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';
@@ -291,12 +291,13 @@ 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) {
+export async function extraRunArgs(common: ResolverParameters, 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) {
+			common.output.write(`GPU support found, add GPU flags to docker call.`);
 			extraArguments.push('--gpus', 'all');
 		}
 	}
@@ -381,7 +382,7 @@ while sleep 1 & wait $!; do :; done`, '-']; // `wait $!` allows for the `trap` t
 		...containerEnv,
 		...containerUser,
 		...(config.runArgs || []),
-		...(await extraRunArgs(params, config) || []),
+		...(await extraRunArgs(common, params, config) || []),
 		...featureArgs,
 		...entrypoint,
 		imageName,