From f533ca1cb8166ffd51535aee903613760f42ca74 Mon Sep 17 00:00:00 2001 From: Chris Major Date: Fri, 11 Nov 2022 08:01:01 +0000 Subject: [PATCH] GPU requirement and auto-detect NVIDIA extensions (#173) --- src/spec-configuration/configuration.ts | 6 +++++ src/spec-node/dockerCompose.ts | 23 ++++++++++++++--- src/spec-node/imageMetadata.ts | 34 +++++++++++++++++++++++-- src/spec-node/singleContainer.ts | 19 ++++++++++++-- src/spec-node/utils.ts | 8 +++++- src/test/imageMetadata.test.ts | 34 +++++++++++++++++++++++++ 6 files changed, 115 insertions(+), 9 deletions(-) diff --git a/src/spec-configuration/configuration.ts b/src/spec-configuration/configuration.ts index 18ae74eea..6004894dc 100644 --- a/src/spec-configuration/configuration.ts +++ b/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 { diff --git a/src/spec-node/dockerCompose.ts b/src/spec-node/dockerCompose.ts index 3cd762ccf..de1ed35b0 100644 --- a/src/spec-node/dockerCompose.ts +++ b/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('')}` : ''} `; diff --git a/src/spec-node/imageMetadata.ts b/src/spec-node/imageMetadata.ts index 836a77bee..de297a0a3 100644 --- a/src/spec-node/imageMetadata.ts +++ b/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) { diff --git a/src/spec-node/singleContainer.ts b/src/spec-node/singleContainer.ts index 135dbf031..751457755 100644 --- a/src/spec-node/singleContainer.ts +++ b/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) | undefined, containerUser: string | undefined, extraLabels: Record) { 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(), diff --git a/src/spec-node/utils.ts b/src/spec-node/utils.ts index 7412191ad..aa81a09b8 100644 --- a/src/spec-node/utils.ts +++ b/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 { + 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); diff --git a/src/test/imageMetadata.test.ts b/src/test/imageMetadata.test.ts index 0c54f5bf3..3e228c26e 100644 --- a/src/test/imageMetadata.test.ts +++ b/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 {