From 3fb3d8c8c14800ff8f1cd76d0052985f355ecfcc Mon Sep 17 00:00:00 2001 From: Jacob Tomlinson Date: Thu, 17 Dec 2020 09:46:11 +0000 Subject: [PATCH] Dask internal inherit config (#4364) * Change update method to ensure config is applied * Add test for serialization methods * Remove note * Implement DASK_INTERNAL_INHERIT_CONFIG for dask-spec --- distributed/cli/dask_scheduler.py | 3 +-- distributed/cli/dask_spec.py | 8 ++++++++ distributed/cli/dask_worker.py | 3 +-- distributed/tests/test_utils.py | 18 ++++++++++++++++++ 4 files changed, 28 insertions(+), 4 deletions(-) diff --git a/distributed/cli/dask_scheduler.py b/distributed/cli/dask_scheduler.py index 45d4f484c8..335fcd63b5 100755 --- a/distributed/cli/dask_scheduler.py +++ b/distributed/cli/dask_scheduler.py @@ -169,8 +169,7 @@ def main( if "DASK_INTERNAL_INHERIT_CONFIG" in os.environ: config = deserialize_for_cli(os.environ["DASK_INTERNAL_INHERIT_CONFIG"]) - # Update the global config given priority to the existing global config - dask.config.update(dask.config.global_config, config, priority="old") + dask.config.update(dask.config.global_config, config) if not host and (tls_ca_file or tls_cert or tls_key): host = "tls://" diff --git a/distributed/cli/dask_spec.py b/distributed/cli/dask_spec.py index 0a224e5b37..299878a3a4 100644 --- a/distributed/cli/dask_spec.py +++ b/distributed/cli/dask_spec.py @@ -1,10 +1,13 @@ import asyncio import click import json +import os import sys import yaml +import dask.config from distributed.deploy.spec import run_spec +from distributed.utils import deserialize_for_cli @click.command(context_settings=dict(ignore_unknown_options=True)) @@ -13,6 +16,11 @@ @click.option("--spec-file", type=str, default=None, help="") @click.version_option() def main(args, spec: str, spec_file: str): + + if "DASK_INTERNAL_INHERIT_CONFIG" in os.environ: + config = deserialize_for_cli(os.environ["DASK_INTERNAL_INHERIT_CONFIG"]) + dask.config.update(dask.config.global_config, config) + if spec and spec_file or not spec and not spec_file: print("Must specify exactly one of --spec and --spec-file") sys.exit(1) diff --git a/distributed/cli/dask_worker.py b/distributed/cli/dask_worker.py index b3b546ea3c..140d0d1ab8 100755 --- a/distributed/cli/dask_worker.py +++ b/distributed/cli/dask_worker.py @@ -397,8 +397,7 @@ def del_pid_file(): if "DASK_INTERNAL_INHERIT_CONFIG" in os.environ: config = deserialize_for_cli(os.environ["DASK_INTERNAL_INHERIT_CONFIG"]) - # Update the global config given priority to the existing global config - dask.config.update(dask.config.global_config, config, priority="old") + dask.config.update(dask.config.global_config, config) nannies = [ t( diff --git a/distributed/tests/test_utils.py b/distributed/tests/test_utils.py index 01332f3883..022e992544 100644 --- a/distributed/tests/test_utils.py +++ b/distributed/tests/test_utils.py @@ -45,6 +45,8 @@ LRU, offload, TimeoutError, + deserialize_for_cli, + serialize_for_cli, ) from distributed.utils_test import loop, loop_in_thread # noqa: F401 from distributed.utils_test import div, has_ipv6, inc, throws, gen_test, captured_logger @@ -606,3 +608,19 @@ def test_lru(): async def test_offload(): assert (await offload(inc, 1)) == 2 assert (await offload(lambda x, y: x + y, 1, y=2)) == 3 + + +def test_cli_serialization(): + # Use context manager without changing the value to ensure test side effects are restored + with dask.config.set( + { + "distributed.comm.default-scheme": dask.config.get( + "distributed.comm.default-scheme" + ) + } + ): + config = deserialize_for_cli( + serialize_for_cli({"distributed": {"comm": {"default-scheme": "tls"}}}) + ) # Take a round trip through the serialization + dask.config.update(dask.config.global_config, config) + assert dask.config.get("distributed.comm.default-scheme") == "tls"