From 959f87ae4bcb0b8f9af89e567bea8a71fd7b04be Mon Sep 17 00:00:00 2001
From: Peter Andreas Entschev <peter@entschev.com>
Date: Thu, 13 Aug 2026 22:31:02 -0700
Subject: [PATCH] Expose UCXX progress mode in cudf-polars options

---
 .../cudf_polars/cudf_polars/engine/options.py |  9 ++++++
 .../tests/streaming/test_options.py           | 28 +++++++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/python/cudf_polars/cudf_polars/engine/options.py b/python/cudf_polars/cudf_polars/engine/options.py
index 611129cf2c0..93b163ba42c 100644
--- a/python/cudf_polars/cudf_polars/engine/options.py
+++ b/python/cudf_polars/cudf_polars/engine/options.py
@@ -218,6 +218,12 @@ class StreamingOptions:
         Env: ``RAPIDSMPF_UNBOUNDED_FILE_READ_CACHE``.
         Default: ``"disabled"``.
         Category: rapidsmpf.
+    ucxx_progress_mode
+        UCXX progress mode (``"polling"``, ``"thread-blocking"``, or
+        ``"thread-polling"``).
+        Env: ``RAPIDSMPF_UCXX_PROGRESS_MODE``.
+        Default: ``"thread-blocking"``.
+        Category: rapidsmpf.
     num_py_executors
         Workers for the internal Python ``ThreadPoolExecutor``.
         Env: ``CUDF_POLARS__EXECUTOR__NUM_PY_EXECUTORS``.
@@ -341,6 +347,9 @@ class StreamingOptions:
     unbounded_file_read_cache: str | Unspecified = _opt(
         "rapidsmpf", "RAPIDSMPF_UNBOUNDED_FILE_READ_CACHE"
     )
+    ucxx_progress_mode: (
+        Literal["polling", "thread-blocking", "thread-polling"] | Unspecified
+    ) = _opt("rapidsmpf", "RAPIDSMPF_UCXX_PROGRESS_MODE")
     # ---- Executor ----
     num_py_executors: int | Unspecified = _opt(
         "executor", "CUDF_POLARS__EXECUTOR__NUM_PY_EXECUTORS", int
diff --git a/python/cudf_polars/tests/streaming/test_options.py b/python/cudf_polars/tests/streaming/test_options.py
index d422c92d5a6..5b7659cef63 100644
--- a/python/cudf_polars/tests/streaming/test_options.py
+++ b/python/cudf_polars/tests/streaming/test_options.py
@@ -118,6 +118,7 @@ def test_rapidsmpf_options_serialized() -> None:
         statistics=True,
         pinned_memory=False,
         num_streaming_threads=8,
+        ucxx_progress_mode="thread-polling",
         log="DEBUG",
         pinned_max_pool_size="4GiB",
         unbounded_file_read_cache="host",
@@ -126,6 +127,7 @@ def test_rapidsmpf_options_serialized() -> None:
     assert strings["statistics"] == "True"
     assert strings["pinned_memory"] == "False"
     assert strings["num_streaming_threads"] == "8"
+    assert strings["ucxx_progress_mode"] == "thread-polling"
     assert strings["log"] == "DEBUG"
     assert strings["pinned_max_pool_size"] == "4GiB"
     assert strings["unbounded_file_read_cache"] == "host"
@@ -159,6 +161,32 @@ def test_rapidsmpf_options_env_var_absent(monkeypatch: pytest.MonkeyPatch) -> No
     assert "log" not in StreamingOptions().to_rapidsmpf_options().get_strings()
 
 
+def test_ucxx_progress_mode_picks_up_env_var(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv("RAPIDSMPF_UCXX_PROGRESS_MODE", "polling")
+    strings = StreamingOptions().to_rapidsmpf_options().get_strings()
+    assert strings["ucxx_progress_mode"] == "polling"
+
+
+def test_ucxx_progress_mode_explicit_overrides_env_var(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv("RAPIDSMPF_UCXX_PROGRESS_MODE", "polling")
+    strings = (
+        StreamingOptions(ucxx_progress_mode="thread-polling")
+        .to_rapidsmpf_options()
+        .get_strings()
+    )
+    assert strings["ucxx_progress_mode"] == "thread-polling"
+
+
+def test_ucxx_progress_mode_absent(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.delenv("RAPIDSMPF_UCXX_PROGRESS_MODE", raising=False)
+    strings = StreamingOptions().to_rapidsmpf_options().get_strings()
+    assert "ucxx_progress_mode" not in strings
+
+
 def test_pinned_max_pool_size_env_var(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setenv("RAPIDSMPF_PINNED_MAX_POOL_SIZE", "4GiB")
     strings = StreamingOptions().to_rapidsmpf_options().get_strings()