Skip to content

Commit

Permalink
fix: avoid remounting GPU drivers
Browse files Browse the repository at this point in the history
  • Loading branch information
sreya committed Feb 13, 2024
1 parent 4ffaa0e commit 9da96fd
Showing 1 changed file with 34 additions and 31 deletions.
65 changes: 34 additions & 31 deletions cli/docker.go
Original file line number Diff line number Diff line change
Expand Up @@ -472,43 +472,13 @@ func runDockerCVM(ctx context.Context, log slog.Logger, client dockerutil.Docker
if flags.hostUsrLibDir == "" {
return xerrors.Errorf("when using GPUs, %q must be specified", EnvUsrLibDir)
}

// Unmount GPU drivers in /proc as it causes issues when creating any
// container in some cases (even the image metadata container).
_, err = xunix.TryUnmountProcGPUDrivers(ctx, log)
if err != nil {
return xerrors.Errorf("unmount /proc GPU drivers: %w", err)
}

devs, binds, err := xunix.GPUs(ctx, log, flags.hostUsrLibDir)
if err != nil {
return xerrors.Errorf("find gpus: %w", err)
}

for _, dev := range devs {
devices = append(devices, container.DeviceMapping{
PathOnHost: dev.Path,
PathInContainer: dev.Path,
CgroupPermissions: "rwm",
})
}

for _, bind := range binds {
// If the bind has a path that points to the host-mounted /usr/lib
// directory we need to remap it to /usr/lib inside the container.
mountpoint := bind.Path
if strings.HasPrefix(mountpoint, flags.hostUsrLibDir) {
mountpoint = filepath.Join(
"/usr/lib",
strings.TrimPrefix(mountpoint, strings.TrimSuffix(flags.hostUsrLibDir, "/")),
)
}
mounts = append(mounts, xunix.Mount{
Source: bind.Path,
Mountpoint: mountpoint,
ReadOnly: slices.Contains(bind.Opts, "ro"),
})
}
envs = append(envs, xunix.GPUEnvs(ctx)...)
}

log.Debug(ctx, "fetching image metadata",
Expand Down Expand Up @@ -598,6 +568,39 @@ func runDockerCVM(ctx context.Context, log slog.Logger, client dockerutil.Docker
}
}

if flags.addGPU {
devs, binds, err := xunix.GPUs(ctx, log, flags.hostUsrLibDir)
if err != nil {
return xerrors.Errorf("find gpus: %w", err)
}

for _, dev := range devs {
devices = append(devices, container.DeviceMapping{
PathOnHost: dev.Path,
PathInContainer: dev.Path,
CgroupPermissions: "rwm",
})
}

for _, bind := range binds {
// If the bind has a path that points to the host-mounted /usr/lib
// directory we need to remap it to /usr/lib inside the container.
mountpoint := bind.Path
if strings.HasPrefix(mountpoint, flags.hostUsrLibDir) {
mountpoint = filepath.Join(
"/usr/lib",
strings.TrimPrefix(mountpoint, strings.TrimSuffix(flags.hostUsrLibDir, "/")),
)
}
mounts = append(mounts, xunix.Mount{
Source: bind.Path,
Mountpoint: mountpoint,
ReadOnly: slices.Contains(bind.Opts, "ro"),
})
}
envs = append(envs, xunix.GPUEnvs(ctx)...)
}

blog.Info("Creating workspace...")

// Create the inner container.
Expand Down

0 comments on commit 9da96fd

Please sign in to comment.