Skip to content

Commit

Permalink
Add option for GPU availability (microsoft/vscode-remote-release#9385)
Browse files Browse the repository at this point in the history
  • Loading branch information
chrmarti committed Sep 10, 2024
1 parent cbfbdb1 commit 4e07a4d
Show file tree
Hide file tree
Showing 4 changed files with 18 additions and 4 deletions.
6 changes: 4 additions & 2 deletions src/spec-node/devContainers.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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',
Expand Down
3 changes: 3 additions & 0 deletions src/spec-node/devContainersSpecCLI.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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.' },
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion src/spec-node/singleContainer.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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)) {
Expand Down
11 changes: 10 additions & 1 deletion src/spec-node/utils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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';
Expand Down Expand Up @@ -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;
Expand Down

0 comments on commit 4e07a4d

Please sign in to comment.