Add device synchonization for local CuPy benchmarks with Dask profiling

[charlesbluca]

This PR ensures CUDA device synchronization for local CuPy benchmarks when profiling with Dask performance reports.

Prior to this, the run times of benchmarks running over UCX would vary drastically depending on if a performance report was generated or not:

>>> python local_cupy.py -p ucx --ucx-net-devices auto
Roundtrip benchmark
--------------------------
Operation          | transpose_sum
User size          | 10000
User second size   | 1000
User chunk-size    | 2500
Compute shape      | (10000, 10000)
Compute chunk-size | (2500, 2500)
Ignore-size        | 1.05 MB
Protocol           | ucx
Device(s)          | 0
Worker Thread(s)   | 1
==========================
Wall-clock         | npartitions
--------------------------
2.78 s             | 16
2.02 s             | 16
2.02 s             | 16
==========================
(w1,w2)            | 25% 50% 75% (total nbytes)
--------------------------

>>> python local_cupy.py -p ucx --ucx-net-devices auto --profile dask-report.html
Roundtrip benchmark
--------------------------
Operation          | transpose_sum
User size          | 10000
User second size   | 1000
User chunk-size    | 2500
Compute shape      | (10000, 10000)
Compute chunk-size | (2500, 2500)
Ignore-size        | 1.05 MB
Protocol           | ucx
Device(s)          | 0
Worker Thread(s)   | 1
==========================
Wall-clock         | npartitions
--------------------------
2.79 s             | 16
49.28 ms           | 16
49.22 ms           | 16
==========================
(w1,w2)            | 25% 50% 75% (total nbytes)
--------------------------

Additionally, I suppressed flake8 rule E402 in docs/source/conf.py, as that had been causing the pre-commit hooks to fail up until now. I'm okay with removing that change if it would be better suited for a separate PR.

[jakirkham]

cc @pentschev

[pentschev]

LGTM, thanks @charlesbluca !

[pentschev]

Just FYI on labels @charlesbluca , you need at least one of each:

• Current state: labels beginning with 0 to 5;
• Breaking/non-breaking
• Bug/feature

Otherwise the label checker won't pass.

[quasiben]

Woohoo! Thanks @charlesbluca !

[jakirkham]

Are you able to select labels Charles? If not, please let us know

[charlesbluca]

No, I don't think I am - do I need additional repo privileges for that?

[codecov-io]

Codecov Report

Merging #533 (3864e7e) into branch-0.18 (32d9d33) will increase coverage by 2.02%.
The diff coverage is 96.35%.

@@               Coverage Diff               @@
##           branch-0.18     #533      +/-   ##
===============================================
+ Coverage        90.42%   92.45%   +2.02%     
===============================================
  Files               15       16       +1     
  Lines             1128     1550     +422     
===============================================
+ Hits              1020     1433     +413     
- Misses             108      117       +9     

Impacted Files	Coverage Δ
dask_cuda/cli/dask_cuda_worker.py	96.92% <ø> (ø)
dask_cuda/cuda_worker.py	78.75% <75.00%> (+1.73%)	⬆️
dask_cuda/device_host_file.py	90.90% <80.00%> (-7.96%)	⬇️
dask_cuda/get_device_memory_objects.py	89.04% <89.04%> (ø)
dask_cuda/proxify_device_objects.py	93.87% <93.87%> (ø)
dask_cuda/proxy_object.py	91.21% <96.07%> (+3.42%)	⬆️
dask_cuda/explicit_comms/dataframe/merge.py	96.24% <96.24%> (ø)
dask_cuda/explicit_comms/dataframe/shuffle.py	98.51% <98.51%> (ø)
dask_cuda/__init__.py	100.00% <100.00%> (ø)
dask_cuda/explicit_comms/comms.py	98.78% <100.00%> (-0.23%)	⬇️
... and 8 more

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update d78c668...3864e7e. Read the comment docs.

[pentschev]

I created a request on your behalf @charlesbluca .

[pentschev]

@gpucibot merge

[jakirkham]

Thanks Peter :)

[pentschev]

@gpucibot merge