Update black to 22.3.0

.pre-commit-config.yaml

@@ -4,7 +4,7 @@ repos:
         hooks:
               - id: isort
       - repo: https://github.com/ambv/black
-        rev: 19.10b0
+        rev: 22.3.0
         hooks:
               - id: black
       - repo: https://gitlab.com/pycqa/flake8

dask_cuda/__init__.py

@@ -19,8 +19,10 @@
 
 
 # Monkey patching Dask to make use of explicit-comms when `DASK_EXPLICIT_COMMS=True`
-dask.dataframe.shuffle.rearrange_by_column_tasks = get_rearrange_by_column_tasks_wrapper(
-    dask.dataframe.shuffle.rearrange_by_column_tasks
+dask.dataframe.shuffle.rearrange_by_column_tasks = (
+    get_rearrange_by_column_tasks_wrapper(
+        dask.dataframe.shuffle.rearrange_by_column_tasks
+    )
 )
 
 

dask_cuda/benchmarks/local_cudf_merge.py

@@ -36,7 +36,7 @@ def generate_chunk(i_chunk, local_size, num_chunks, chunk_type, frac_match, gpu)
         import numpy as xp
         import pandas as xdf
 
-    xp.random.seed(2 ** 32 - 1)
+    xp.random.seed(2**32 - 1)
 
     chunk_type = chunk_type or "build"
     frac_match = frac_match or 1.0
@@ -258,7 +258,10 @@ def main(args):
         for (w1, w2), v in bandwidths.items()
     }
     total_nbytes = {
-        (scheduler_workers[w1].name, scheduler_workers[w2].name,): format_bytes(sum(nb))
+        (
+            scheduler_workers[w1].name,
+            scheduler_workers[w2].name,
+        ): format_bytes(sum(nb))
         for (w1, w2), nb in total_nbytes.items()
     }
 
@@ -379,21 +382,30 @@ def main(args):
 def parse_args():
     special_args = [
         {
-            "name": ["-b", "--backend",],
+            "name": [
+                "-b",
+                "--backend",
+            ],
             "choices": ["dask", "explicit-comms"],
             "default": "dask",
             "type": str,
             "help": "The backend to use.",
         },
         {
-            "name": ["-t", "--type",],
+            "name": [
+                "-t",
+                "--type",
+            ],
             "choices": ["cpu", "gpu"],
             "default": "gpu",
             "type": str,
             "help": "Do merge with GPU or CPU dataframes",
         },
         {
-            "name": ["-c", "--chunk-size",],
+            "name": [
+                "-c",
+                "--chunk-size",
+            ],
             "default": 1_000_000,
             "metavar": "n",
             "type": int,
@@ -444,9 +456,17 @@ def parse_args():
             "action": "store_true",
             "help": "Write output as markdown",
         },
-        {"name": "--runs", "default": 3, "type": int, "help": "Number of runs",},
         {
-            "name": ["-s", "--set-index",],
+            "name": "--runs",
+            "default": 3,
+            "type": int,
+            "help": "Number of runs",
+        },
+        {
+            "name": [
+                "-s",
+                "--set-index",
+            ],
             "action": "store_true",
             "help": "Call set_index on the key column to sort the joined dataframe.",
         },

dask_cuda/benchmarks/local_cudf_shuffle.py

@@ -138,7 +138,10 @@ def main(args):
         for (w1, w2), v in bandwidths.items()
     }
     total_nbytes = {
-        (scheduler_workers[w1].name, scheduler_workers[w2].name,): format_bytes(sum(nb))
+        (
+            scheduler_workers[w1].name,
+            scheduler_workers[w2].name,
+        ): format_bytes(sum(nb))
         for (w1, w2), nb in total_nbytes.items()
     }
 
@@ -251,14 +254,20 @@ def parse_args():
             "help": "Number of input partitions (default '100')",
         },
         {
-            "name": ["-b", "--backend",],
+            "name": [
+                "-b",
+                "--backend",
+            ],
             "choices": ["dask", "explicit-comms"],
             "default": "dask",
             "type": str,
             "help": "The backend to use.",
         },
         {
-            "name": ["-t", "--type",],
+            "name": [
+                "-t",
+                "--type",
+            ],
             "choices": ["cpu", "gpu"],
             "default": "gpu",
             "type": str,
@@ -276,7 +285,12 @@ def parse_args():
             "action": "store_true",
             "help": "Write output as markdown",
         },
-        {"name": "--runs", "default": 3, "type": int, "help": "Number of runs",},
+        {
+            "name": "--runs",
+            "default": 3,
+            "type": int,
+            "help": "Number of runs",
+        },
     ]
 
     return parse_benchmark_args(

dask_cuda/benchmarks/local_cupy.py

@@ -229,9 +229,10 @@ async def run(args):
                 for (w1, w2), v in bandwidths.items()
             }
             total_nbytes = {
-                (scheduler_workers[w1].name, scheduler_workers[w2].name,): format_bytes(
-                    sum(nb)
-                )
+                (
+                    scheduler_workers[w1].name,
+                    scheduler_workers[w2].name,
+                ): format_bytes(sum(nb))
                 for (w1, w2), nb in total_nbytes.items()
             }
 
@@ -318,35 +319,50 @@ async def run(args):
 def parse_args():
     special_args = [
         {
-            "name": ["-s", "--size",],
+            "name": [
+                "-s",
+                "--size",
+            ],
             "default": "10000",
             "metavar": "n",
             "type": int,
             "help": "The array size n in n^2 (default 10000). For 'svd' operation "
             "the second dimension is given by --second-size.",
         },
         {
-            "name": ["-2", "--second-size",],
+            "name": [
+                "-2",
+                "--second-size",
+            ],
             "default": "1000",
             "type": int,
             "help": "The second dimension size for 'svd' operation (default 1000).",
         },
         {
-            "name": ["-t", "--type",],
+            "name": [
+                "-t",
+                "--type",
+            ],
             "choices": ["cpu", "gpu"],
             "default": "gpu",
             "type": str,
             "help": "Do merge with GPU or CPU dataframes.",
         },
         {
-            "name": ["-o", "--operation",],
+            "name": [
+                "-o",
+                "--operation",
+            ],
             "default": "transpose_sum",
             "type": str,
             "help": "The operation to run, valid options are: "
             "'transpose_sum' (default), 'dot', 'fft', 'svd', 'sum', 'mean', 'slice'.",
         },
         {
-            "name": ["-c", "--chunk-size",],
+            "name": [
+                "-c",
+                "--chunk-size",
+            ],
             "default": "2500",
             "type": int,
             "help": "Chunk size (default 2500).",

dask_cuda/benchmarks/local_cupy_map_overlap.py

@@ -118,9 +118,10 @@ async def run(args):
                 for (w1, w2), v in bandwidths.items()
             }
             total_nbytes = {
-                (scheduler_workers[w1].name, scheduler_workers[w2].name,): format_bytes(
-                    sum(nb)
-                )
+                (
+                    scheduler_workers[w1].name,
+                    scheduler_workers[w2].name,
+                ): format_bytes(sum(nb))
                 for (w1, w2), nb in total_nbytes.items()
             }
 
@@ -198,28 +199,40 @@ async def run(args):
 def parse_args():
     special_args = [
         {
-            "name": ["-s", "--size",],
+            "name": [
+                "-s",
+                "--size",
+            ],
             "default": "10000",
             "metavar": "n",
             "type": int,
             "help": "The size n in n^2 (default 10000)",
         },
         {
-            "name": ["-t", "--type",],
+            "name": [
+                "-t",
+                "--type",
+            ],
             "choices": ["cpu", "gpu"],
             "default": "gpu",
             "type": str,
             "help": "Use GPU or CPU arrays",
         },
         {
-            "name": ["-c", "--chunk-size",],
+            "name": [
+                "-c",
+                "--chunk-size",
+            ],
             "default": "128 MiB",
             "metavar": "nbytes",
             "type": str,
             "help": "Chunk size (default '128 MiB')",
         },
         {
-            "name": ["-k", "--kernel-size",],
+            "name": [
+                "-k",
+                "--kernel-size",
+            ],
             "default": "1",
             "metavar": "k",
             "type": int,
@@ -232,7 +245,12 @@ def parse_args():
             "type": parse_bytes,
             "help": "Ignore messages smaller than this (default '1 MB')",
         },
-        {"name": "--runs", "default": 3, "type": int, "help": "Number of runs",},
+        {
+            "name": "--runs",
+            "default": 3,
+            "type": int,
+            "help": "Number of runs",
+        },
     ]
 
     return parse_benchmark_args(

dask_cuda/benchmarks/utils.py

@@ -62,7 +62,9 @@ def parse_benchmark_args(description="Generic dask-cuda Benchmark", args_list=[]
         "Logging is only enabled if RMM memory pool is enabled.",
     )
     parser.add_argument(
-        "--all-to-all", action="store_true", help="Run all-to-all before computation",
+        "--all-to-all",
+        action="store_true",
+        help="Run all-to-all before computation",
     )
     parser.add_argument(
         "--enable-tcp-over-ucx",
@@ -241,7 +243,10 @@ def get_scheduler_workers(dask_scheduler=None):
 
 
 def setup_memory_pool(
-    dask_worker=None, pool_size=None, disable_pool=False, log_directory=None,
+    dask_worker=None,
+    pool_size=None,
+    disable_pool=False,
+    log_directory=None,
 ):
     import cupy
 

dask_cuda/cli/dask_cuda_worker.py

@@ -126,7 +126,10 @@
     allows querying the amount of memory allocated by RMM.""",
 )
 @click.option(
-    "--pid-file", type=str, default="", help="File to write the process PID.",
+    "--pid-file",
+    type=str,
+    default="",
+    help="File to write the process PID.",
 )
 @click.option(
     "--resources",
@@ -314,7 +317,9 @@ def main(
 ):
     if tls_ca_file and tls_cert and tls_key:
         security = Security(
-            tls_ca_file=tls_ca_file, tls_worker_cert=tls_cert, tls_worker_key=tls_key,
+            tls_ca_file=tls_ca_file,
+            tls_worker_cert=tls_cert,
+            tls_worker_key=tls_key,
         )
     else:
         security = None

dask_cuda/explicit_comms/comms.py

@@ -123,7 +123,7 @@ async def server_handler(ep):
 
 
 async def _create_endpoints(session_state, peers):
-    """ Each worker creates a UCX endpoint to all workers with greater rank"""
+    """Each worker creates a UCX endpoint to all workers with greater rank"""
     assert session_state["loop"] is asyncio.get_event_loop()
 
     myrank = session_state["rank"]

dask_cuda/local_cuda_cluster.py

@@ -378,7 +378,9 @@ def new_worker_spec(self):
         visible_devices = cuda_visible_devices(worker_count, self.cuda_visible_devices)
         spec["options"].update(
             {
-                "env": {"CUDA_VISIBLE_DEVICES": visible_devices,},
+                "env": {
+                    "CUDA_VISIBLE_DEVICES": visible_devices,
+                },
                 "plugins": {
                     CPUAffinity(
                         get_cpu_affinity(nvml_device_index(0, visible_devices))

dask_cuda/proxy_object.py

@@ -378,7 +378,9 @@ def __del__(self):
         pxy.manager.remove(self)
 
     def _pxy_serialize(
-        self, serializers: Iterable[str], proxy_detail: ProxyDetail = None,
+        self,
+        serializers: Iterable[str],
+        proxy_detail: ProxyDetail = None,
     ) -> None:
         """Inplace serialization of the proxied object using the `serializers`
 

dask_cuda/tests/test_dask_cuda_worker.py

@@ -206,7 +206,12 @@ def test_pre_import(loop):  # noqa: F811
 
     with popen(["dask-scheduler", "--port", "9369", "--no-dashboard"]):
         with popen(
-            ["dask-cuda-worker", "127.0.0.1:9369", "--pre-import", module,]
+            [
+                "dask-cuda-worker",
+                "127.0.0.1:9369",
+                "--pre-import",
+                module,
+            ]
         ):
             with Client("127.0.0.1:9369", loop=loop) as client:
                 assert wait_workers(client, n_gpus=get_n_gpus())