From ac550aeb021c2dfef0e119ee157cabf7b70b6a6d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Robert=20G=C3=BCnzler?= <r@gnzler.io>
Date: Fri, 5 Jun 2020 16:53:50 +0200
Subject: [PATCH] Add label to expose gpu to container
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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:
https://github.com/compose-spec/compose-spec/issues/74
https://github.com/docker/compose/issues/6691

Change-type: patch
Connects-to: https://github.com/balena-os/balena-jetson/pull/75
Signed-off-by: Robert Günzler <robertg@balena.io>
---
 src/compose/service.ts       |  3 +++
 src/compose/types/service.ts |  1 +
 src/compose/utils.ts         |  9 +++++++++
 test/04-service.spec.ts      | 23 +++++++++++++++++++++++
 4 files changed, 36 insertions(+)

diff --git a/src/compose/service.ts b/src/compose/service.ts
index b6ce0ae123..615023f6a4 100644
--- a/src/compose/service.ts
+++ b/src/compose/service.ts
@@ -371,6 +371,7 @@ export class Service {
 			command: [],
 			cgroupParent: '',
 			devices,
+			deviceRequests: [],
 			dnsOpt: [],
 			entrypoint: '',
 			extraHosts: [],
@@ -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,
@@ -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,
diff --git a/src/compose/types/service.ts b/src/compose/types/service.ts
index 62ba845394..f8fa8d1d59 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 2cb30ded2d..47a6090baa 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 cd708c1ed0..17701ee575 100644
--- a/test/04-service.spec.ts
+++ b/test/04-service.spec.ts
@@ -464,6 +464,29 @@ describe('compose/service', () => {
 		});
 	});
 
+	it('should include a device request for a gpu device when the gpu feature label is set', () => {
+		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([
+				{
+					Count: 1,
+					Capabilities: [['gpu']],
+				},
+			]);
+	});
+
 	describe('Docker <-> Compose config', () => {
 		const omitConfigForComparison = (config: ServiceConfig) =>
 			_.omit(config, ['running', 'networks']);