Merge pull request #1362 from balena-io/rgz/device-request-gpu

Add label to expose gpu to container

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,

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[];

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) => {

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']);