ucx

.github/workflows/python.yaml

@@ -117,15 +117,16 @@ jobs:
  MODULE: ${{ matrix.module }}
  if: ${{ matrix.module != 'gpu' }}
  run: |
- pip install numpy scipy cython coverage flaky
+ pip install numpy scipy cython coverage flaky ucxx-cu12
  pip install -e ".[dev,extra]"
+ ucx_info -v
  working-directory: ./python
 
- - name: Install ucx dependencies
- if: ${{ (matrix.module != 'gpu') && (matrix.os == 'ubuntu-latest') && (matrix.python-version != '3.11') }}
- run: |
- conda install -c conda-forge -c rapidsai ucx-proc=*=cpu ucx ucx-py
- pip install -U numpy
+ # - name: Install ucx dependencies
+ #  if: ${{ (matrix.module != 'gpu') && (matrix.os == 'ubuntu-latest') && (matrix.python-version != '3.11') }}
+ #  run: |
+ #  conda install -c conda-forge -c rapidsai ucx-proc=*=cpu ucx ucx-py
+ #  pip install -U numpy
 
  - name: Install fury
  if: ${{ (matrix.module != 'gpu') && (matrix.os == 'ubuntu-latest') && (matrix.python-version == '3.9') }}

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}"