[STF] Add python bindings

.github/CODEOWNERS

@@ -23,7 +23,7 @@ nvbench_helper/ @nvidia/cccl-infra-codeowners
 **/cmake/ @nvidia/cccl-cmake-codeowners
 # make an exception for the cudax test cmake file
 cudax/test/CMakeLists.txt @nvidia/cccl-cudax-codeowners
-# make an exception for STF test cmake files, following the PR #5406 cudax test precedent
+# make an exception for STF test cmake files
 cudax/test/stf/CMakeLists.txt @nvidia/cccl-cudax-codeowners
 cudax/test/stf/static_error_checks/CMakeLists.txt @nvidia/cccl-cudax-codeowners
 

AGENTS.md

@@ -207,6 +207,7 @@ Supported versions: `3.10`, `3.11`, `3.12`, `3.13`
 
 * **cuda.compute** — Device-level algorithms, iterators, custom GPU types
 * **cuda.coop._experimental** — Block/warp-level primitives for Numba CUDA
+* **cuda.stf._experimental** — Sequential Task Flow (CUDASTF) Python bindings in the `cuda-stf` package (Linux only)
 * **cuda.cccl.headers** — Programmatic access to headers
 
 ### Installation
@@ -262,14 +263,16 @@ include_paths = headers.get_include_paths()
 ./ci/test_cuda_coop_python.sh -py-version 3.10
 ./ci/test_cuda_cccl_headers_python.sh -py-version 3.10
 ./ci/test_cuda_cccl_examples_python.sh -py-version 3.10
+./ci/test_cuda_stf_python.sh -py-version 3.10  # Linux only
 ```
 
 Test organization:
 
 * `tests/compute` — Algorithms and iterators
 * `tests/coop` — Cooperative primitives
 * `tests/headers` — Header integration
-* `test_examples.py` — Runs compute/coop examples
+* `python/cuda_stf/tests/stf` — Sequential Task Flow (separate `cuda-stf` package, Linux only)
+* `test_examples.py` — Runs compute/coop examples (STF examples live in `python/cuda_stf/tests/test_examples.py`)
 
 ---
 

ci/build_cuda_stf_combined_python.sh

@@ -0,0 +1,37 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Combined producer for the cuda.stf._experimental test lane.
+#
+# cuda-stf tests must run against the *current branch's* cuda-cccl as well as
+# the current branch's cuda-stf. The workflow model lets a consumer job depend
+# on only a single producer, so this one producer builds and uploads BOTH
+# wheels under distinct artifact names (via get_wheel_artifact_name.sh):
+#
+#   * cuda-cccl  -> wheel-cccl-stf-<os>-<arch>-py<ver>  (CCCL_WHEEL_KIND=cccl-stf)
+#   * cuda-stf   -> wheel-stf-<os>-<arch>-py<ver>       (CCCL_WHEEL_KIND=stf)
+#
+# The cuda-cccl wheel uses the dedicated 'cccl-stf' kind so it does NOT collide
+# with the 'wheel-cccl-...' artifact produced by the regular build_py_wheel job
+# (which also runs in project 'python' with the same py/os/arch).
+#
+# ci/test_cuda_stf_python.sh then downloads and installs both local wheels in a
+# single resolver invocation. Build cuda-cccl first: the cuda-stf build leaves
+# any co-located cuda_cccl wheel in wheelhouse/ untouched.
+
+ci_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+usage="Usage: $0 -py-version <python_version> [additional options...]"
+
+# shellcheck source=ci/util/python/common_arg_parser.sh
+source "$ci_dir/util/python/common_arg_parser.sh"
+parse_python_args "$@"
+require_py_version "$usage" || exit 1
+
+echo "::group::⚒️ Building current-branch cuda-cccl wheel"
+CCCL_WHEEL_KIND=cccl-stf "$ci_dir/build_cuda_cccl_python.sh" "$@"
+echo "::endgroup::"
+
+echo "::group::⚒️ Building current-branch cuda-stf wheel"
+"$ci_dir/build_cuda_stf_python.sh" "$@"
+echo "::endgroup::"

ci/build_cuda_stf_python.sh

@@ -0,0 +1,184 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ci_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+usage="Usage: $0 -py-version <python_version> [additional options...]"
+
+# shellcheck source=ci/util/python/common_arg_parser.sh
+source "$ci_dir/util/python/common_arg_parser.sh"
+parse_python_args "$@"
+
+# Check if py_version was provided (this script requires it)
+require_py_version "$usage" || exit 1
+
+echo "Docker socket: " "$(ls /var/run/docker.sock)"
+
+if [[ -n "${GITHUB_ACTIONS:-}" ]]; then
+  # Prepare mount points etc for getting artifacts in/out of the container.
+  # shellcheck source=ci/util/artifacts/common.sh
+  source "$ci_dir/util/artifacts/common.sh"
+  # Note that these mounts use the runner (not the devcontainer) filesystem for
+  # source directories because of docker-out-of-docker quirks.
+  # The workflow-job GH actions make sure that they exist before running any
+  # scripts.
+  action_mounts=(
+    --mount "type=bind,source=${ARTIFACT_ARCHIVES},target=${ARTIFACT_ARCHIVES}"
+    --mount "type=bind,source=${ARTIFACT_UPLOAD_STAGE},target=${ARTIFACT_UPLOAD_STAGE}"
+  )
+else
+  # If not running in GitHub Actions, we don't need to set up artifact mounts.
+  action_mounts=()
+fi
+
+# cuda_stf must be built in a container that can produce manylinux wheels, and
+# has the CUDA toolkit installed. We use the rapidsai/ci-wheel image for this.
+# We build separate wheels using separate containers for each CUDA version,
+# then merge them into a single wheel. CUDASTF is Linux-only.
+
+readonly cuda12_version=12.9.1
+readonly cuda13_version=13.1.1
+readonly devcontainer_version=26.04
+readonly devcontainer_distro=rockylinux8
+# Use a baseline Python tag for the rapidsai ci-wheel image. The requested
+# py_version is installed inside the container by setup_python_env (uv).
+readonly devcontainer_python_version=3.10
+
+if [[ "$(uname -m)" == "aarch64" ]]; then
+  cuda12_image="rapidsai/ci-wheel:${devcontainer_version}-cuda${cuda12_version}-${devcontainer_distro}-py${devcontainer_python_version}-arm64"
+  cuda13_image="rapidsai/ci-wheel:${devcontainer_version}-cuda${cuda13_version}-${devcontainer_distro}-py${devcontainer_python_version}-arm64"
+else
+  cuda12_image="rapidsai/ci-wheel:${devcontainer_version}-cuda${cuda12_version}-${devcontainer_distro}-py${devcontainer_python_version}"
+  cuda13_image="rapidsai/ci-wheel:${devcontainer_version}-cuda${cuda13_version}-${devcontainer_distro}-py${devcontainer_python_version}"
+fi
+# shellcheck disable=SC2034
+readonly cuda12_image
+# shellcheck disable=SC2034
+readonly cuda13_image
+
+mkdir -p wheelhouse
+
+# Clear stale STF wheels from a previous run so the per-CTK wheel selection
+# below is unambiguous. Leave any co-located wheels (e.g. a cuda_cccl wheel
+# staged by a combined producer job) untouched.
+rm -f wheelhouse/cuda_stf-*.whl
+
+# Shared caches across the cu12 + cu13 wheel builds. Both jobs compile an
+# identical LLVM/clang tree (LLVM has no CUDA dep), so a shared ccache cuts
+# the second build's LLVM phase substantially; a shared CPM source cache skips
+# the second LLVM git clone entirely.
+mkdir -p ./.ccache ./.cpm-cache
+host_ccache_dir="${HOST_WORKSPACE:?}/.ccache"
+host_cpm_cache_dir="${HOST_WORKSPACE:?}/.cpm-cache"
+
+for ctk in 12 13; do
+  image="cuda${ctk}_image"
+  image="${!image}"
+  echo "::group::⚒️ Building CUDA $ctk cuda-stf wheel on $image"
+  (
+    set -x
+    docker pull "$image"
+    docker run --rm -i \
+        --workdir /workspace/python/cuda_stf \
+        --mount "type=bind,source=${HOST_WORKSPACE:?},target=/workspace/" \
+        --mount "type=bind,source=${host_ccache_dir},target=/root/.ccache" \
+        --mount "type=bind,source=${host_cpm_cache_dir},target=/root/.cpm-cache" \
+        "${action_mounts[@]}" \
+        --env "py_version=${py_version}" \
+        --env "GITHUB_ACTIONS=${GITHUB_ACTIONS:-}" \
+        --env "GITHUB_RUN_ID=${GITHUB_RUN_ID:-}" \
+        --env "JOB_ID=${JOB_ID:-}" \
+        --env "CCACHE_DIR=/root/.ccache" \
+        --env "CPM_SOURCE_CACHE=/root/.cpm-cache" \
+        "$image" \
+        /workspace/ci/build_cuda_stf_wheel.sh
+    # Prevent GHA runners from exhausting available storage with leftover images:
+    if [[ -n "${GITHUB_ACTIONS:-}" ]]; then
+      docker rmi -f "$image"
+    fi
+  )
+  echo "::endgroup::"
+done
+
+echo "Merging CUDA wheels..."
+
+# Set up a Python environment for the merge/repair steps.
+source "$ci_dir/pyenv_helper.sh"
+setup_python_env "${py_version}"
+
+# Needed for unpacking and repacking wheels.
+python -m pip install wheel
+
+# Find the built wheels, requiring exactly one match per CUDA version so a
+# stale or duplicate wheel cannot be silently merged.
+require_single_wheel() {
+  local pattern="$1" desc="$2"
+  local matches=()
+  while IFS= read -r match; do
+    matches+=("$match")
+  done < <(find wheelhouse -maxdepth 1 -name "$pattern" | sort)
+  if [[ ${#matches[@]} -eq 0 ]]; then
+    echo "Error: no $desc cuda-stf wheel found in wheelhouse/ (pattern: $pattern)" >&2
+    ls -la wheelhouse/ >&2
+    exit 1
+  fi
+  if [[ ${#matches[@]} -gt 1 ]]; then
+    echo "Error: expected exactly one $desc cuda-stf wheel, found ${#matches[@]}:" >&2
+    printf '  %s\n' "${matches[@]}" >&2
+    exit 1
+  fi
+  printf '%s\n' "${matches[0]}"
+}
+
+cu12_wheel="$(require_single_wheel 'cuda_stf-*cu12*.whl' 'CUDA 12')"
+cu13_wheel="$(require_single_wheel 'cuda_stf-*cu13*.whl' 'CUDA 13')"
+
+echo "Found CUDA 12 wheel: $cu12_wheel"
+echo "Found CUDA 13 wheel: $cu13_wheel"
+
+# Merge the wheels
+python python/cuda_stf/merge_cuda_wheels.py "$cu12_wheel" "$cu13_wheel" --output-dir wheelhouse_merged
+
+# Install auditwheel and repair the merged wheel
+python -m pip install patchelf auditwheel
+for wheel in wheelhouse_merged/cuda_stf-*.whl; do
+    echo "Repairing merged wheel: $wheel"
+    python -m auditwheel repair \
+        --exclude 'libnvrtc.so.12' \
+        --exclude 'libnvrtc.so.13' \
+        --exclude 'libnvJitLink.so.12' \
+        --exclude 'libnvJitLink.so.13' \
+        --exclude 'libcudart.so.12' \
+        --exclude 'libcudart.so.13' \
+        --exclude 'libcuda.so.1' \
+        "$wheel" \
+        --wheel-dir wheelhouse_final
+done
+
+# Drop only the per-CTK STF inputs we just merged; keep any unrelated wheels
+# (e.g. a co-located cuda_cccl wheel) intact.
+rm -f "$cu12_wheel" "$cu13_wheel"
+mkdir -p wheelhouse
+
+# Move only the final repaired merged wheel
+if ls wheelhouse_final/cuda_stf-*.whl 1> /dev/null 2>&1; then
+    mv wheelhouse_final/cuda_stf-*.whl wheelhouse/
+    echo "Final merged wheel moved to wheelhouse"
+else
+    echo "No final repaired wheel found, moving unrepaired merged wheel"
+    mv wheelhouse_merged/cuda_stf-*.whl wheelhouse/
+fi
+
+# Clean up temporary directories
+rm -rf wheelhouse_merged wheelhouse_final
+
+echo "Final wheels in wheelhouse:"
+ls -la wheelhouse/
+
+if [[ -n "${GITHUB_ACTIONS:-}" ]]; then
+  # Upload under a distinct artifact name so it does not clobber the cuda-cccl
+  # wheel (both build jobs run in project 'python').
+  wheel_artifact_name="$(CCCL_WHEEL_KIND=stf ci/util/workflow/get_wheel_artifact_name.sh)"
+  # Upload only the final STF wheel, not any co-located wheels in wheelhouse/.
+  ci/util/artifacts/upload.sh "$wheel_artifact_name" 'wheelhouse/cuda_stf-.*\.whl'
+fi

ci/build_cuda_stf_wheel.sh

@@ -0,0 +1,82 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Target script for `docker run` command in build_cuda_stf_python.sh
+# The /workspace pathnames are hard-wired here.
+
+# Install GCC 13 toolset (needed for the build) and ccache (shared between
+# cu12 and cu13 builds via /root/.ccache bind-mount from the host).
+/workspace/ci/util/retry.sh 5 30 dnf -y install \
+  gcc-toolset-13-gcc gcc-toolset-13-gcc-c++ ccache
+
+# When the caller bind-mounts a ccache dir, wire it through to CMake. This
+# transparently caches every compile, so the second wheel build (cu13 after
+# cu12, or vice versa) reuses the entire LLVM/clang object tree.
+if [[ -n "${CCACHE_DIR:-}" ]]; then
+  export CMAKE_C_COMPILER_LAUNCHER=ccache
+  export CMAKE_CXX_COMPILER_LAUNCHER=ccache
+  export CMAKE_CUDA_COMPILER_LAUNCHER=ccache
+  echo "ccache enabled: CCACHE_DIR=${CCACHE_DIR}"
+  ccache --version 2>&1 | head -1 || true
+  ccache --show-stats 2>&1 | head -5 || true
+fi
+echo -e "#!/usr/bin/env bash\nsource /opt/rh/gcc-toolset-13/enable" >/etc/profile.d/enable_devtools.sh
+# shellcheck disable=SC1091
+source /etc/profile.d/enable_devtools.sh
+
+# Check what's available
+command -v gcc
+gcc --version
+command -v nvcc
+nvcc --version
+
+# Set up Python environment
+# shellcheck source=ci/pyenv_helper.sh
+source /workspace/ci/pyenv_helper.sh
+# shellcheck disable=SC2154
+setup_python_env "${py_version}"
+command -v python
+python --version
+echo "Done setting up python env"
+
+# Figure out the version to use for the package, we need repo history
+if "$(git rev-parse --is-shallow-repository)"; then
+  git fetch --unshallow
+fi
+# Match the cuda-cccl version prefix so cuda-stf and cuda-cccl stay in lockstep.
+export PACKAGE_VERSION_PREFIX="0.1."
+package_version=$(/workspace/ci/generate_version.sh)
+echo "Using package version ${package_version}"
+# Override the version used by setuptools_scm to the custom version
+export SETUPTOOLS_SCM_PRETEND_VERSION_FOR_CUDA_STF="${package_version}"
+
+cd /workspace/python/cuda_stf
+
+# Determine CUDA version from nvcc
+cuda_version=$(nvcc --version | grep -oP 'release \K[0-9]+\.[0-9]+' | cut -d. -f1)
+echo "Detected CUDA version: ${cuda_version}"
+
+# Configure compilers:
+CXX="$(command -v g++)"
+export CXX
+CUDACXX="$(command -v nvcc)"
+export CUDACXX
+CUDAHOSTCXX="$(command -v g++)"
+export CUDAHOSTCXX
+
+# Build the wheel
+python -m pip wheel --no-deps --verbose --wheel-dir dist .
+
+# Rename wheel to include CUDA version suffix
+for wheel in dist/cuda_stf-*.whl; do
+    if [[ -f "$wheel" ]]; then
+        base_name=$(basename "$wheel" .whl)
+        new_name="${base_name}.cu${cuda_version}.whl"
+        mv "$wheel" "dist/${new_name}"
+        echo "Renamed wheel to: ${new_name}"
+    fi
+done
+
+# Move wheel to output directory
+mkdir -p /workspace/wheelhouse
+mv dist/cuda_stf-*.cu*.whl /workspace/wheelhouse/