Introduce distributed.comm.ucx.environment config slot #7164

wence- · 2022-10-20T10:47:46Z

Can be used for setting arbitrary UCX configuration in addition to the specific high-level options in those cases where one needs fine-grained control over the UCX configuration.

Tests added / passed
Passes pre-commit run --all-files

cc @pentschev, @quasiben

Can be used for setting arbitrary UCX configuration in addition to the specific high-level options in those cases where one needs fine-grained control over the UCX configuration.

distributed/comm/ucx.py

This means that we don't accidentally skip validation for names in two parts of the hierarchy that happen to collide with a skipped name.

distributed/comm/ucx.py

github-actions · 2022-10-20T12:29:50Z

Unit Test Results

See test report for an extended history of previous test failures. This is useful for diagnosing flaky tests.

      15 files ±0       15 suites ±0 6h 24m 54s ⏱️ ±0s
  3 150 tests ±0   3 060 ✔️ +2   85 💤 - 1   5 ❌ - 1
23 307 runs ±0 22 372 ✔️ ±0 915 💤 +1 20 ❌ - 1

For more details on these failures, see this check.

Results for commit f653811. ± Comparison against base commit 6afce9c.

♻️ This comment has been updated with latest results.

We cannot control (say) transport-level (uct) options when calling ucp.init, but if the relevant UCX environment variable lives in the environment at the time we make the call, then it will be respected. So take any user-provided low-level settings and temporarily insert into the environment. Precedence (highest to lowest): 1. Externally specified environment overrides 2. High-level distributed ucx settings 3. Low-level distributed ucx.environment settings

jrbourbeau · 2022-10-20T19:46:49Z

distributed/tests/test_config.py

-        "bokeh-application",
-        "environ",
-        "pre-spawn-environ",
+        "distributed.scheduler.default-task-durations",


Why are these changes needed?

The schema validation in this test traverses the full yaml schema from its root to any leaves. Normally, any object property (mapping to a dict in python) is required to have properties set (indicating which keys are valid in the dict). In some cases (for example the nanny environment) the keys can be arbitrary, so we skip validation.

Previously this skipping was just done on the leaf key name (disregarding the path to the leaf), so could have resulted in false positive skipping (e.g. previously any key name environ would not have had its properties validated even in the case where they existed).

To avoid this, I'm changing the skip list to specify the fully-qualified key names, and remembering the path to the leaf when traversing.

pentschev

Thanks @wence- , overall this looks great. I've left a few comments that we probably need to change.

pentschev · 2022-10-20T16:45:09Z

distributed/comm/ucx.py

+    with patch.dict(os.environ, ucx_environment):
+        # We carefully ensure that ucx_environment only contains things
+        # that don't override ucx_config or existing slots in the
+        # environment, so the user's external environment can safely
+        # override things here.
+        ucp.init(options=ucx_config, env_takes_precedence=True)


Nice handling!

distributed/comm/ucx.py

So that users have some useful information if they specifying overlapping configuration options in both the high level ucx configuration and directly by setting the environment, report when the ucx.environment is overridden by either high level options or the external environment.

This checking was a leftover from dask#3515 when users could set arbitrary options in the top-level ucx configuration slots. Now we construct the high level options ourselves and can therefore be sure that the configuration names we will pass to UCX are valid. Conversely, we cannot check the names of the arbitrary ucx.environment options because ucp.get_config() will only return UCP-level configuration variables (and not transport-level UCT options).

distributed/comm/ucx.py

This is not complete since UCX_TLS=all may enable CUDA transports, but we can certainly catch user-specified settings.

pentschev

LGTM, thanks so much for all the work here @wence- !

quasiben · 2022-10-21T18:57:46Z

Generally looks good. @pentschev @wence-, now that TLS=ALL generally works what do you think about adding that as a control flag to the standard UCX config here as well ? So instead of having to set all the config options: nvlink, cuda-copy, infiniband, ... we can (just like with dask-cuda), turn it all on. Would that be too confusing ?

wence- · 2022-10-21T19:01:07Z

Generally looks good. @pentschev @wence-, now that TLS=ALL generally works what do you think about adding that as a control flag to the standard UCX config here as well ? So instead of having to set all the config options: nvlink, cuda-copy, infiniband, ... we can (just like with dask-cuda), turn it all on. Would that be too confusing ?

If no UCX_TLS is specified, that implies all, so the current code works I think.

pentschev · 2022-10-21T19:07:36Z

So instead of having to set all the config options: nvlink, cuda-copy, infiniband, ... we can (just like with dask-cuda), turn it all on. Would that be too confusing ?

The only purpose for those flags today is to have a user-friendly way of explicitly enabling/disabling transports that we generally care about. One could just as well set distributed.comm.ucx.environment.tls=tcp,cuda_copy,cuda_ipc for example (equivalent to enable_tcp_over_ucx=True enable_nvlink=True enable_infiniband=False). Generally I think we want to make things more widely available to users, but as a developer I wouldn't mind removing them, as I can just set UCX_TLS specifically to what I need, but that is unlikely to be true for every user, which would require a more detailed understanding at what each transport does.

quasiben · 2022-10-24T17:05:28Z

I'm rerunning the failed jobs. Many seemed to have failed test_connection_made_with_extra_conn_args and I'm being a bit cautious. I'll plan to merge after CI finishes

quasiben · 2022-10-24T19:37:46Z

Looks like the failures are a known issue:
#7168

pentschev · 2022-10-24T20:02:44Z

Thanks @wence- for working on this, and @quasiben for merging! 😄

Introduce distributed.comm.ucx.environment config slot

fd028fb

Can be used for setting arbitrary UCX configuration in addition to the specific high-level options in those cases where one needs fine-grained control over the UCX configuration.

wence- commented Oct 20, 2022

View reviewed changes

distributed/comm/ucx.py Show resolved Hide resolved

wence- commented Oct 20, 2022

View reviewed changes

distributed/comm/ucx.py Outdated Show resolved Hide resolved

wence- mentioned this pull request Oct 20, 2022

Configuration Profiles rapidsai/dask-cuda#1020

Open

Use fully qualified config names in schema validation

4302984

This means that we don't accidentally skip validation for names in two parts of the hierarchy that happen to collide with a skipped name.

pentschev reviewed Oct 20, 2022

View reviewed changes

distributed/comm/ucx.py Outdated Show resolved Hide resolved

distributed/comm/ucx.py Outdated Show resolved Hide resolved

wence- added 2 commits October 20, 2022 17:13

Adapt tests to new _prepare_ucx_config calling convention

5cb69a2

wence- force-pushed the wence/feature/ucx-environment branch from f8e83a6 to 217cb29 Compare October 20, 2022 16:13

Test config setting precedence

6ca877c

wence- force-pushed the wence/feature/ucx-environment branch from 217cb29 to 6ca877c Compare October 20, 2022 16:27

jrbourbeau reviewed Oct 20, 2022

View reviewed changes

pentschev suggested changes Oct 20, 2022

View reviewed changes

wence- added 3 commits October 21, 2022 10:48

Be more careful with ucx_config tests

a276bfe

ucx: Use logger.warning rather than warnings.warn

e01c9e0

wence- force-pushed the wence/feature/ucx-environment branch from 97a3cb4 to b9c4bcd Compare October 21, 2022 10:10

wence- commented Oct 21, 2022

View reviewed changes

distributed/comm/ucx.py Show resolved Hide resolved

ucx: check if CUDA transport is requested more exhaustively

f653811

This is not complete since UCX_TLS=all may enable CUDA transports, but we can certainly catch user-specified settings.

pentschev approved these changes Oct 21, 2022

View reviewed changes

quasiben merged commit ecf5e93 into dask:main Oct 24, 2022

wence- deleted the wence/feature/ucx-environment branch November 2, 2022 15:21

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Introduce distributed.comm.ucx.environment config slot #7164

Introduce distributed.comm.ucx.environment config slot #7164

wence- commented Oct 20, 2022

github-actions bot commented Oct 20, 2022 •

edited

Loading

jrbourbeau Oct 20, 2022

wence- Oct 21, 2022 •

edited

Loading

pentschev left a comment

pentschev Oct 20, 2022

pentschev left a comment

quasiben commented Oct 21, 2022

wence- commented Oct 21, 2022

pentschev commented Oct 21, 2022

quasiben commented Oct 24, 2022

quasiben commented Oct 24, 2022

pentschev commented Oct 24, 2022

Introduce distributed.comm.ucx.environment config slot #7164

Introduce distributed.comm.ucx.environment config slot #7164

Conversation

wence- commented Oct 20, 2022

github-actions bot commented Oct 20, 2022 • edited Loading

Unit Test Results

jrbourbeau Oct 20, 2022

Choose a reason for hiding this comment

wence- Oct 21, 2022 • edited Loading

Choose a reason for hiding this comment

pentschev left a comment

Choose a reason for hiding this comment

pentschev Oct 20, 2022

Choose a reason for hiding this comment

pentschev left a comment

Choose a reason for hiding this comment

quasiben commented Oct 21, 2022

wence- commented Oct 21, 2022

pentschev commented Oct 21, 2022

quasiben commented Oct 24, 2022

quasiben commented Oct 24, 2022

pentschev commented Oct 24, 2022

github-actions bot commented Oct 20, 2022 •

edited

Loading

wence- Oct 21, 2022 •

edited

Loading