Add label to expose gpu to container

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]>
balena-os · Jun 9, 2020 · ac550ae · ac550ae
1 parent dde1644
commit ac550ae
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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']);