Update cuDF merge benchmark

[pentschev]

This adds several bugfixes and improvements:

• Correctly assign CUDA_VISIBLE_DEVICES;
• Create CUDA contexts before initializing UCX;
• Ensure merge uses the "key" column and asserts expected result size;
• Allow running the benchmark for multiple iterations;
• Allow running multiple warmup iterations;
• Prints results for each iteration;
• Allow enabling garbage collection after each itearation.

[pentschev]

Since this benchmark is broken due to changes in CUDA context handling, it would be good if this still makes 0.27 (22.08), despite the ongoing burndown. As this is only a benchmark it should not impact release in any way, as this is not executed anywhere as part of that process.

[wence-]

Since this benchmark is broken due to changes in CUDA context handling

Can you explain what you mean by this?

[wence-]

Minor nitpicks, otherwise looks good.

[pentschev]

Since this benchmark is broken due to changes in CUDA context handling

Can you explain what you mean by this?

cuDF creates a CUDA context at import time, which wasn't the case in the past, and is the reason we have rapidsai/dask-cuda#379 in Dask-CUDA, if we don't set CUDA_VISIBLE_DEVICES then a context will always be created on device 0. Added to that, we now rely on UCX to ensure IB<->GPU mapping automatically, but that requires the CUDA context to be created before import ucp.

[wence-]

FWIW, the only niggle I have left here is that the --chunks-per-dev argument is a bit misnamed. The total number of workers started up is chunks-per-dev * num_devices (so when running with all GPUs and more than chunks-per-dev > 1 we get errors because we try and build 2 (or more) RMM pools per device and things break).

[pentschev]

FWIW, the only niggle I have left here is that the --chunks-per-dev argument is a bit misnamed. The total number of workers started up is chunks-per-dev * num_devices (so when running with all GPUs and more than chunks-per-dev > 1 we get errors because we try and build 2 (or more) RMM pools per device and things break).

The RMM pool being a problem is expected with the default value, that's when the user will have to adjust --rmm-init-pool-size. So --chunks-per-dev still sounds reasonable, it would be misnamed if it was --chunks-per-process or --chunks-per-worker. Do you have a proposal on how to make this better?

[wence-]

FWIW, the only niggle I have left here is that the --chunks-per-dev argument is a bit misnamed. The total number of workers started up is chunks-per-dev * num_devices (so when running with all GPUs and more than chunks-per-dev > 1 we get errors because we try and build 2 (or more) RMM pools per device and things break).

The RMM pool being a problem is expected with the default value, that's when the user will have to adjust --rmm-init-pool-size. So --chunks-per-dev still sounds reasonable, it would be misnamed if it was --chunks-per-process or --chunks-per-worker. Do you have a proposal on how to make this better?

I had naively thought that we would get num_devices workers and chunks-per-dev would say "how many chunks does each worker have". But actually, each worker only ever has a single chunk, and we always have num_devices * chunks_per_device workers (i.e. we oversubscribe the GPUs).

Since this isn't a substantive change from the current behaviour, let's not pollute this PR with too many additional pieces (sorry, this is my fault for noticing!) and leave as-is for now.

[pentschev]

I think all changes are in now, could you check one last time/approve the PR @wence- ?

[wence-]

@gpucibot merge

[wence-]

I am happy to merge this, though does it need further approval to go to 0.27 (rather than 0.28) ?

[pentschev]

I am happy to merge this, though does it need further approval to go to 0.27 (rather than 0.28) ?

UCX-Py doesn't follow exactly the same conditions as the remaining of RAPIDS, and given there are no other required reviewers, we should be good to merge. Also gpucibot has no power here. :)

[wence-]

OK, let's wait for tests.

[pentschev]

For the very small benchmark in CI, the assertion failed:

$ python cudf-merge.py --chunks-per-dev 4 --chunk-size 10000 --rm
m-init-pool-size 2097152
Process SpawnProcess-1:
Traceback (most recent call last):
  File "/datasets/pentschev/miniconda3/envs/rn-220803/lib/python3.8/multiprocessing/process.py", line 315, in _bootstrap
    self.run()
  File "/datasets/pentschev/miniconda3/envs/rn-220803/lib/python3.8/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/datasets/pentschev/miniconda3/envs/rn-220803/lib/python3.8/site-packages/ucp/utils.py", line 131, in _worker_process
    ret = loop.run_until_complete(run())
  File "/datasets/pentschev/miniconda3/envs/rn-220803/lib/python3.8/asyncio/base_events.py", line 616, in run_until_complete
    return future.result()
  File "/datasets/pentschev/miniconda3/envs/rn-220803/lib/python3.8/site-packages/ucp/utils.py", line 126, in run
    return await func(rank, eps, args)
  File "/datasets/pentschev/src/ucx-py/benchmarks/cudf-merge.py", line 216, in worker
    assert abs(len(ret) - expected_len) <= expected_len_err

Maximum error was 30, and the actual value was 31. Since this is not critical, I increased the tolerance to 2% now.

[pentschev]

And it seems I was wrong. We have to ask for approval.

[msadang]

LGTM