fix: Add resource requirements (kaito-project#227)

NVIDIA Device Plugin requires `resources` to set up container environment to correctly interface with GPUs --------- Co-authored-by: Fei Guo <[email protected]>
bangqipropel · Feb 1, 2024 · fbdd2ec · fbdd2ec
1 parent 3b7eebe
commit fbdd2ec
Show file tree

Hide file tree

Showing 6 changed files with 21 additions and 1 deletion.
diff --git a/.github/workflows/e2e-preset-configs.json → .github/e2e-preset-configs.json b/.github/workflows/e2e-preset-configs.json → .github/e2e-preset-configs.json
diff --git a/.github/workflows/e2e-preset-test.yml b/.github/workflows/e2e-preset-test.yml
@@ -34,7 +34,7 @@ jobs:
             MATRIX=$(yq e -o=json '.models' presets/models/supported_models.yaml | jq -c)
             
             # Read the additional configurations from e2e-preset-configs.json
-            CONFIGS=$(cat .github/workflows/e2e-preset-configs.json | jq -c '.matrix.image')
+            CONFIGS=$(cat .github/e2e-preset-configs.json | jq -c '.matrix.image')
             
             # Pseudocode for combining matrices
             # COMBINED_MATRIX = []

diff --git a/presets/test/manifests/falcon-40b-instruct/falcon-40b-instruct-statefulset.yaml b/presets/test/manifests/falcon-40b-instruct/falcon-40b-instruct-statefulset.yaml
@@ -20,6 +20,11 @@ spec:
             - /bin/sh
             - -c
             - accelerate launch --num_processes 1 --num_machines 1 --machine_rank 0 --gpu_ids all inference-api.py --pipeline text-generation --torch_dtype bfloat16
+          resources:
+            requests:
+              nvidia.com/gpu: 4
+            limits:
+              nvidia.com/gpu: 4  # Requesting 4 GPUs
           livenessProbe:
             httpGet:
               path: /healthz

diff --git a/presets/test/manifests/falcon-40b/falcon-40b-statefulset.yaml b/presets/test/manifests/falcon-40b/falcon-40b-statefulset.yaml
@@ -20,6 +20,11 @@ spec:
             - /bin/sh
             - -c
             - accelerate launch --num_processes 1 --num_machines 1 --machine_rank 0 --gpu_ids all inference-api.py --pipeline text-generation --torch_dtype bfloat16
+          resources:
+            requests:
+              nvidia.com/gpu: 4
+            limits:
+              nvidia.com/gpu: 4  # Requesting 4 GPUs
           livenessProbe:
             httpGet:
               path: /healthz

diff --git a/presets/test/manifests/falcon-7b-instruct/falcon-7b-instruct-statefulset.yaml b/presets/test/manifests/falcon-7b-instruct/falcon-7b-instruct-statefulset.yaml
@@ -20,6 +20,11 @@ spec:
             - /bin/sh
             - -c
             - accelerate launch --num_processes 1 --num_machines 1 --machine_rank 0 --gpu_ids all inference-api.py --pipeline text-generation --torch_dtype bfloat16
+          resources:
+            requests:
+              nvidia.com/gpu: 2
+            limits:
+              nvidia.com/gpu: 2  # Requesting 2 GPUs
           livenessProbe:
             httpGet:
               path: /healthz

diff --git a/presets/test/manifests/falcon-7b/falcon-7b-statefulset.yaml b/presets/test/manifests/falcon-7b/falcon-7b-statefulset.yaml
@@ -20,6 +20,11 @@ spec:
             - /bin/sh
             - -c
             - accelerate launch --num_processes 1 --num_machines 1 --machine_rank 0 --gpu_ids all inference-api.py --pipeline text-generation --torch_dtype bfloat16
+          resources:
+            requests:
+              nvidia.com/gpu: 2
+            limits:
+              nvidia.com/gpu: 2  # Requesting 2 GPUs
           livenessProbe:
             httpGet:
               path: /healthz