ENH: unpin "numpy<2.0" & move ucx-py to ucxx

.github/workflows/python.yaml

@@ -106,9 +106,9 @@ jobs:
         python-version: ${{ matrix.python-version }}
         activate-environment: ${{ env.CONDA_ENV }}
 
-    # Fix "version `GLIBCXX_3.4.30' not found (required by xoscar_store.cpython-311-x86_64-linux-gnu.so)" issue in Python 3.11
-    - name: Install libstdcxx-ng for Python 3.11
-      if: ${{ (matrix.module != 'gpu') && (matrix.os == 'ubuntu-latest') && (matrix.python-version == '3.11') }}
+    # Fix "version `GLIBCXX_3.4.30' not found (required by xoscar_store.cpython-311-x86_64-linux-gnu.so)" issue
+    - name: Install libstdcxx-ng
+      if: ${{ (matrix.module != 'gpu') && (matrix.os == 'ubuntu-latest') }}
       run: |
         conda install -c conda-forge libstdcxx-ng
 
@@ -122,9 +122,11 @@ jobs:
       working-directory: ./python
 
     - name: Install ucx dependencies
-      if: ${{ (matrix.module != 'gpu') && (matrix.os == 'ubuntu-latest') && (matrix.python-version != '3.11') }}
+      if: ${{ (matrix.os == 'ubuntu-latest') && (matrix.python-version >= '3.9')}}
       run: |
-        conda install -c conda-forge -c rapidsai ucx-proc=*=cpu ucx ucx-py
+        # ucx-py move to ucxx and ucxx-cu12 can be run on CPU
+        # conda install -c conda-forge -c rapidsai ucx-proc=*=cpu ucx ucx-py
+        pip install ucxx-cu12
 
     - name: Install fury
       if: ${{ (matrix.module != 'gpu') && (matrix.os == 'ubuntu-latest') && (matrix.python-version == '3.9') }}
@@ -137,7 +139,7 @@ jobs:
         source activate ${{ env.CONDA_ENV }}
         conda install -y conda-forge::nccl=2.22.3
         pip install --extra-index-url=https://pypi.nvidia.com cudf-cu12
-        pip install ucx-py-cu12 cloudpickle psutil tblib uvloop packaging "numpy<2.0.0" scipy cython coverage flaky
+        pip install ucxx-cu12 cloudpickle psutil tblib uvloop packaging "numpy<2.0.0" scipy cython coverage flaky
         python setup.py clean --all
         pip install -e ./
       working-directory: ./python

CI/conda-environment.yml

@@ -2,7 +2,7 @@ name: xoscar-test
 channels:
   - defaults
 dependencies:
-  - numpy<2.0.0
+  - numpy
   - cloudpickle
   - coverage
   - cython

CI/requirements-wheel.txt

@@ -1,6 +1,6 @@
 oldest-supported-numpy
 
-numpy<2.0.0
+numpy
 packaging
 wheel
 

python/pyproject.toml

@@ -21,7 +21,7 @@ requires = [
     "pandas==1.4.0; python_version>='3.10' and python_version<'3.11' and platform_machine=='arm64'",
     "pandas==1.5.1; python_version>='3.11' and python_version<'3.12'",
     "pandas>=2.1.1; python_version>'3.11'",
-    "numpy<2.0.0",
+    "numpy",
     "cython>=0.29.33",
     "requests>=2.4.0",
     "cloudpickle>=2.2.1; python_version>='3.11'",

python/setup.cfg

@@ -24,7 +24,7 @@ zip_safe = False
 include_package_data = True
 packages = find:
 install_requires =
-    numpy>=1.14.0,<2.0.0
+    numpy>=1.14.0
     pandas>=1.0.0
     scipy>=1.0.0; sys_platform!="win32" or python_version>="3.10"
     scipy>=1.0.0,<=1.9.1; sys_platform=="win32" and python_version<"3.10"

python/xoscar/backends/communication/ucx.py

@@ -34,7 +34,7 @@
 from .core import register_client, register_server
 from .errors import ChannelClosed
 
-ucp = lazy_import("ucp")
+ucp = lazy_import("ucxx")
 numba_cuda = lazy_import("numba.cuda")
 rmm = lazy_import("rmm")
 
@@ -86,7 +86,7 @@ def _get_options(ucx_config: dict) -> Tuple[dict, dict]:
                 tls += ",cuda_copy"
 
             if ucx_config.get("infiniband"):  # pragma: no cover
-                tls = "rc," + tls
+                tls = "ib," + tls
             if ucx_config.get("nvlink"):  # pragma: no cover
                 tls += ",cuda_ipc"
 
@@ -177,7 +177,8 @@ def init(ucx_config: dict):
         new_environ.update(envs)
         os.environ = new_environ  # type: ignore
         try:
-            ucp.init(options=options, env_takes_precedence=True)
+            # let UCX determine the appropriate transports
+            ucp.init()
         finally:
             os.environ = original_environ
 
@@ -313,7 +314,7 @@ async def send_buffers(self, buffers: list, meta: Optional[_MessageBase] = None)
                         await self.ucp_endpoint.send(buf)
                 for buffer in buffers:
                     await self.ucp_endpoint.send(buffer)
-        except ucp.exceptions.UCXBaseException:  # pragma: no cover
+        except ucp.exceptions.UCXError:  # pragma: no cover
             self.abort()
             raise ChannelClosed("While writing, the connection was closed")
 
@@ -516,7 +517,7 @@ async def connect(
 
         try:
             ucp_endpoint = await ucp.create_endpoint(host, port)
-        except ucp.exceptions.UCXBaseException as e:  # pragma: no cover
+        except ucp.exceptions.UCXError as e:  # pragma: no cover
             raise ChannelClosed(
                 f"Connection closed before handshake completed, "
                 f"local address: {local_address}, dest address: {dest_address}"

python/xoscar/serialization/core.pyx

@@ -55,7 +55,6 @@ from .pyfury import get_fury
 
 BUFFER_PICKLE_PROTOCOL = max(pickle.DEFAULT_PROTOCOL, 5)
 cdef bint HAS_PICKLE_BUFFER = pickle.HIGHEST_PROTOCOL >= 5
-cdef bint _PANDAS_HAS_MGR = hasattr(pd.Series([0]), "_mgr")
 
 cdef TypeDispatcher _serial_dispatcher = TypeDispatcher()
 cdef dict _deserializers = dict()
@@ -260,16 +259,7 @@ def unpickle_buffers(list buffers):
     else:
         result = cloudpickle.loads(buffers[0], buffers=buffers[1:])
 
-    # as pandas prior to 1.1.0 use _data instead of _mgr to hold BlockManager,
-    # deserializing from high versions may produce mal-functioned pandas objects,
-    # thus the patch is needed
-    if _PANDAS_HAS_MGR:
-        return result
-    else:  # pragma: no cover
-        if hasattr(result, "_mgr") and isinstance(result, (pd.DataFrame, pd.Series)):
-            result._data = getattr(result, "_mgr")
-            delattr(result, "_mgr")
-        return result
+    return result
 
 
 cdef class PickleSerializer(Serializer):