diff --git a/CLAUDE.md b/CLAUDE.md index 46a9382..0e0c5d1 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -376,6 +376,61 @@ torchtitan job where capture died ~10 s in, exactly when inductor spun up its compile workers. Regression tests: `tests/test_fd_capture.py::TestForkedChildCannotDisableParentCapture`. +### NCCL environment capture + +Distributed failures are almost always configuration failures, so the NCCL +environment is recorded automatically on every run — no user code change. What +is captured: + +| Matched | Examples | +|---|---| +| `NCCL_*` | `NCCL_DEBUG`, `NCCL_SOCKET_IFNAME`, `NCCL_IB_HCA`, `NCCL_ALGO` | +| `TORCH_NCCL_*` | `TORCH_NCCL_ASYNC_ERROR_HANDLING`, `TORCH_NCCL_BLOCKING_WAIT` | +| `FI_*`, `OFI_*` | `FI_PROVIDER`, `FI_EFA_USE_DEVICE_RDMA` (libfabric / aws-ofi-nccl) | +| `UCX_*` | `UCX_TLS`, `UCX_NET_DEVICES` | +| Exact keys | `TORCH_DISTRIBUTED_DEBUG`, `TORCH_CPP_LOG_LEVEL`, `TORCH_SHOW_CPP_STACKTRACES`, `GLOO_SOCKET_IFNAME` | + +Runs are selectable by what they ran with, since `systemMetadata.*` is an +accepted filter field: + +```python +import pluto.query as pq + +runs = pq.list_runs( + 'my-project', + filters={'systemMetadata.nccl.nccl_env.NCCL_ALGO': 'Tree'}, +) +``` + +The environment is read at `pluto.init()` time, so vars exported later — e.g. +by a launcher configuring NCCL right before `init_process_group()` — are not +recorded. (This section doubles as the user-facing description: the Docusaurus +site that would have carried it was removed in #145, and the live docs are +Mintlify, built outside this repo.) + +- `collect_nccl_env()` (`pluto/sys.py`) is the single source of truth: it scans + `NCCL_ENV_PREFIXES` (`NCCL_`, `TORCH_NCCL_`, `FI_`, `OFI_`, `UCX_`) plus the + exact `NCCL_ENV_KEYS`, sorts, and masks credential-looking keys with + `MASKED_VALUE`. Read at call time, never cached at import. +- **Two sinks, deliberately.** `System.get_nccl_info()` puts it in + `systemMetadata.nccl.nccl_env` at run create (queryable via the + `systemMetadata.` filter prefix), and `Op._log_nccl_env()` logs one line from + `Op.start()`. The log line is not redundant: resuming ranks go through + `/api/runs/resume`, whose payload carries **no** system info + (`make_compat_resume_v1`), so in multi-node runs the metadata only ever + describes the rank that created the run. The console line is per-rank, which + is where a misconfigured worker actually shows up. +- `_log_nccl_env()` is called from `start()`, i.e. after `setup_logger()` and + the sync manager are up, so the line is captured and uploaded. It is + suppressed under `mode == 'noop'` and `disable_system_metrics` (backfill), + swallows collection errors, and truncates at `NCCL_ENV_LOG_MAX_CHARS` — + hosts with large `FI_*`/`UCX_*` sets would otherwise push multi-KB lines into + the console stream. +- Helper naming avoids credential words (`_is_masked_key`, not + `_is_secret_key`) for the same CodeQL reason documented under "Auth failures" + above: `py/clear-text-logging-sensitive-data` classifies a call's result by + callee name, and this result *is* logged. + ### Network Filesystems (NFS/Lustre/SMB) and SQLite WAL The sync DB uses SQLite WAL mode (`pluto/sync/store.py`), which relies on POSIX diff --git a/docs-api/config.mdx b/docs-api/config.mdx index 314fb41..7e344bf 100644 --- a/docs-api/config.mdx +++ b/docs-api/config.mdx @@ -69,7 +69,7 @@ Collect CUDA and cuDNN version information from system and PyTorch. get_nccl_info() -> Dict[str, Any] ``` -Collect NCCL version information and NCCL environment variables. +Collect NCCL versions and the NCCL/fabric environment. #### `System.get_infiniband_info` diff --git a/pluto/op.py b/pluto/op.py index b0151c3..4f5d7d1 100644 --- a/pluto/op.py +++ b/pluto/op.py @@ -34,7 +34,7 @@ from .store import DataStore from .sync import SyncProcessManager from .sync.store import HEALTH_METRIC_KEYS -from .sys import System +from .sys import System, collect_nccl_env from .util import ( ANSI, deep_merge, @@ -54,6 +54,11 @@ # logged, regardless of how many distinct values the series has. STRING_SERIES_MAX_LEN = 200 +# Cap on the rendered NCCL-environment log line. Hosts with a large fabric +# configuration (many FI_*/UCX_* vars) would otherwise push a multi-KB line +# into the console stream; the full set is still sent in systemMetadata. +NCCL_ENV_LOG_MAX_CHARS = 2048 + def _is_distributed_environment() -> bool: """Check if running in a distributed (DDP/FSDP) environment.""" @@ -723,6 +728,12 @@ def start(self) -> None: # Print URL where users can view the run logger.info(f'{tag}: {self._view_run_message()}') + # NCCL/fabric configuration of this rank. It also rides along in + # systemMetadata, but the log line lands in the run's console output + # (uploaded like any other log), which is where anyone debugging a + # hung all-reduce or an unexpected TCP fallback looks first. + self._log_nccl_env() + # Register excepthook to detect unhandled exceptions and mark runs as FAILED _register_excepthook() @@ -735,6 +746,29 @@ def start(self) -> None: pluto.ops.append(self) pluto.log, pluto.alert, pluto.watch = self.log, self.alert, self.watch + def _log_nccl_env(self) -> None: + """Log the NCCL-relevant environment at run start. + + Emitted once per rank, after the logger and sync process are up so the + line is captured and uploaded with the rest of the console output. + Suppressed under ``disable_system_metrics`` (backfill/migration), where + this host's environment says nothing about the run being written. + """ + if self.settings.mode == 'noop' or self.settings.disable_system_metrics: + return + try: + nccl_env = collect_nccl_env() + except Exception as e: # never let diagnostics break a run + logger.debug('%s: NCCL environment capture skipped: %s', tag, e) + return + if not nccl_env: + logger.debug('%s: no NCCL environment variables set', tag) + return + rendered = ', '.join(f'{k}={v}' for k, v in nccl_env.items()) + if len(rendered) > NCCL_ENV_LOG_MAX_CHARS: + rendered = f'{rendered[:NCCL_ENV_LOG_MAX_CHARS]}... (truncated)' + logger.info('%s: NCCL environment (%d vars): %s', tag, len(nccl_env), rendered) + def log( self, data: Dict[str, Any], diff --git a/pluto/sys.py b/pluto/sys.py index ee9271c..0c9f3f3 100644 --- a/pluto/sys.py +++ b/pluto/sys.py @@ -18,6 +18,59 @@ logger = logging.getLogger(f'{__name__.split(".")[0]}') tag = 'System' +# Environment variables that shape NCCL behaviour. NCCL_*/TORCH_NCCL_* are +# NCCL's own knobs and PyTorch's ProcessGroupNCCL wrappers around them; the +# fabric prefixes cover the transport plugins NCCL dispatches to +# (libfabric/aws-ofi-nccl on EFA, UCX), which decide whether a job runs over +# IB/RoCE or silently falls back to TCP. +NCCL_ENV_PREFIXES = ('NCCL_', 'TORCH_NCCL_', 'FI_', 'OFI_', 'UCX_') + +# Vars outside those prefixes that still change collective behaviour, or are +# needed to make sense of NCCL's own debug output. +NCCL_ENV_KEYS = ( + 'TORCH_DISTRIBUTED_DEBUG', + 'TORCH_CPP_LOG_LEVEL', + 'TORCH_SHOW_CPP_STACKTRACES', + 'GLOO_SOCKET_IFNAME', +) + +# Substrings marking a value that must not leave the machine verbatim. NCCL_* +# is not where credentials normally live, but the prefix scan is broad enough +# that a stray NCCL_..._TOKEN would otherwise be shipped and logged as-is. +_MASKED_KEY_MARKERS = ( + 'TOKEN', + 'SECRET', + 'PASSWORD', + 'PASSWD', + 'API_KEY', + 'APIKEY', + 'CREDENTIAL', + 'PRIVATE_KEY', + 'ACCESS_KEY', +) +MASKED_VALUE = '' + + +def _is_masked_key(key: str) -> bool: + upper = key.upper() + return any(marker in upper for marker in _MASKED_KEY_MARKERS) + + +def collect_nccl_env( + environ: Optional[Mapping[str, str]] = None, +) -> Dict[str, str]: + """NCCL-relevant environment variables, sorted, credential values masked. + + Read at call time rather than cached at import, so a var exported between + interpreter start and ``pluto.init()`` is still picked up. + """ + env: Mapping[str, str] = os.environ if environ is None else environ + d: Dict[str, str] = {} + for key, value in env.items(): + if key.startswith(NCCL_ENV_PREFIXES) or key in NCCL_ENV_KEYS: + d[key] = MASKED_VALUE if _is_masked_key(key) else value + return dict(sorted(d.items())) + class System: def __init__(self, settings: Settings) -> None: @@ -468,7 +521,7 @@ def get_cuda_info(self) -> Dict[str, Any]: return d def get_nccl_info(self) -> Dict[str, Any]: - """Collect NCCL version information and NCCL environment variables.""" + """Collect NCCL versions and the NCCL/fabric environment.""" d: Dict[str, Any] = {} # PyTorch NCCL version @@ -502,11 +555,10 @@ def get_nccl_info(self) -> Dict[str, Any]: except Exception: pass - # NCCL environment variables - nccl_env: Dict[str, str] = {} - for key, value in os.environ.items(): - if key.startswith('NCCL_'): - nccl_env[key] = value + # NCCL environment variables (plus the fabric/plugin vars NCCL + # dispatches through) — shipped in systemMetadata on run create, so a + # job's collective configuration is queryable per run. + nccl_env = collect_nccl_env() if nccl_env: d['nccl_env'] = nccl_env diff --git a/tests/test_fork_e2e.py b/tests/test_fork_e2e.py index 05de806..478f22e 100644 --- a/tests/test_fork_e2e.py +++ b/tests/test_fork_e2e.py @@ -69,6 +69,14 @@ def _poll_metric_names( ) +def _metric_values(project: str, run_id: int, metric: str) -> List[float]: + """Values of a single metric series, in server order.""" + metrics = pq.get_metrics(project, run_id, metric_names=[metric]) + if hasattr(metrics, 'to_dict'): # pandas DataFrame + return list(metrics['value'].tolist()) + return [m['value'] for m in metrics] + + def _poll_max_step( project: str, run_id: int, @@ -277,11 +285,14 @@ def test_fork_e2e_log_metrics(parent_run): metric_names = _poll_metric_names(FORK_PROJECT, run_id, ['fork/loss']) assert 'fork/loss' in metric_names - metrics = pq.get_metrics(FORK_PROJECT, run_id, metric_names=['fork/loss']) - if hasattr(metrics, 'to_dict'): - values = metrics['value'].tolist() - else: - values = [m['value'] for m in metrics] + # The name is queryable as soon as the *first* point lands, so reading the + # series straight after that races ingest (seen in CI: len(values) == 1). + # Poll for the full series the way the parent fixture does with + # _poll_max_step; the assertions below still have to hold at the deadline. + values = _poll( + fn=lambda: _metric_values(FORK_PROJECT, run_id, 'fork/loss'), + check=lambda vals: len(vals) >= 5, + ) assert len(values) == 5 assert values[0] == pytest.approx(0.5, abs=1e-6) diff --git a/tests/test_nccl_env.py b/tests/test_nccl_env.py new file mode 100644 index 0000000..7a2ea5c --- /dev/null +++ b/tests/test_nccl_env.py @@ -0,0 +1,120 @@ +"""Unit tests for the startup NCCL environment log line. + +NCCL's configuration is the first thing anyone asks for when a distributed +job hangs, runs slow, or quietly falls back to TCP. Pluto ships it two ways +on startup: inside ``systemMetadata`` on run create (see +``tests/test_system_info.py``) and as a console line emitted by +``Op.start()``, which is captured and uploaded like any other log output. +These tests pin the log-line half. +""" + +from __future__ import annotations + +import atexit +import logging +import os +from unittest import mock + +import pluto +from pluto.op import NCCL_ENV_LOG_MAX_CHARS, Op +from pluto.sets import Settings +from pluto.sys import MASKED_VALUE + + +def _make_op(tmp_path) -> Op: + """An Op with no server/sync side effects, ready to log.""" + settings = Settings() + settings.mode = 'noop' # skips login/run-create in Op.__init__ + settings.dir = str(tmp_path) + settings.meta = [] # shadow the class-level shared list (test isolation) + os.makedirs(os.path.join(settings.get_dir(), 'files'), exist_ok=True) + op = Op(config={}, settings=settings) + # Past construction, act like a live run: 'noop' means nothing is sent + # anywhere, so the log line is intentionally suppressed for it. + op.settings.mode = 'online' + return op + + +def _messages(caplog) -> str: + return '\n'.join(r.getMessage() for r in caplog.records) + + +class TestLogNcclEnv: + def test_logs_nccl_env_at_info(self, tmp_path, caplog): + op = _make_op(tmp_path) + env = {'NCCL_DEBUG': 'INFO', 'NCCL_SOCKET_IFNAME': 'eth0'} + with mock.patch.dict(os.environ, env, clear=False): + with caplog.at_level(logging.INFO, logger='pluto'): + op._log_nccl_env() + out = _messages(caplog) + assert 'NCCL_DEBUG=INFO' in out + assert 'NCCL_SOCKET_IFNAME=eth0' in out + + def test_nothing_logged_when_no_nccl_vars_set(self, tmp_path, caplog): + op = _make_op(tmp_path) + with mock.patch('pluto.op.collect_nccl_env', return_value={}): + with caplog.at_level(logging.INFO, logger='pluto'): + op._log_nccl_env() + assert 'NCCL environment' not in _messages(caplog) + + def test_credentials_are_masked_in_the_log_line(self, tmp_path, caplog): + op = _make_op(tmp_path) + with mock.patch.dict(os.environ, {'NCCL_AUTH_TOKEN': 'hunter2'}): + with caplog.at_level(logging.INFO, logger='pluto'): + op._log_nccl_env() + out = _messages(caplog) + assert 'hunter2' not in out + assert MASKED_VALUE in out + + def test_long_env_is_truncated(self, tmp_path, caplog): + op = _make_op(tmp_path) + env = {f'NCCL_PAD_{i:03d}': 'x' * 64 for i in range(100)} + with mock.patch('pluto.op.collect_nccl_env', return_value=env): + with caplog.at_level(logging.INFO, logger='pluto'): + op._log_nccl_env() + out = _messages(caplog) + assert '(truncated)' in out + # The cap bounds the rendered vars, not the whole formatted message. + assert len(out) < NCCL_ENV_LOG_MAX_CHARS + 200 + + def test_suppressed_under_disable_system_metrics(self, tmp_path, caplog): + """Backfill/migration: this host's env says nothing about the run.""" + op = _make_op(tmp_path) + op.settings.disable_system_metrics = True + with mock.patch.dict(os.environ, {'NCCL_DEBUG': 'INFO'}): + with caplog.at_level(logging.INFO, logger='pluto'): + op._log_nccl_env() + assert 'NCCL environment' not in _messages(caplog) + + def test_suppressed_in_noop_mode(self, tmp_path, caplog): + op = _make_op(tmp_path) + op.settings.mode = 'noop' + with mock.patch.dict(os.environ, {'NCCL_DEBUG': 'INFO'}): + with caplog.at_level(logging.INFO, logger='pluto'): + op._log_nccl_env() + assert 'NCCL environment' not in _messages(caplog) + + def test_collection_failure_never_raises(self, tmp_path, caplog): + op = _make_op(tmp_path) + with mock.patch('pluto.op.collect_nccl_env', side_effect=RuntimeError('boom')): + with caplog.at_level(logging.INFO, logger='pluto'): + op._log_nccl_env() # must not propagate + assert 'NCCL environment (' not in _messages(caplog) + + def test_start_emits_the_line(self, tmp_path, caplog): + """The hook is wired into Op.start(), not just callable on its own.""" + op = _make_op(tmp_path) + op._monitor = mock.MagicMock() + op._sync_manager = None + op._iface = None + # start() publishes module-level globals and appends to pluto.ops; put + # them back so a later test in this worker sees an untouched module. + saved = (pluto.ops, pluto.log, pluto.alert, pluto.watch) + try: + with mock.patch.dict(os.environ, {'NCCL_DEBUG': 'INFO'}, clear=False): + with caplog.at_level(logging.INFO, logger='pluto'): + op.start() + finally: + atexit.unregister(op.finish) + pluto.ops, pluto.log, pluto.alert, pluto.watch = saved + assert 'NCCL_DEBUG=INFO' in _messages(caplog) diff --git a/tests/test_system_info.py b/tests/test_system_info.py index f4c5b8d..daaf2b3 100644 --- a/tests/test_system_info.py +++ b/tests/test_system_info.py @@ -7,7 +7,13 @@ import pytest -from pluto.sys import System +from pluto.sys import ( + MASKED_VALUE, + NCCL_ENV_KEYS, + NCCL_ENV_PREFIXES, + System, + collect_nccl_env, +) class TestSystemInfoHelper: @@ -194,10 +200,34 @@ def test_collects_nccl_env_vars(self): assert result['nccl_env']['NCCL_SOCKET_IFNAME'] == 'eth0' assert result['nccl_env']['NCCL_IB_DISABLE'] == '0' + def test_collects_torch_nccl_and_fabric_env_vars(self): + """TORCH_NCCL_*/FI_*/UCX_* and the extra distributed vars ride along.""" + env = { + 'TORCH_NCCL_ASYNC_ERROR_HANDLING': '1', + 'TORCH_NCCL_BLOCKING_WAIT': '0', + 'FI_PROVIDER': 'efa', + 'FI_EFA_USE_DEVICE_RDMA': '1', + 'UCX_TLS': 'rc,cuda_copy', + 'TORCH_DISTRIBUTED_DEBUG': 'DETAIL', + 'GLOO_SOCKET_IFNAME': 'eth0', + } + with patch.dict(os.environ, env, clear=False): + sys_obj = self._make_system() + nccl_env = sys_obj.get_nccl_info()['nccl_env'] + for key, value in env.items(): + assert nccl_env[key] == value + + def test_unrelated_env_vars_not_collected(self): + """The scan stays scoped — no wholesale environ dump.""" + env = {'MY_SECRET_TRAINING_FLAG': 'nope', 'PATH_TO_NCCL': 'also-nope'} + with patch.dict(os.environ, env, clear=False): + nccl_env = collect_nccl_env() + assert 'MY_SECRET_TRAINING_FLAG' not in nccl_env + assert 'PATH_TO_NCCL' not in nccl_env + def test_no_nccl_env_when_none_set(self): - """nccl_env key absent when no NCCL_* vars exist.""" - env = {k: v for k, v in os.environ.items() if not k.startswith('NCCL_')} - with patch.dict(os.environ, env, clear=True): + """nccl_env key absent when no NCCL-relevant vars exist.""" + with patch.dict(os.environ, _env_without_nccl(), clear=True): sys_obj = self._make_system() result = sys_obj.get_nccl_info() assert 'nccl_env' not in result @@ -209,6 +239,66 @@ def test_get_info_includes_nccl_when_available(self): if result: assert 'nccl' in info + def test_get_info_carries_nccl_env(self): + """The env vars reach the systemMetadata payload sent on run create.""" + with patch.dict(os.environ, {'NCCL_ALGO': 'Ring'}, clear=False): + sys_obj = self._make_system() + info = sys_obj.get_info() + assert info['nccl']['nccl_env']['NCCL_ALGO'] == 'Ring' + + +def _env_without_nccl(): + """Current environ minus everything collect_nccl_env() would pick up.""" + return { + k: v + for k, v in os.environ.items() + if not k.startswith(NCCL_ENV_PREFIXES) and k not in NCCL_ENV_KEYS + } + + +class TestCollectNcclEnv: + """Tests for the standalone collect_nccl_env() helper.""" + + def test_reads_os_environ_by_default(self): + with patch.dict(os.environ, {'NCCL_DEBUG': 'WARN'}, clear=False): + assert collect_nccl_env()['NCCL_DEBUG'] == 'WARN' + + def test_accepts_an_explicit_mapping(self): + assert collect_nccl_env({'NCCL_DEBUG': 'INFO', 'HOME': '/root'}) == { + 'NCCL_DEBUG': 'INFO' + } + + def test_keys_are_sorted(self): + env = {'NCCL_SOCKET_IFNAME': 'eth0', 'FI_PROVIDER': 'efa', 'NCCL_ALGO': 'Tree'} + assert list(collect_nccl_env(env)) == [ + 'FI_PROVIDER', + 'NCCL_ALGO', + 'NCCL_SOCKET_IFNAME', + ] + + @pytest.mark.parametrize( + 'key', + [ + 'NCCL_AUTH_TOKEN', + 'FI_EFA_SECRET', + 'NCCL_NET_PLUGIN_PASSWORD', + 'UCX_API_KEY', + 'NCCL_ACCESS_KEY_ID', + ], + ) + def test_credential_bearing_values_are_masked(self, key): + result = collect_nccl_env({key: 'hunter2'}) + assert result[key] == MASKED_VALUE + assert 'hunter2' not in str(result) + + def test_masking_is_case_insensitive(self): + assert collect_nccl_env({'NCCL_auth_token': 'x'})['NCCL_auth_token'] == ( + MASKED_VALUE + ) + + def test_empty_when_nothing_relevant_is_set(self): + assert collect_nccl_env({'PATH': '/usr/bin', 'HOME': '/root'}) == {} + class TestGetInfinibandInfo(TestSystemInfoHelper): """Tests for System.get_infiniband_info() method."""