From 4839ded90bddd24856f633f9c704b556624dd12e Mon Sep 17 00:00:00 2001
From: Mihir Patel <mihir.v.patel7@gmail.com>
Date: Thu, 3 Oct 2024 13:26:37 -0400
Subject: [PATCH] Add cluster as an arg to GPU workflow (#32)

* add cluster

* gpu type
---
 .github/actions/pytest-gpu/action.yaml | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/.github/actions/pytest-gpu/action.yaml b/.github/actions/pytest-gpu/action.yaml
index af23b19..8abdc0b 100644
--- a/.github/actions/pytest-gpu/action.yaml
+++ b/.github/actions/pytest-gpu/action.yaml
@@ -28,6 +28,14 @@ inputs:
   mcloud_api_key:
     required: true
     description: "MosaicML API key - Pass in as repo secret"
+  cluster:
+    required: false
+    type: string
+    default: r1z4
+  gpu_type:
+    required: false
+    type: string
+    default: a100_40gb
   python_version:
     required: false
     type: string
@@ -116,6 +124,8 @@ runs:
         fi
 
         python .github/mcli/mcli_pytest.py \
+          --cluster '${{ inputs.cluster }}' \
+          --gpu_type '${{ inputs.gpu_type }}' \
           --image '${{ inputs.container }}' \
           --git_repo '${{ inputs.git_repo }}' \
           --pip_deps '${{ inputs.pip_deps }}' \