From d9c1204a90f68e76dfa7d6162250bf3d23f86959 Mon Sep 17 00:00:00 2001
From: Raunak Bhagat <rabhagat31@gmail.com>
Date: Wed, 4 Dec 2024 14:07:24 -0800
Subject: [PATCH] Submit daft repo instead of using wheel

---
 .github/ci-scripts/templatize_ray_config.py | 27 ++++-----------------
 .github/workflows/run-cluster.yaml          | 23 +++++-------------
 2 files changed, 11 insertions(+), 39 deletions(-)

diff --git a/.github/ci-scripts/templatize_ray_config.py b/.github/ci-scripts/templatize_ray_config.py
index 2a444c7eec..bc3e904a54 100644
--- a/.github/ci-scripts/templatize_ray_config.py
+++ b/.github/ci-scripts/templatize_ray_config.py
@@ -14,7 +14,7 @@
 
 CLUSTER_NAME_PLACEHOLDER = "\\{{CLUSTER_NAME}}"
 DAFT_INSTALL_PLACEHOLDER = "\\{{DAFT_INSTALL}}"
-OTHER_INSTALL_PLACEHOLDER = "\{{OTHER_INSTALLS}}"
+OTHER_INSTALL_PLACEHOLDER = "\\{{OTHER_INSTALLS}}"
 PYTHON_VERSION_PLACEHOLDER = "\\{{PYTHON_VERSION}}"
 CLUSTER_PROFILE__NODE_COUNT = "\\{{CLUSTER_PROFILE/node_count}}"
 CLUSTER_PROFILE__INSTANCE_TYPE = "\\{{CLUSTER_PROFILE/instance_type}}"
@@ -68,28 +68,14 @@ class Metadata(BaseModel, extra="allow"):
 
     parser = ArgumentParser()
     parser.add_argument("--cluster-name", required=True)
-    parser.add_argument("--daft-wheel-url")
     parser.add_argument("--daft-version")
     parser.add_argument("--python-version", required=True)
     parser.add_argument("--cluster-profile", required=True, choices=["debug_xs-x86", "medium-x86"])
-    parser.add_argument("--working-dir", required=True)
-    parser.add_argument("--entrypoint-script", required=True)
+    parser.add_argument("--script", required=True)
     args = parser.parse_args()
 
     content = content.replace(CLUSTER_NAME_PLACEHOLDER, args.cluster_name)
-
-    if args.daft_wheel_url and args.daft_version:
-        raise ValueError(
-            "Cannot specify both the `daft-wheel-name` and the `daft-version`; please choose one or the other"
-        )
-    elif args.daft_wheel_url:
-        daft_install = args.daft_wheel_url
-    elif args.daft_version:
-        daft_install = f"getdaft=={args.daft_version}"
-    else:
-        daft_install = "getdaft"
-    content = content.replace(DAFT_INSTALL_PLACEHOLDER, daft_install)
-
+    content = content.replace(DAFT_INSTALL_PLACEHOLDER, f"getdaft=={args.daft_version}" if args.daft_version else "")
     content = content.replace(PYTHON_VERSION_PLACEHOLDER, args.python_version)
 
     profile = profiles[args.cluster_profile]
@@ -101,11 +87,8 @@ class Metadata(BaseModel, extra="allow"):
         CLUSTER_PROFILE__VOLUME_MOUNT, profile.volume_mount if profile.volume_mount else NOOP_STEP
     )
 
-    working_dir = Path(args.working_dir)
-    assert working_dir.exists() and working_dir.is_dir()
-    entrypoint_script_fullpath: Path = working_dir / args.entrypoint_script
-    assert entrypoint_script_fullpath.exists() and entrypoint_script_fullpath.is_file()
-    with open(entrypoint_script_fullpath) as f:
+    assert Path(args.script).exists()
+    with open(args.script) as f:
         metadata = Metadata(**read_inline_metadata.read(f.read()))
 
     content = content.replace(OTHER_INSTALL_PLACEHOLDER, " ".join(metadata.dependencies))
diff --git a/.github/workflows/run-cluster.yaml b/.github/workflows/run-cluster.yaml
index 903706b8cc..b830d9dbee 100644
--- a/.github/workflows/run-cluster.yaml
+++ b/.github/workflows/run-cluster.yaml
@@ -3,10 +3,6 @@ name: run-cluster
 on:
   workflow_dispatch:
     inputs:
-      daft_wheel_url:
-        description: Daft python-wheel URL
-        type: string
-        required: false
       daft_version:
         description: Daft version (errors if both this and "Daft python-wheel URL" are provided)
         type: string
@@ -24,16 +20,11 @@ on:
         - debug_xs-x86
         required: false
         default: medium-x86
-      working_dir:
-        description: Working directory
-        type: string
-        required: false
-        default: .github/working-dir
-      entrypoint_script:
+      script:
         description: Entry-point python script (must be inside of the working directory)
         type: string
         required: true
-      entrypoint_args:
+      args:
         description: Entry-point arguments
         type: string
         required: false
@@ -75,12 +66,10 @@ jobs:
             --python 3.12 \
             .github/ci-scripts/templatize_ray_config.py \
             --cluster-name "ray-ci-run-${{ github.run_id }}_${{ github.run_attempt }}" \
-            --daft-wheel-url '${{ inputs.daft_wheel_url }}' \
             --daft-version '${{ inputs.daft_version }}' \
             --python-version '${{ inputs.python_version }}' \
             --cluster-profile '${{ inputs.cluster_profile }}' \
-            --working-dir '${{ inputs.working_dir }}' \
-            --entrypoint-script '${{ inputs.entrypoint_script }}'
+            --script '${{ inputs.script }}'
         ) >> .github/assets/ray.yaml
         cat .github/assets/ray.yaml
     - name: Download private ssh key
@@ -99,15 +88,15 @@ jobs:
     - name: Submit job to ray cluster
       run: |
         source .venv/bin/activate
-        if [[ -z '${{ inputs.entrypoint_script }}' ]]; then
+        if [[ -z '${{ inputs.script }}' ]]; then
           echo 'Invalid command submitted; command cannot be empty'
           exit 1
         fi
         ray job submit \
-          --working-dir ${{ inputs.working_dir }} \
+          --working-dir . \
           --address http://localhost:8265 \
           --runtime-env-json '{"env_vars": {"DAFT_ENABLE_RAY_TRACING": "1"}}' \
-          -- python ${{ inputs.entrypoint_script }} ${{ inputs.entrypoint_args }}
+          -- python ${{ inputs.script }} ${{ inputs.args }}
     - name: Download log files from ray cluster
       run: |
         source .venv/bin/activate