docker/buildx
1
0
Fork 0
mirror of https://gitea.com/Lydanne/buildx.git synced 2025-05-18 00:47:48 +08:00
Code Issues Packages Projects Releases 1 Wiki Activity

Merge pull request #3063 from crazy-max/driver-ctn-gpu-request

Browse Source 
driver: request gpu when creating container builder
This commit is contained in:
Tõnis Tiigi 2025-03-13 09:56:10 -07:00 committed by GitHub
commit 97f1d47464
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 35 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
driver/docker-container/driver.go
View File

@ -56,6 +56,7 @@ type Driver struct {
restartPolicy container.RestartPolicy
env []string
defaultLoad bool
gpus []container.DeviceRequest
}
func (d *Driver) IsMobyDriver() bool {
@ -158,6 +159,9 @@ func (d *Driver) create(ctx context.Context, l progress.SubLogger) error {
if d.cpusetMems != "" {
hc.Resources.CpusetMems = d.cpusetMems
}
if len(d.gpus) > 0 && d.hasGPUCapability(ctx, cfg.Image, d.gpus) {
hc.Resources.DeviceRequests = d.gpus
}
if info, err := d.DockerAPI.Info(ctx); err == nil {
if info.CgroupDriver == "cgroupfs" {
// Place all buildkit containers inside this cgroup by default so limits can be attached
@ -429,6 +433,31 @@ func (d *Driver) HostGatewayIP(ctx context.Context) (net.IP, error) {
return nil, errors.New("host-gateway is not supported by the docker-container driver")
}
// hasGPUCapability checks if docker daemon has GPU capability. We need to run
// a dummy container with GPU device to check if the daemon has this capability
// because there is no API to check it yet.
func (d *Driver) hasGPUCapability(ctx context.Context, image string, gpus []container.DeviceRequest) bool {
cfg := &container.Config{
Image: image,
Entrypoint: []string{"/bin/true"},
}
hc := &container.HostConfig{
NetworkMode: container.NetworkMode(container.IPCModeNone),
AutoRemove: true,
Resources: container.Resources{
DeviceRequests: gpus,
},
}
resp, err := d.DockerAPI.ContainerCreate(ctx, cfg, hc, &network.NetworkingConfig{}, nil, "")
if err != nil {
return false
}
if err := d.DockerAPI.ContainerStart(ctx, resp.ID, container.StartOptions{}); err != nil {
return false
}
return true
}
func demuxConn(c net.Conn) net.Conn {
pr, pw := io.Pipe()
// TODO: rewrite parser with Reader() to avoid goroutine switch

6
driver/docker-container/factory.go
View File

@ -51,6 +51,12 @@ func (f *factory) New(ctx context.Context, cfg driver.InitConfig) (driver.Driver
InitConfig: cfg,
restartPolicy: rp,
}
var gpus dockeropts.GpuOpts
if err := gpus.Set("all"); err == nil {
if v := gpus.Value(); len(v) > 0 {
d.gpus = v
}
}
for k, v := range cfg.DriverOpts {
switch {
case k == "network":