feat: add --gpus option

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.argv.gpus) {
+                dockerCmd += `--gpus ${this.argv.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) {
@@ -1579,6 +1583,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.argv.gpus) {
+            dockerCmd += `--gpus ${this.argv.gpus} --ipc=host --ulimit memlock=-1 --ulimit stack=67108864`;
+        }
+
         for (const volume of this.argv.volume) {
             dockerCmd += `--volume ${volume} `;
         }