Skip to content

Commit dc2a79e

Browse files
committed
feat: add --gpus automatically if tag specified
1 parent 4bc7b53 commit dc2a79e

1 file changed

Lines changed: 17 additions & 4 deletions

File tree

src/job.ts

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -986,8 +986,8 @@ If you know what you're doing and would like to suppress this warning, use one o
986986
dockerCmd += `--cpus=${cpuConfig} `;
987987
}
988988

989-
if (this.argv.gpus) {
990-
dockerCmd += `--gpus="${this.argv.gpus}" `;
989+
if (this.gpus) {
990+
dockerCmd += `--gpus="${this.gpus}" `;
991991
dockerCmd += `--ipc=host `;
992992
dockerCmd += `--ulimit memlock=-1 `;
993993
dockerCmd += `--ulimit stack=67108864 `;
@@ -1226,6 +1226,19 @@ If you know what you're doing and would like to suppress this warning, use one o
12261226
return image.entrypoint;
12271227
}
12281228

1229+
get gpus (): string | null {
1230+
if (this.argv.gpus) {
1231+
return this.argv.gpus;
1232+
} else if ("tags" in this.jobData) {
1233+
for (const tag of this.jobData["tags"]) {
1234+
if (tag.match(/^(.*-)?gpu(-.*)?$/)) {
1235+
return "all";
1236+
}
1237+
}
1238+
}
1239+
return null;
1240+
}
1241+
12291242
private async validateCiDependencyProxyServerAuthentication (imageName: string) {
12301243
const CI_DEPENDENCY_PROXY_SERVER = this._variables["CI_DEPENDENCY_PROXY_SERVER"];
12311244
if (!imageName.startsWith(CI_DEPENDENCY_PROXY_SERVER)) {
@@ -1586,8 +1599,8 @@ If you know what you're doing and would like to suppress this warning, use one o
15861599
dockerCmd += `--shm-size=${this.argv.shmSize} `;
15871600
}
15881601

1589-
if (this.argv.gpus) {
1590-
dockerCmd += `--gpus="${this.argv.gpus}" `;
1602+
if (this.gpus) {
1603+
dockerCmd += `--gpus="${this.gpus}" `;
15911604
dockerCmd += `--ipc=host `;
15921605
dockerCmd += `--ulimit memlock=-1 `;
15931606
dockerCmd += `--ulimit stack=67108864 `;

0 commit comments

Comments
 (0)