A support for GPU (#1866)

Author: zapadinsky

src/argv.ts

@@ -388,6 +388,10 @@ export class Argv {
         return this.map.get("containerEmulate") ?? null;
     }
 
+    get gpus (): string | null {
+        return this.map.get("gpus") ?? null;
+    }
+
     get concurrency (): number | null {
         const concurrency = this.map.get("concurrency");
         if (!concurrency) return null;

src/index.ts

@@ -366,6 +366,11 @@ process.on("SIGUSR2", async () => {
             description: "The name, without the architecture, of a gitlab hosted runner to emulate. See here: https://docs.gitlab.com/ee/ci/runners/hosted_runners/linux.html#machine-types-available-for-linux---x86-64",
             choices: GitlabRunnerPresetValues,
         })
+        .option("gpus", {
+            type: "string",
+            description: "A list of GPU indexes or 'all'",
+            requiresArg: true,
+        })
         .option("color", {
             requiresArg: false,
             default: true,

src/job.ts

@@ -986,6 +986,10 @@ If you know what you're doing and would like to suppress this warning, use one o
                 dockerCmd += `--cpus=${cpuConfig} `;
             }
 
+            if (this.gpus) {
+                dockerCmd += `--gpus ${this.gpus} --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 `;
+            }
+
             // host and none networks have to be specified using --network, since they cannot be used with
             // `docker network connect`.
             for (const network of this.argv.network) {
@@ -1219,6 +1223,19 @@ If you know what you're doing and would like to suppress this warning, use one o
         return image.entrypoint;
     }
 
+    get gpus (): string | null {
+        if (this.argv.gpus) {
+            return this.argv.gpus;
+        } else if ("tags" in this.jobData) {
+            for (const tag of this.jobData["tags"]) {
+                if (tag.match(/^(.*-)?gpu(-.*)?$/)) {
+                    return "all";
+                }
+            }
+        }
+        return null;
+    }
+
     private async validateCiDependencyProxyServerAuthentication (imageName: string) {
         const CI_DEPENDENCY_PROXY_SERVER = this._variables["CI_DEPENDENCY_PROXY_SERVER"];
         if (!imageName.startsWith(CI_DEPENDENCY_PROXY_SERVER)) {
@@ -1579,6 +1596,10 @@ If you know what you're doing and would like to suppress this warning, use one o
             dockerCmd += `--shm-size=${this.argv.shmSize} `;
         }
 
+        if (this.gpus) {
+            dockerCmd += `--gpus ${this.gpus} --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 `;
+        }
+
         for (const volume of this.argv.volume) {
             dockerCmd += `--volume ${volume} `;
         }