From 54d32f38cbd695487abc2e94171c13cac78901e0 Mon Sep 17 00:00:00 2001
From: Raunak Bhagat <rabhagat31@gmail.com>
Date: Wed, 4 Dec 2024 16:44:21 -0800
Subject: [PATCH] Add env-var support

---
 .github/ci-scripts/format_env_vars.py        | 25 ++++++++++++++++++++
 .github/workflows/run-cluster.yaml           | 12 +++++++++-
 .github/working-dir/uv_run_script_example.py |  3 ---
 3 files changed, 36 insertions(+), 4 deletions(-)
 create mode 100644 .github/ci-scripts/format_env_vars.py
 delete mode 100644 .github/working-dir/uv_run_script_example.py

diff --git a/.github/ci-scripts/format_env_vars.py b/.github/ci-scripts/format_env_vars.py
new file mode 100644
index 0000000000..7028e3a79e
--- /dev/null
+++ b/.github/ci-scripts/format_env_vars.py
@@ -0,0 +1,25 @@
+import argparse
+import json
+
+
+def parse_env_var_str(env_var_str: str) -> dict:
+    iter = map(
+        lambda s: s.strip().split("="),
+        filter(lambda s: s, env_var_str.split(",")),
+    )
+    return {k: v for k, v in iter}
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--env-vars", required=True)
+    args = parser.parse_args()
+
+    env_vars = parse_env_var_str(args.env_vars)
+    ray_env_vars = {
+        "env_vars": {
+            "DAFT_ENABLE_RAY_TRACING": "1",
+            **env_vars,
+        },
+    }
+    print(json.dumps(ray_env_vars))
diff --git a/.github/workflows/run-cluster.yaml b/.github/workflows/run-cluster.yaml
index 903706b8cc..149af3ab21 100644
--- a/.github/workflows/run-cluster.yaml
+++ b/.github/workflows/run-cluster.yaml
@@ -38,6 +38,11 @@ on:
         type: string
         required: false
         default: ""
+      env_vars:
+        description: Environment variables
+        type: string
+        required: false
+        default: ""
 
 jobs:
   run-command:
@@ -83,6 +88,11 @@ jobs:
             --entrypoint-script '${{ inputs.entrypoint_script }}'
         ) >> .github/assets/ray.yaml
         cat .github/assets/ray.yaml
+    - name: Setup ray env vars
+      run: |
+        ray_env_var=$(python .github/ci-scripts/format_env_vars.py --env-vars '${{ inputs.env_vars }}')
+        echo $ray_env_var
+        echo "ray_env_var=$ray_env_var" >> $GITHUB_ENV
     - name: Download private ssh key
       run: |
         KEY=$(aws secretsmanager get-secret-value --secret-id ci-github-actions-ray-cluster-key-3 --query SecretString --output text)
@@ -106,7 +116,7 @@ jobs:
         ray job submit \
           --working-dir ${{ inputs.working_dir }} \
           --address http://localhost:8265 \
-          --runtime-env-json '{"env_vars": {"DAFT_ENABLE_RAY_TRACING": "1"}}' \
+          --runtime-env-json "$ray_env_var" \
           -- python ${{ inputs.entrypoint_script }} ${{ inputs.entrypoint_args }}
     - name: Download log files from ray cluster
       run: |
diff --git a/.github/working-dir/uv_run_script_example.py b/.github/working-dir/uv_run_script_example.py
deleted file mode 100644
index 0d7611736e..0000000000
--- a/.github/working-dir/uv_run_script_example.py
+++ /dev/null
@@ -1,3 +0,0 @@
-import daft
-
-print(daft)