[ci] Limit num gpus if specified for A100 (#5694)

Signed-off-by: kevin <kevin@anyscale.com>
2024-06-19 16:30:03 -07:00
parent 4a30d7e3cc
commit 949e49a685
2 changed files with 2 additions and 1 deletions
--- a/.buildkite/test-template-aws.j2
+++ b/.buildkite/test-template-aws.j2
@@ -75,7 +75,7 @@ steps:
            - "'cd {{ (step.working_dir or default_working_dir) | safe }} && {{ step.command or (step.commands | join(' && ')) | safe }}'"
            resources:
              limits:
-                nvidia.com/gpu: 8
+                nvidia.com/gpu: {{ step.num_gpus or 1 }}
            volumeMounts:
            - name: devshm
              mountPath: /dev/shm