Skip to content

Commit

Permalink
Add label to expose gpu to container
Browse files Browse the repository at this point in the history
In the absence of an upstream implementation of the DeviceRequest API introduced
as part of Docker API v1.40 we roll our own using a feature label.

As per my comment in the code, we fall back to the default behavior of
docker cli's `--gpu` and request single device with the `gpu` capabilty.
The only implementation at the moment is the NVIDIA driver; here:
https://github.com/balena-os/balena-engine/blob/master/daemon/nvidia_linux.go

Background on the composefile implementation:
compose-spec/compose-spec#74
docker/compose#6691

Change-type: patch
Connects-to: balena-os/balena-jetson#75
Signed-off-by: Robert Günzler <[email protected]>
  • Loading branch information
robertgzr committed Jun 11, 2020
1 parent 5512654 commit b2e8fe8
Show file tree
Hide file tree
Showing 4 changed files with 47 additions and 0 deletions.
3 changes: 3 additions & 0 deletions src/compose/service.ts
Original file line number Diff line number Diff line change
Expand Up @@ -371,6 +371,7 @@ export class Service {
command: [],
cgroupParent: '',
devices,
deviceRequests: [],
dnsOpt: [],
entrypoint: '',
extraHosts: [],
Expand Down Expand Up @@ -515,6 +516,7 @@ export class Service {
capAdd: container.HostConfig.CapAdd || [],
capDrop: container.HostConfig.CapDrop || [],
devices: container.HostConfig.Devices || [],
deviceRequests: container.HostConfig.DeviceRequests || [],
networks,
memLimit: container.HostConfig.Memory || 0,
memReservation: container.HostConfig.MemoryReservation || 0,
Expand Down Expand Up @@ -635,6 +637,7 @@ export class Service {
Binds: binds,
CgroupParent: this.config.cgroupParent,
Devices: this.config.devices,
DeviceRequests: this.config.deviceRequests,
Dns: this.config.dns,
DnsOptions: this.config.dnsOpt,
DnsSearch: this.config.dnsSearch,
Expand Down
1 change: 1 addition & 0 deletions src/compose/types/service.ts
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,7 @@ export interface ServiceConfig {
command: string[];
cgroupParent: string;
devices: DockerDevice[];
deviceRequests: Dockerode.DeviceRequest[];
dns: string | string[];
dnsOpt: string[];
dnsSearch: string | string[];
Expand Down
9 changes: 9 additions & 0 deletions src/compose/utils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -373,6 +373,15 @@ export function addFeaturesFromLabels(
'io.balena.features.sysfs': () => service.config.volumes.push('/sys:/sys'),
'io.balena.features.procfs': () =>
service.config.volumes.push('/proc:/proc'),
'io.balena.features.gpu': () =>
// TODO once the compose-spec has an implementation we
// should probably follow that, for now we copy the
// bahavior of docker cli
// https://github.com/balena-os/balena-engine-cli/blob/19.03-balena/opts/gpus.go#L81-L89
service.config.deviceRequests.push({
Count: 1,
Capabilities: [['gpu']],
} as Dockerode.DeviceRequest),
};

_.each(features, (fn, label) => {
Expand Down
34 changes: 34 additions & 0 deletions test/04-service.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -464,6 +464,40 @@ describe('compose/service', () => {
});
});

describe('io.balena.features.gpu: Docker <-> Compose config', () => {
const gpuDeviceRequest = {
Count: 1,
Capabilities: [['gpu']],
};
it('should succeed from compose object', () => {
const s = Service.fromComposeObject(
{
appId: 123,
serviceId: 123,
serviceName: 'test',
labels: {
'io.balena.features.gpu': '1',
},
},
{ appName: 'test' } as any,
);

expect(s.config)
.to.have.property('deviceRequests')
.that.deep.equals([gpuDeviceRequest]);
});

it('should succeed from docker container', () => {
const dockerCfg = require('./data/docker-states/simple/inspect.json');
dockerCfg.HostConfig.DeviceRequests = [gpuDeviceRequest];
const s = Service.fromDockerContainer(dockerCfg);

expect(s.config)
.to.have.property('deviceRequests')
.that.deep.equals([gpuDeviceRequest]);
});
});

describe('Docker <-> Compose config', () => {
const omitConfigForComparison = (config: ServiceConfig) =>
_.omit(config, ['running', 'networks']);
Expand Down

0 comments on commit b2e8fe8

Please sign in to comment.