ci: fix TPC-H benchmark workflows (#5123)

## Changes Made

Fixes the UV setup in the benchmark workflows so that they work with our
new dependency setup.

I was having issues with running a Daft script as a Ray job and it turns
out that we broke the behavior of automatically connecting to the Ray
cluster when running in a Ray job in #4567.

The `RAY_ADDRESS` env var is set to `<local ip>:6379` in a job script,
but in that PR we prepended `ray://` to it, which causes `ray.init` to
try and fail to create a Ray client instead of connecting to the
cluster. I've removed that logic, as well as the logic in Daft to read
`RAY_ADDRESS` because `ray.init` does that automatically already.

## Related Issues

<!-- Link to related GitHub issues, e.g., "Closes #123" -->

## Checklist

- [x] Documented in API Docs (if applicable)
- [x] Documented in User Guide (if applicable)
- [x] If adding a new documentation page, doc is added to
`docs/mkdocs.yml` navigation
- [x] Documentation builds and is formatted properly (tag @/ccmao1130
for docs review)

Author: kevinzwang

.github/ci-scripts/distributed_tpch.py

@@ -6,7 +6,6 @@
 import os
 import sys
 import time
-import urllib.request
 
 from ray.job_submission import JobDetails, JobStatus, JobSubmissionClient
 
@@ -39,13 +38,10 @@ def run_benchmark():
 
     client = JobSubmissionClient(address="http://localhost:8265")
 
-    # Create the directory if it doesn't exist
-    os.makedirs("traces", exist_ok=True)
-
     for q in range(1, 23):
         print(f"Running TPC-H Q{q}... ", end="", flush=True)
 
-        start = time.perf_counter()
+        start: float = time.perf_counter()
 
         submission_id = client.submit_job(
             entrypoint=f"DAFT_RUNNER=ray DAFT_PROGRESS_BAR=0 python answers_sql.py {parquet_path} {q}",
@@ -56,10 +52,6 @@ def run_benchmark():
 
         end = time.perf_counter()
 
-        # Download the trace file
-        trace_url = f"http://localhost:8265/api/v0/tasks/timeline?download=1&job_id={job_details.job_id}"
-        urllib.request.urlretrieve(trace_url, f"traces/q{q}-trace.json")
-
         if job_details.status != JobStatus.SUCCEEDED:
             print(f"\nRay job did not succeed, received job status: {job_details.status}\nJob details: {job_details}")
             sys.exit(1)

.github/workflows/benchmark-distributed-tpch.yml

@@ -75,9 +75,13 @@ jobs:
       uses: astral-sh/setup-uv@v6
       with:
         python-version: ${{ env.PYTHON_VERSION }}
-        version: "0.6.17"
-    - name: Install Daft and dev dependencies
+        enable-cache: true
+        cache-dependency-glob: "**/pyproject.toml"
+    - name: Setup Virtual Env
       run: |
+        uv venv --seed .venv
+        echo "$GITHUB_WORKSPACE/.venv/bin" >> $GITHUB_PATH
+        source .venv/bin/activate
         rm -rf daft
         uv pip install daft --pre --extra-index-url ${{ env.DAFT_INDEX_URL }}
         uv pip install gspread ray[default] boto3
@@ -98,18 +102,17 @@ jobs:
         chmod 600 ~/.ssh/ci-github-actions-ray-cluster-key.pem
     - name: Spin up and connect to ray cluster
       run: |
+        source .venv/bin/activate
         ray up ray.yaml -y
         ray dashboard ray.yaml &
 
     - name: Run benchmark and upload results to Google Sheets
-      run: PYTHONPATH=. python .github/ci-scripts/distributed_tpch.py
-
-    - name: Upload traces
-      uses: actions/upload-artifact@v4
-      with:
-        name: traces
-        path: traces
+      run: |
+        source .venv/bin/activate
+        PYTHONPATH=. python .github/ci-scripts/distributed_tpch.py
 
     - name: Spin down ray cluster
       if: always()
-      run: ray down ray.yaml -y
+      run: |
+        source .venv/bin/activate
+        ray down ray.yaml -y

.github/workflows/benchmark-local-tpch.yml

@@ -58,17 +58,23 @@ jobs:
       uses: astral-sh/setup-uv@v6
       with:
         python-version: ${{ env.PYTHON_VERSION }}
-        version: "0.6.17"
-    - name: Install Daft and dev dependencies
+        enable-cache: true
+        cache-dependency-glob: "**/pyproject.toml"
+    - name: Setup Virtual Env
       run: |
+        uv venv --seed .venv
+        echo "$GITHUB_WORKSPACE/.venv/bin" >> $GITHUB_PATH
+        source .venv/bin/activate
         rm -rf daft
         uv pip install daft --pre --extra-index-url ${{ env.DAFT_INDEX_URL }}
-        uv pip install gspread
+        uv pip install gspread ray[default] boto3
     - name: Write service account secret file
       run: |
         mkdir -p ~/.config/gspread
         cat << EOF > ~/.config/gspread/service_account.json
         ${{ secrets.GOOGLE_SHEETS_SERVICE_ACCOUNT }}
         EOF
     - name: Run benchmark and upload results to Google Sheets
-      run: PYTHONPATH=. DAFT_RUNNER=native python .github/ci-scripts/local_tpch.py
+      run: |
+        source .venv/bin/activate
+        PYTHONPATH=. DAFT_RUNNER=native python .github/ci-scripts/local_tpch.py

benchmarking/tpch/answers_sql.py

@@ -80,6 +80,10 @@ def get_answer(q: int, get_df) -> daft.DataFrame:
 
 
 def main(parquet_path, q):
+    if q in (11, 22):
+        # TODO: remove this once we support cross joins in Flotilla
+        daft.set_execution_config(use_legacy_ray_runner=True)
+
     s3_config_from_env = S3Config.from_env()
     io_config = IOConfig(s3=s3_config_from_env)
 

daft/runners/ray_runner.py

@@ -7,7 +7,6 @@
 import threading
 import time
 import uuid
-import warnings
 from collections.abc import Generator, Iterable, Iterator
 from datetime import datetime
 from queue import Full, Queue
@@ -268,7 +267,7 @@ def _to_pandas_ref(df: pd.DataFrame | ray.ObjectRef) -> ray.ObjectRef:
     elif isinstance(df, ray.ObjectRef):
         return df
     else:
-        raise ValueError("Expected a Ray object ref or a Pandas DataFrame, " f"got {type(df)}")
+        raise ValueError(f"Expected a Ray object ref or a Pandas DataFrame, got {type(df)}")
 
 
 class RayPartitionSet(PartitionSet[ray.ObjectRef]):
@@ -887,8 +886,7 @@ def _run_plan(
 
         start = datetime.now()
         profile_filename = (
-            f"profile_RayRunner.run()_"
-            f"{datetime.replace(datetime.now(), second=0, microsecond=0).isoformat()[:-3]}.json"
+            f"profile_RayRunner.run()_{datetime.replace(datetime.now(), second=0, microsecond=0).isoformat()[:-3]}.json"
         )
 
         with profiler(profile_filename), ray_tracing.ray_tracer(result_uuid, daft_execution_config) as runner_tracer:
@@ -1234,16 +1232,11 @@ def __init__(
                     address,
                 )
         else:
-            if address is not None:
-                if not address.endswith("10001"):
-                    warnings.warn(
-                        f"The address to a Ray client server is typically at port :10001, but instead we found: {address}"
-                    )
-                if not address.startswith("ray://"):
-                    warnings.warn(
-                        f"Expected Ray address to start with 'ray://' protocol but found: {address}. Automatically prefixing your address with the protocol to make a Ray connection: ray://{address}"
-                    )
-                    address = "ray://" + address
+            if address is not None and address.startswith("ray://"):
+                logger.warning(
+                    "Specifying a Ray address with the 'ray://' prefix uses the Ray Client, which may impact performance. If this is running in a Ray job, you may not need to specify the address at all."
+                )
+
             ray.init(address=address)
 
         # Check if Ray is running in "client mode" (connected to a Ray cluster via a Ray client)

src/daft-context/src/lib.rs

@@ -312,19 +312,17 @@ fn parse_usize_env_var(var_name: &str) -> Option<usize> {
 #[cfg(feature = "python")]
 fn get_ray_runner_config_from_env() -> RunnerConfig {
     const DAFT_RAY_ADDRESS: &str = "DAFT_RAY_ADDRESS";
-    const RAY_ADDRESS: &str = "RAY_ADDRESS";
     const DAFT_DEVELOPER_RAY_MAX_TASK_BACKLOG: &str = "DAFT_DEVELOPER_RAY_MAX_TASK_BACKLOG";
     const DAFT_RAY_FORCE_CLIENT_MODE: &str = "DAFT_RAY_FORCE_CLIENT_MODE";
 
     let address = if let Ok(address) = std::env::var(DAFT_RAY_ADDRESS) {
         log::warn!(
-            "Detected usage of the ${} environment variable. This will be deprecated, please use ${} instead.",
+            "Detected usage of the ${} environment variable. This will be deprecated, please use $RAY_ADDRESS instead.",
             DAFT_RAY_ADDRESS,
-            RAY_ADDRESS
         );
         Some(address)
     } else {
-        std::env::var(RAY_ADDRESS).ok()
+        None
     };
     let max_task_backlog = parse_usize_env_var(DAFT_DEVELOPER_RAY_MAX_TASK_BACKLOG);
     let force_client_mode = parse_bool_env_var(DAFT_RAY_FORCE_CLIENT_MODE);