Support CUDA 12.2

[jameslamb]

Description

• switches to CUDA 12.2.2 for building conda packages and wheels
• adds new tests running against CUDA 12.2.2

Notes for Reviewers

This is part of ongoing work to build and test packages against CUDA 12.2.2 across all of RAPIDS.

For more details see:

• Add support for CUDA 12.2 wheels build-planning#7
• add CUDA 12.2 conda and wheel jobs shared-workflows#166

Planning a second round of PRs to revert these references back to a proper branch-24.{nn} release branch of shared-workflows once rapidsai/shared-workflows#166 is merged.

(created with rapids-reviser)

[jameslamb]

conda builds and tests on CUDA 12.0.1 and 12.2.2 are segfaulting (build link).

I don't see similar errors on recent builds on branch-0.36 or other PRs.

[pentschev]

conda builds and tests on CUDA 12.0.1 and 12.2.2 are segfaulting (build link).

I don't see similar errors on recent builds on branch-0.36 or other PRs.

Those are indeed flaky tests, they don't always happen in same form. Let's first try rerunning them, based on the log for 12.2 almost all tests passed, the one that failed is one of the very last ones and likely to pass soon.

[jameslamb]

Ok thank you!

[jakirkham]

Thanks James & Peter! 🙏

Looks like that cleared things up and builds now pass

Though am going to mark this "do not merge" for the moment while we work through issues in the other PRs

[jakirkham]

Looks like there is a conflict here

It is a generated file. So can just regenerate it with rapids-dependency-file-generator

[jameslamb]

Looks like there is a conflict here

fixed in 3dbdb61

[jameslamb]

I see a new test failure on the v11.8.0 conda tests.

There are quite a few stacktraces and things that look like errors in the logs, but only a single unit test case failure.

I think these snippets summarize it well:

2024-01-25 20:38:01,988 - distributed.worker - ERROR - Failed to communicate with scheduler during heartbeat.
Traceback (most recent call last):
  File "/tmp/distributed/distributed/core.py", line 1564, in _connect
    comm = await connect(
  File "/tmp/distributed/distributed/comm/core.py", line 342, in connect
    comm = await wait_for(
  File "/tmp/distributed/distributed/utils.py", line 1940, in wait_for
    return await asyncio.wait_for(fut, timeout)
  File "/opt/conda/envs/test/lib/python3.10/asyncio/tasks.py", line 432, in wait_for
    await waiter
asyncio.exceptions.CancelledError
...
distributed.comm.core.CommClosedError: ConnectionPool closing.
asyncio.exceptions.InvalidStateError: invalid state
asyncio.exceptions.InvalidStateError: invalid state
...
_____________________ ERROR at teardown of test_transpose ______________________

    @pytest.fixture(scope="function")
    def ucxx_loop():
...
        with check_thread_leak():
            yield loop
>           ucxx.reset()
...
>               raise UCXError(msg)
E               ucxx.UCXXError: Trying to reset UCX but not all Endpoints and/or Listeners are closed(). The following objects are still referencing ApplicationContext: 
E                 <frame at 0x7fdfe400e4c0, file '/opt/conda/envs/test/lib/python3.10/site-packages/ucxx/_lib_async/listener.py', line 155, code _listener_handler_coroutine>
E                 <frame at 0x55ddca529910, file '/opt/conda/envs/test/lib/python3.10/site-packages/ucxx/_lib_async/listener.py', line 155, code _listener_handler_coroutine>
...
/opt/conda/envs/test/lib/python3.10/site-packages/ucxx/core.py:111: UCXXError
ERROR python/distributed-ucxx/distributed_ucxx/tests/test_ucxx.py::test_transpose - ucxx.UCXXError: Trying to reset UCX but not all Endpoints and/or Listeners are closed(). The following objects are still referencing ApplicationContext: 
  <frame at 0x7fdfe400e4c0, file '/opt/conda/envs/test/lib/python3.10/site-packages/ucxx/_lib_async/listener.py', line 155, code _listener_handler_coroutine>
  <frame at 0x55ddca529910, file '/opt/conda/envs/test/lib/python3.10/site-packages/ucxx/_lib_async/listener.py', line 155, code _listener_handler_coroutine>

(build link)

NOTE: as of #174, rapids-dask-dependency is getting pulled in, so the tests are getting the latest main branch of dask and distributed, not releases. Maybe this error lies in some change between the latest distributed / dask releases and current main branch of those projects.

[jakirkham]

Vyas restarted CI as we believe this is an unrelated flaky test

[vyasr]

I'm guessing it's because this PR was started before #145, but it needs updating for the wheel jobs. The conda outputs look sensible.

I was going to ask some questions about the necessity of the CUDA dependency in ucx-py because it's pure Python, but I realize those are less relevant here since there is an actual C++ dependency in ucxx since it is not a pure Python package and we do need to compile against CUDA.

[jakirkham]

/merge