fix(#592): forward S3/storage env vars to remote MPI ranks via -x

[russfellows]

Fixes #592 (Problem A — S3 credential forwarding).

The bug

OpenMPI does not propagate arbitrary env vars to remote hosts by
default. For a multi-host object-storage run, the launcher host sets
AWS_ENDPOINT_URL, AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY,
S3DLIO_FOLLOW_REDIRECTS, STORAGE_LIBRARY, and BUCKET — but only
DLIO_DROP_CACHES_TIMEOUT was being forwarded to remote ranks via
-x. Remote ranks had none of the S3 config and would hang or fail
with auth errors on their first S3 operation.

The fix

In DLIOBenchmark.generate_dlio_command(), after the existing
DLIO_DROP_CACHES_TIMEOUT forwarding, iterate os.environ (sorted
for determinism) and append -x VARNAME for every variable whose name
starts with AWS_ or S3DLIO_, or is exactly STORAGE_LIBRARY or
BUCKET.

for _v in sorted(os.environ):
    if (_v.startswith('AWS_') or _v.startswith('S3DLIO_')
            or _v in ('STORAGE_LIBRARY', 'BUCKET')):
        mpi_prefix += f" -x {_v}"

Tests

11 new tests in tests/unit/test_mpi_s3_env_forwarding.py. All call
the real generate_dlio_command() with generate_mpi_prefix_cmd
patched to a fixed prefix string.

TestS3EnvForwardingMissing (6 tests) — RED before fix, GREEN after:

• test_aws_endpoint_url_forwarded — S3 endpoint must reach remote ranks
• test_aws_access_key_id_forwarded — access key must reach remote ranks
• test_aws_secret_access_key_forwarded — secret key must reach remote ranks
• test_s3dlio_follow_redirects_forwarded — AIStore redirect flag must reach remote ranks
• test_storage_library_forwarded — storage library selection must reach remote ranks
• test_bucket_forwarded — bucket name must reach remote ranks

TestAbsentVarsNotForwarded (2 tests):

• Vars not in the launcher env must not produce spurious -x flags

TestDropCachesTimeoutStillForwarded (2 tests):

• Pre-existing -x DLIO_DROP_CACHES_TIMEOUT forwarding must not regress
• Both DLIO and S3 vars appear together when all are set

TestNonMpiRunsUnchanged (1 test):

• No -x flags when exec_type is not MPI

Verification

• uv run pytest tests/unit/test_mpi_s3_env_forwarding.py -v →
  7 FAILED (bug confirmed), then 11 passed after fix
• uv run pytest tests/unit -q → 2266 passed, no regressions
• On-cluster: multi-host object-storage run should no longer
  hang/auth-fail on remote ranks

Files

• mlpstorage_py/benchmarks/dlio.py — extend MPI env forwarding loop
• tests/unit/test_mpi_s3_env_forwarding.py — 11 new tests

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅