code: def test_decide_worker_coschedule_order_neighbors(ndeps, nthreads):
      @gen_cluster(
          client=True,
          nthreads=nthreads,
          config={
              "distributed.scheduler.work-stealing": False,
              "distributed.scheduler.worker-saturation": float("inf"),
          },
      )
      async def test_decide_worker_coschedule_order_neighbors_(
          c, s, *workers
      ):
          r"""
          Ensure that sibling root tasks are scheduled to the same node, reducing future
          data transfer.
  
          We generate a wide layer of "root" tasks (random NumPy arrays). All of those
          tasks share 0-5 trivial dependencies. The ``ndeps=0`` and ``ndeps=1`` cases are
          most common in real-world use (``ndeps=1`` is basically ``da.from_array(...,
          inline_array=False)`` or ``da.from_zarr``). The graph is structured like this
          (though the number of tasks and workers is different):
  
              |-W1-|  |-W2-| |-W3-|  |-W4-|   < ---- ideal task scheduling
  
                q       r       s       t      < --- `sum-aggregate-`
               / \     / \     / \     / \
              i   j   k   l   m   n   o   p    < --- `sum-`
              |   |   |   |   |   |   |   |
              a   b   c   d   e   f   g   h    < --- `random-`
              \   \   \   |   |   /   /   /
                     TRIVIAL * 0..5
  
          Neighboring `random-` tasks should be scheduled on the same worker. We test that
          generally, only one worker holds each row of the array, that the `random-` tasks
          are never transferred, and that there are few transfers overall.
          """
          da = pytest.importorskip("dask.array")
          np = pytest.importorskip("numpy")
  
          if ndeps == 0:
              x = da.random.random((100, 100), chunks=(10, 10))
              random_keys = set(flatten(x.__dask_keys__()))
              trivial_deps = {}
          else:
  
              def random(**kwargs):
                  assert len(kwargs) == ndeps
                  return np.random.random((10, 10))
  
              trivial_deps = {
                  f"k{i}": delayed(object(), name=f"object-{i}")
                  for i in range(ndeps)
              }
  
              # TODO is there a simpler (non-blockwise) way to make this sort of graph?
              x = da.blockwise(
                  random,
                  "yx",
                  new_axes={"y": (10,) * 10, "x": (10,) * 10},
                  dtype=float,
                  **trivial_deps,
              )
              random_keys = set(flatten(x.__dask_keys__()))
  
          xx, xsum = dask.persist(x, x.sum(axis=1, split_every=20))
          await xsum
  
          # Check that each chunk-row of the array is (mostly) stored on the same worker
          primary_worker_key_fractions = []
          secondary_worker_key_fractions = []
          for keys in x.__dask_keys__():
              # Iterate along rows of the array.
              keys = set(keys)
  
              # No more than 2 workers should have any keys
              assert (
                  sum(any(k in w.data for k in keys) for w in workers)
                  <= 2
              )
  
              # What fraction of the keys for this row does each worker hold?
              key_fractions = [
                  len(set(w.data).intersection(keys)) / len(keys)
                  for w in workers
              ]
              key_fractions.sort()
              # Primary worker: holds the highest percentage of keys
              # Secondary worker: holds the second highest percentage of keys
              primary_worker_key_fractions.append(key_fractions[-1])
              secondary_worker_key_fractions.append(key_fractions[-2])
  
          # There may be one or two rows that were poorly split across workers,
          # but the vast majority of rows should only be on one worker.
          assert np.mean(primary_worker_key_fractions) >= 0.9
          assert np.median(primary_worker_key_fractions) == 1.0
          assert np.mean(secondary_worker_key_fractions) <= 0.1
          assert np.median(secondary_worker_key_fractions) == 0.0
  
          # Check that there were few transfers
          unexpected_transfers = []
          for worker in workers:
              for log in worker.transfer_incoming_log:
                  keys = log["keys"]
                  # The root-ish tasks should never be transferred
                  assert not random_keys.intersection(keys)
                  # `object-` keys (the trivial deps of the root random tasks) should be
                  # transferred
                  if any(
                      (
                          not isinstance(k, str)
                          or not k.startswith("object")
                      )
                      for k in keys
                  ):
                      # But not many other things should be
                      unexpected_transfers.append(list(keys))
  
          # A transfer at the very end to move aggregated results is fine (necessary with
          # unbalanced workers in fact), but generally there should be very very few
          # transfers.
          assert len(unexpected_transfers) <= 3, unexpected_transfers
  
      test_decide_worker_coschedule_order_neighbors_()
  
imports: from __future__ import annotations
  import asyncio
  import itertools
  import json
  import logging
  import math
  import operator
  import pickle
  import random
  import re
  import sys
  from collections.abc import Collection
  from itertools import product
  from textwrap import dedent
  from time import sleep
  from unittest import mock
  import cloudpickle
  import psutil
  import pytest
  from tlz import concat, first, merge
  from tornado.ioloop import IOLoop
  import dask
  from dask import bag, delayed
  from dask.core import flatten
  from dask.highlevelgraph import HighLevelGraph, MaterializedLayer
  from dask.utils import parse_timedelta, tmpfile, typename
  from distributed import CancelledError, Client, Event, Future, Lock, Nanny, SchedulerPlugin, Worker, fire_and_forget, wait
  from distributed.comm.addressing import parse_host_port
  from distributed.compatibility import LINUX, MACOS, WINDOWS, PeriodicCallback
  from distributed.core import ConnectionPool, Status, clean_exception, connect, rpc
  from distributed.metrics import time
  from distributed.protocol import serialize
  from distributed.protocol.pickle import dumps, loads
  from distributed.protocol.serialize import ToPickle
  from distributed.scheduler import KilledWorker, MemoryState, Scheduler, WorkerState
  from distributed.utils import TimeoutError, wait_for
  from distributed.utils_test import NO_AMM, BlockedGatherDep, BlockedGetData, BrokenComm, NoSchedulerDelayWorker, assert_story, async_poll_for, captured_handler, captured_logger, cluster, dec, div, freeze_batched_send, freeze_data_fetching, gen_cluster, gen_test, inc, nodebug, padded_time, raises_with_cause, slowadd, slowdec, slowidentity, slowinc, tls_only_security, varying, wait_for_state
  from distributed.worker import dumps_function, secede
  from distributed import WorkerPlugin
label_desc: Use image node_id 1 for calling a global function with example usage: test_decide_worker_coschedule_order_neighbors(ndeps, nthreads) and returns: np