From 152e8cacf8ef5aa5ceab91b030f2bc34faa41bd0 Mon Sep 17 00:00:00 2001
From: Anatoly Zapadinsky <anatoly@bladerunner.io>
Date: Thu, 6 Mar 2025 18:37:31 +0300
Subject: [PATCH 1/2] feat: add --gpus option

---
 src/argv.ts  | 4 ++++
 src/index.ts | 5 +++++
 src/job.ts   | 8 ++++++++
 3 files changed, 17 insertions(+)

diff --git a/src/argv.ts b/src/argv.ts
index 1d360caad..cb9a7928c 100644
--- a/src/argv.ts
+++ b/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;
diff --git a/src/index.ts b/src/index.ts
index beb1002de..b57f9ab24 100644
--- a/src/index.ts
+++ b/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,
diff --git a/src/job.ts b/src/job.ts
index f37d82cbb..c4fbaa9bb 100644
--- a/src/job.ts
+++ b/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} `;
         }

From baaca82bcb5c8b79f549d5e716466d8641df7179 Mon Sep 17 00:00:00 2001
From: Anatoly Zapadinsky <anatoly@bladerunner.io>
Date: Wed, 12 Mar 2025 20:13:07 +0300
Subject: [PATCH 2/2] feat: add --gpus automatically if tag specified

---
 src/job.ts | 21 +++++++++++++++++----
 1 file changed, 17 insertions(+), 4 deletions(-)

diff --git a/src/job.ts b/src/job.ts
index c4fbaa9bb..0b6dfab5d 100644
--- a/src/job.ts
+++ b/src/job.ts
@@ -986,8 +986,8 @@ 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 `;
+            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
@@ -1223,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)) {
@@ -1583,8 +1596,8 @@ 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`;
+        if (this.gpus) {
+            dockerCmd += `--gpus ${this.gpus} --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 `;
         }
 
         for (const volume of this.argv.volume) {