AZP: Build UCXX conda + wheel packages (Phase #2)

buildlib/azure-pipelines-pr.yml

@@ -31,5 +31,13 @@ pr:
     - buildlib/tools/perf_results.py
     - buildlib/tools/perf-common.yml
 
+resources:
+  repositories:
+    - repository: ucxx
+      type: github
+      name: rapidsai/ucxx
+      endpoint: Mellanox-lab
+      ref: refs/heads/main
+
 extends:
   template: pr/main.yml

buildlib/dockers/rapidsai-ci-conda.Dockerfile

@@ -0,0 +1,7 @@
+# Azure Pipelines wrapper around rapidsai/ci-conda.
+# Opens /opt/conda so the Azure-injected step user can use conda/python.
+
+ARG BASE_IMAGE=rapidsai/ci-conda:26.06-latest
+FROM ${BASE_IMAGE}
+
+RUN chmod -R o+rwX /opt/conda

buildlib/dockers/rapidsai-ci-wheel.Dockerfile

@@ -0,0 +1,7 @@
+# Azure Pipelines wrapper around rapidsai/ci-wheel.
+# Opens /pyenv so the Azure-injected step user can write shims.
+
+ARG BASE_IMAGE=rapidsai/ci-wheel:26.06-latest
+FROM ${BASE_IMAGE}
+
+RUN chmod -R o+rwX /pyenv

buildlib/pr/main.yml

@@ -249,23 +249,40 @@ resources:
     - container: centos10stream
       image: rdmz-harbor.rdmz.labs.mlnx/hpcx/x86_64/centos10stream/builder:inbox
       options: $(DOCKER_OPT_ARGS) $(DOCKER_OPT_VOLUMES)
+    - container: ucxx_rapidsai_ci_conda
+      # Thin wrapper of rapidsai/ci-conda; see buildlib/dockers/rapidsai-ci-conda.Dockerfile.
+      image: rdmz-harbor.rdmz.labs.mlnx/ucx/rapidsai-ci-conda:26.06-azp-1
+      options: $(DOCKER_OPT_ARGS) $(DOCKER_OPT_VOLUMES)
+    - container: ucxx_rapidsai_ci_conda_gpu
+      # Same image as above, +GPU. Bind-mount host MPS dirs so the container's
+      # CUDA client can reach the host MPS daemon (host GPUs are in
+      # Exclusive_Process compute mode and need MPS for shared access).
+      image: rdmz-harbor.rdmz.labs.mlnx/ucx/rapidsai-ci-conda:26.06-azp-1
+      options: $(DOCKER_OPT_ARGS) $(DOCKER_OPT_VOLUMES) $(DOCKER_OPT_GPU) -v /tmp/nvidia-mps:/tmp/nvidia-mps
+    - container: ucxx_rapidsai_ci_wheel
+      # Thin wrapper of rapidsai/ci-wheel; see buildlib/dockers/rapidsai-ci-wheel.Dockerfile.
+      image: rdmz-harbor.rdmz.labs.mlnx/ucx/rapidsai-ci-wheel:26.06-azp-1
+      options: $(DOCKER_OPT_ARGS) $(DOCKER_OPT_VOLUMES)
 
 stages:
   - stage: Codestyle
     jobs:
     - template: codestyle.yml
 
   - stage: Basic_compile
+    condition: false  # temp: skip non-UCXX stage; restore before merge
     dependsOn: []
     jobs:
     - template: basic_compile.yml
 
   - stage: Static_check
+    condition: false  # temp: skip non-UCXX stage; restore before merge
     dependsOn: [Basic_compile]
     jobs:
     - template: static_checks.yml
 
   - stage: Build
+    condition: false  # temp: skip non-UCXX stage; restore before merge
     dependsOn: [Basic_compile]
     jobs:
       - template: build_job.yml
@@ -283,6 +300,7 @@ stages:
           displayName: Build on aarch64
 
   - stage: ucx_perftest_mad_rte
+    condition: false  # temp: skip non-UCXX stage; restore before merge
     dependsOn: [Basic_compile]
     displayName: ucx_perftest over MAD RTE
     lockBehavior: sequential
@@ -292,6 +310,7 @@ stages:
     - template: mad_tests.yml
 
   - stage: WireCompat
+    condition: false  # temp: skip non-UCXX stage; restore before merge
     dependsOn: [Basic_compile]
     jobs:
     - template: wire_compat.yml
@@ -323,14 +342,126 @@ stages:
     #    demands: ucx_iodemo -equals yes
 
   - stage: Coverity
+    condition: false  # temp: skip non-UCXX stage; restore before merge
     dependsOn: [Static_check]
     jobs:
       - template: coverity.yml
         parameters:
           demands: ucx_docker -equals yes
           container: coverity_rh7
 
+  - template: ucxx_tests_stage.yml
+    parameters:
+      dependsOn: []  # temp: skip Static_check to test stage faster on this PR; restore [Static_check] before merge
+
+  - stage: UCXX_build
+    dependsOn: []  # temp: skip Static_check while iterating Phase 2; restore [Static_check] before merge
+    jobs:
+      # 2.1: conda-cpp-build. Mirrors upstream conda-cpp-build matrix (4 legs).
+      - template: ucxx_conda_cpp_build.yml
+        parameters:
+          container: ucxx_rapidsai_ci_conda
+          slices:
+            - name: x86_64_cuda12_py311
+              demands: ucx_docker -equals yes
+              rapids_cuda_version: '12.9.1'
+              rapids_py_version: '3.11'
+            - name: x86_64_cuda13_py311
+              demands: ucx_docker -equals yes
+              rapids_cuda_version: '13.2.0'
+              rapids_py_version: '3.11'
+            - name: aarch64_cuda12_py311
+              demands: ucx_arm64
+              rapids_cuda_version: '12.9.1'
+              rapids_py_version: '3.11'
+            - name: aarch64_cuda13_py311
+              demands: ucx_arm64
+              rapids_cuda_version: '13.2.0'
+              rapids_py_version: '3.11'
+      # 2.2: conda-python-build. Depends on 2.1's artifacts (per-slice).
+      - template: ucxx_conda_python_build.yml
+        parameters:
+          container: ucxx_rapidsai_ci_conda
+          slices:
+            - name: x86_64_cuda12_py311
+              demands: ucx_docker -equals yes
+              rapids_cuda_version: '12.9.1'
+              rapids_py_version: '3.11'
+            - name: x86_64_cuda13_py311
+              demands: ucx_docker -equals yes
+              rapids_cuda_version: '13.2.0'
+              rapids_py_version: '3.11'
+            - name: aarch64_cuda12_py311
+              demands: ucx_arm64
+              rapids_cuda_version: '12.9.1'
+              rapids_py_version: '3.11'
+            - name: aarch64_cuda13_py311
+              demands: ucx_arm64
+              rapids_cuda_version: '13.2.0'
+              rapids_py_version: '3.11'
+      # 2.3: wheel-build-libucxx. Upstream filter = max_by(PY_VER) per (ARCH, CUDA major).
+      - template: ucxx_wheel_libucxx_build.yml
+        parameters:
+          container: ucxx_rapidsai_ci_wheel
+          slices:
+            - name: x86_64_cuda12_py314
+              demands: ucx_docker -equals yes
+              rapids_cuda_version: '12.9.1'
+              rapids_py_version: '3.14'
+            - name: x86_64_cuda13_py314
+              demands: ucx_docker -equals yes
+              rapids_cuda_version: '13.0.2'
+              rapids_py_version: '3.14'
+            - name: aarch64_cuda12_py314
+              demands: ucx_arm64
+              rapids_cuda_version: '12.9.1'
+              rapids_py_version: '3.14'
+            - name: aarch64_cuda13_py314
+              demands: ucx_arm64
+              rapids_cuda_version: '13.0.2'
+              rapids_py_version: '3.14'
+      # 2.4: wheel-build-ucxx. Upstream filter = min_by(PY_VER) per (CUDA, ARCH).
+      # Each ucxx wheel slice consumes a libucxx wheel slice; pair by (ARCH, CUDA).
+      - template: ucxx_wheel_ucxx_build.yml
+        parameters:
+          container: ucxx_rapidsai_ci_wheel
+          slices:
+            - name: x86_64_cuda12_py311
+              libucxx_slice: x86_64_cuda12_py314
+              demands: ucx_docker -equals yes
+              rapids_cuda_version: '12.9.1'
+              rapids_py_version: '3.11'
+            - name: x86_64_cuda13_py311
+              libucxx_slice: x86_64_cuda13_py314
+              demands: ucx_docker -equals yes
+              rapids_cuda_version: '13.0.2'
+              rapids_py_version: '3.11'
+            - name: aarch64_cuda12_py311
+              libucxx_slice: aarch64_cuda12_py314
+              demands: ucx_arm64
+              rapids_cuda_version: '12.9.1'
+              rapids_py_version: '3.11'
+            - name: aarch64_cuda13_py311
+              libucxx_slice: aarch64_cuda13_py314
+              demands: ucx_arm64
+              rapids_cuda_version: '13.0.2'
+              rapids_py_version: '3.11'
+      # 2.5: wheel-build-distributed-ucxx. Pure-python; upstream emits 2 amd64 legs.
+      - template: ucxx_wheel_distributed_ucxx_build.yml
+        parameters:
+          container: ucxx_rapidsai_ci_wheel
+          slices:
+            - name: x86_64_cuda12_py314
+              demands: ucx_docker -equals yes
+              rapids_cuda_version: '12.9.1'
+              rapids_py_version: '3.14'
+            - name: x86_64_cuda13_py314
+              demands: ucx_docker -equals yes
+              rapids_cuda_version: '13.0.2'
+              rapids_py_version: '3.14'
+
   - stage: Tests
+    condition: false  # temp: skip non-UCXX stage; restore before merge
     dependsOn: [Basic_compile]
     jobs:
     - template: tests.yml
@@ -410,6 +541,7 @@ stages:
             demands: ucx_gpu -equals yes
 
   - stage: Build_Static
+    condition: false  # temp: skip non-UCXX stage; restore before merge
     dependsOn: [Basic_compile]
     jobs:
       - job: build_source
@@ -456,12 +588,14 @@ stages:
 
 
   - stage: Cuda
+    condition: false  # temp: skip non-UCXX stage; restore before merge
     dependsOn: [Basic_compile]
     jobs:
       - template: cuda/cuda.yml
 
 
   - stage: AddressSanitizer
+    condition: false  # temp: skip non-UCXX stage; restore before merge
     dependsOn: [Static_check]
     jobs:
     - template: tests.yml

buildlib/pr/ucxx_conda_cpp_build.yml

@@ -0,0 +1,40 @@
+parameters:
+  container: ucxx_rapidsai_ci_conda
+  slices: []
+
+jobs:
+  - ${{ each slice in parameters.slices }}:
+    - job: ucxx_conda_cpp_build_${{ slice.name }}
+      workspace:
+        clean: all
+      pool:
+        name: MLNX
+        demands: ${{ slice.demands }}
+      displayName: 'UCXX conda-cpp-build (${{ slice.name }})'
+      container: ${{ parameters.container }}
+      timeoutInMinutes: 60
+      variables:
+        UCX_DIR: $(Agent.BuildDirectory)/ucx
+        UCXX_DIR: $(Agent.BuildDirectory)/ucxx
+        RAPIDS_CUDA_VERSION: ${{ slice.rapids_cuda_version }}
+        RAPIDS_PY_VERSION: ${{ slice.rapids_py_version }}
+        RAPIDS_BLD_OUTPUT_DIR: $(Build.ArtifactStagingDirectory)
+
+      steps:
+        - checkout: self
+          path: ucx
+          fetchDepth: 100
+          retryCountOnTaskFailure: 5
+
+        - checkout: ucxx
+          path: ucxx
+          retryCountOnTaskFailure: 5
+
+        - bash: bash $(UCX_DIR)/buildlib/tools/build_ucxx.sh conda_cpp
+          displayName: Build UCXX conda C++ package
+
+        - task: PublishBuildArtifacts@1
+          displayName: Publish conda-cpp artifact
+          inputs:
+            pathToPublish: $(RAPIDS_BLD_OUTPUT_DIR)
+            artifactName: ucxx-conda-cpp-${{ slice.name }}

buildlib/pr/ucxx_conda_python_build.yml

@@ -0,0 +1,53 @@
+parameters:
+  container: ucxx_rapidsai_ci_conda
+  slices: []
+
+jobs:
+  - ${{ each slice in parameters.slices }}:
+    - job: ucxx_conda_python_build_${{ slice.name }}
+      workspace:
+        clean: all
+      pool:
+        name: MLNX
+        demands: ${{ slice.demands }}
+      displayName: 'UCXX conda-python-build (${{ slice.name }})'
+      container: ${{ parameters.container }}
+      timeoutInMinutes: 60
+      dependsOn: ucxx_conda_cpp_build_${{ slice.name }}
+      variables:
+        UCX_DIR: $(Agent.BuildDirectory)/ucx
+        UCXX_DIR: $(Agent.BuildDirectory)/ucxx
+        RAPIDS_CUDA_VERSION: ${{ slice.rapids_cuda_version }}
+        RAPIDS_PY_VERSION: ${{ slice.rapids_py_version }}
+        RAPIDS_BLD_OUTPUT_DIR: $(Build.ArtifactStagingDirectory)
+
+      steps:
+        - checkout: self
+          path: ucx
+          fetchDepth: 100
+          retryCountOnTaskFailure: 5
+
+        - checkout: ucxx
+          path: ucxx
+          retryCountOnTaskFailure: 5
+
+        - task: DownloadBuildArtifacts@1
+          displayName: Fetch conda-cpp artifact
+          inputs:
+            buildType: current
+            artifactName: ucxx-conda-cpp-${{ slice.name }}
+            downloadPath: $(System.DefaultWorkingDirectory)/_dl
+
+        - bash: |
+            rm -rf "$(RAPIDS_BLD_OUTPUT_DIR)"
+            mv "$(System.DefaultWorkingDirectory)/_dl/ucxx-conda-cpp-${{ slice.name }}" "$(RAPIDS_BLD_OUTPUT_DIR)"
+          displayName: Stage conda-cpp artifact
+
+        - bash: bash $(UCX_DIR)/buildlib/tools/build_ucxx.sh conda_python
+          displayName: Build UCXX conda Python package
+
+        - task: PublishBuildArtifacts@1
+          displayName: Publish conda-python artifact
+          inputs:
+            pathToPublish: $(RAPIDS_BLD_OUTPUT_DIR)
+            artifactName: ucxx-conda-python-${{ slice.name }}