From 25bcbe8261f494a16022c55e250500c3c435a661 Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Wed, 6 Sep 2023 16:58:16 -0700
Subject: [PATCH] Pin to working ucx

---
 conda/recipes/ucx-py/meta.yaml |  7 +++++++
 dependencies.yaml              | 16 +++++++++++++---
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/conda/recipes/ucx-py/meta.yaml b/conda/recipes/ucx-py/meta.yaml
index 1b88500a7..854eb1187 100644
--- a/conda/recipes/ucx-py/meta.yaml
+++ b/conda/recipes/ucx-py/meta.yaml
@@ -5,6 +5,7 @@
 {% set version = environ.get('GIT_DESCRIBE_TAG', '0.0.0.dev').lstrip('v') %}
 {% set py_version = environ['CONDA_PY'] %}
 {% set cuda_version = '.'.join(environ['RAPIDS_CUDA_VERSION'].split('.')[:2]) %}
+{% set cuda_major = cuda_version.split('.')[0] %}
 {% set date_string = environ['RAPIDS_DATE_STRING'] %}
 
 package:
@@ -27,7 +28,13 @@ requirements:
   host:
     - python
     - pip
+    {% if cuda_major == "11" %}
     - ucx
+    {% else %}
+    # Temporary pin because the subsequent build (h64cca9d_3) is missing proper
+    # CUDA support
+    - ucx==1.14.1=h195a15c_3
+    {% endif %}
     {% for r in data.get("build-system", {}).get("requires", []) %}
     - {{ r }}
     {% endfor %}
diff --git a/dependencies.yaml b/dependencies.yaml
index 1634ae829..8b90862b5 100644
--- a/dependencies.yaml
+++ b/dependencies.yaml
@@ -108,9 +108,19 @@ dependencies:
         packages:
           - numpy>=1.21
           - pynvml>=11.4.1
-      - output_types: [conda]
-        packages:
-          - ucx
+    specific:
+      - output_types: conda
+        matrices:
+          - matrix:
+              cuda: "12.0"
+            packages:
+              # Temporary pin because the subsequent build (h64cca9d_3) is
+              # missing proper CUDA support
+              - ucx==1.14.1=h195a15c_3
+          - matrix:
+              cuda: "11.8"
+            packages:
+              - ucx
   test_python:
     common:
       - output_types: [conda, requirements, pyproject]