diff --git a/src/compose/service.ts b/src/compose/service.ts index 84a92136..60984ed4 100644 --- a/src/compose/service.ts +++ b/src/compose/service.ts @@ -370,6 +370,7 @@ export class Service { command: [], cgroupParent: '', devices, + deviceRequests: [], dnsOpt: [], entrypoint: '', extraHosts: [], @@ -520,6 +521,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, @@ -640,6 +642,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, diff --git a/src/compose/types/service.ts b/src/compose/types/service.ts index 62ba8453..f8fa8d1d 100644 --- a/src/compose/types/service.ts +++ b/src/compose/types/service.ts @@ -100,6 +100,7 @@ export interface ServiceConfig { command: string[]; cgroupParent: string; devices: DockerDevice[]; + deviceRequests: Dockerode.DeviceRequest[]; dns: string | string[]; dnsOpt: string[]; dnsSearch: string | string[]; diff --git a/src/compose/utils.ts b/src/compose/utils.ts index 2cb30ded..47a6090b 100644 --- a/src/compose/utils.ts +++ b/src/compose/utils.ts @@ -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) => { diff --git a/test/04-service.spec.ts b/test/04-service.spec.ts index cd708c1e..14667d8f 100644 --- a/test/04-service.spec.ts +++ b/test/04-service.spec.ts @@ -464,6 +464,42 @@ 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 = _.cloneDeep( + 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']);