ENH: unpin "numpy<2.0" & move ucx-py to ucxx (xorbitsai#106)

.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):