Add ability to specify custom daft python-wheel to run

.github/assets/.template.yaml

@@ -48,4 +48,4 @@ setup_commands:
 - uv v
 - echo "source $HOME/.venv/bin/activate" >> $HOME/.bashrc
 - source .venv/bin/activate
-- uv pip install pip ray[default] py-spy getdaft\{{DAFT_VERSION}}
+- uv pip install pip ray[default] py-spy \{{DAFT_INSTALL}}

.github/ci-scripts/templatize_ray_config.py

@@ -4,7 +4,7 @@
 from typing import Optional
 
 CLUSTER_NAME_PLACEHOLDER = "\\{{CLUSTER_NAME}}"
-DAFT_VERSION_PLACEHOLDER = "\\{{DAFT_VERSION}}"
+DAFT_INSTALL_PLACEHOLDER = "\\{{DAFT_INSTALL}}"
 PYTHON_VERSION_PLACEHOLDER = "\\{{PYTHON_VERSION}}"
 CLUSTER_PROFILE__NODE_COUNT = "\\{{CLUSTER_PROFILE/node_count}}"
 CLUSTER_PROFILE__INSTANCE_TYPE = "\\{{CLUSTER_PROFILE/instance_type}}"
@@ -51,6 +51,7 @@ class Profile:
 
     parser = ArgumentParser()
     parser.add_argument("--cluster-name")
+    parser.add_argument("--daft-wheel-url")
     parser.add_argument("--daft-version")
     parser.add_argument("--python-version")
     parser.add_argument("--cluster-profile")
@@ -59,10 +60,17 @@ class Profile:
     if args.cluster_name:
         content = content.replace(CLUSTER_NAME_PLACEHOLDER, args.cluster_name)
 
-    if args.daft_version:
-        content = content.replace(DAFT_VERSION_PLACEHOLDER, f"=={args.daft_version}")
+    if args.daft_wheel_url and args.daft_version:
+        raise Exception(
+            "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:
-        content = content.replace(DAFT_VERSION_PLACEHOLDER, "")
+        daft_install = "getdaft"
+    content = content.replace(DAFT_INSTALL_PLACEHOLDER, daft_install)
 
     if args.python_version:
         content = content.replace(PYTHON_VERSION_PLACEHOLDER, args.python_version)

.github/workflows/run-cluster.yaml

@@ -3,9 +3,13 @@ name: run-cluster
 on:
   workflow_dispatch:
     inputs:
+      daft_wheel_url:
+        type: string
+        description: The url pointing to a daft python-wheel for which to install and run
+        required: false
       daft_version:
         type: string
-        description: The version of daft to install
+        description: The version of daft to install (errors if both this and `daft_wheel_url` are provided)
         required: false
       python_version:
         type: string
@@ -62,6 +66,7 @@ jobs:
         (cat .github/assets/.template.yaml \
           | python .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 }}'

pyproject.toml

@@ -90,6 +90,9 @@ typeCheckingMode = "off"
 venv = ".venv"
 venvPath = "."
 
+[[tool.pyright.executionEnvironments]]
+root = ".github/ci-scripts"
+
 [tool.pytest.ini_options]
 addopts = "-m 'not (integration or benchmark or hypothesis)'"
 minversion = "6.0"