Fix inconsistent hashing for Nanny-spawned workers (#8400)

distributed/nanny.py

@@ -233,6 +233,13 @@ def __init__( # type: ignore[no-untyped-def]
  self.Worker = Worker if worker_class is None else worker_class
 
  self.pre_spawn_env = _get_env_variables("distributed.nanny.pre-spawn-environ")
+ # To get consistent hashing on subprocesses, we need to set a consistent seed for
+ # the Python hash algorithm; xref https://github.com/dask/distributed/pull/8400
+ if self.pre_spawn_env.get("PYTHONHASHSEED") in (None, "0"):
+ # This number is arbitrary; it was chosen to commemorate
+ # https://github.com/dask/dask/issues/6640.
+ self.pre_spawn_env.update({"PYTHONHASHSEED": "6640"})
+
  self.env = merge(
  self.pre_spawn_env,
  _get_env_variables("distributed.nanny.environ"),

distributed/tests/test_dask_collections.py

@@ -13,6 +13,7 @@
 import dask.dataframe as dd
 
 from distributed.client import wait
+from distributed.nanny import Nanny
 from distributed.utils_test import gen_cluster
 
 dfs = [
@@ -124,6 +125,18 @@ async def test_bag_groupby_tasks_default(c, s, a, b):
  assert not any("partd" in k[0] for k in b2.dask)
 
 
+@gen_cluster(client=True, Worker=Nanny)
+async def test_bag_groupby_key_hashing(c, s, a, b):
+ # https://github.com/dask/distributed/issues/4141
+ dsk = {("x", 0): (range, 5), ("x", 1): (range, 5), ("x", 2): (range, 5)}
+ grouped = db.Bag(dsk, "x", 3).groupby(lambda x: "even" if x % 2 == 0 else "odd")
+ remote = c.compute(grouped)
+ result = await remote
+ assert len(result) == 2
+ assert ("odd", [1, 3] * 3) in result
+ assert ("even", [0, 2, 4] * 3) in result
+
+
 @pytest.mark.parametrize("wait", [wait, lambda x: None])
 def test_dataframe_set_index_sync(wait, client):
  df = dask.datasets.timeseries(

distributed/tests/test_nanny.py

@@ -340,10 +340,13 @@ async def test_environment_variable_config(c, s, monkeypatch):
  },
 )
 async def test_environment_variable_pre_post_spawn(c, s, n):
- assert n.env == {"PRE-SPAWN": "1", "POST-SPAWN": "2"}
+ assert n.env == {"PRE-SPAWN": "1", "POST-SPAWN": "2", "PYTHONHASHSEED": "6640"}
  results = await c.run(lambda: os.environ)
  assert results[n.worker_address]["PRE-SPAWN"] == "1"
  assert results[n.worker_address]["POST-SPAWN"] == "2"
+ # if unset in pre-spawn-environ config, PYTHONHASHSEED defaults to "6640" to ensure
+ # consistent hashing across workers; https://github.com/dask/distributed/issues/4141
+ assert results[n.worker_address]["PYTHONHASHSEED"] == "6640"
 
  del os.environ["PRE-SPAWN"]
  assert "POST-SPAWN" not in os.environ