PyTorch Distributed Elastic Launch Segmentation Fault with Python 3.12

[PaulZhang12]

🐛 Describe the bug

With Python 3.12, using torch.distributed elastic_launch results in segmentation fault. Python 3.11 with the same code works.

import torch.distributed.launcher as pet
import uuid
import tempfile
import os


def get_launch_config(world_size: int, rdzv_endpoint: str) -> pet.LaunchConfig:
    return pet.LaunchConfig(
        min_nodes=1,
        max_nodes=1,
        nproc_per_node=world_size,
        run_id=str(uuid.uuid4()),
        rdzv_backend="c10d",
        rdzv_endpoint=rdzv_endpoint,
        rdzv_configs={"store_type": "file"},
        start_method="spawn",
        monitor_interval=1,
        max_restarts=0,
    )

def entrypoint():
    print("Here!")


if __name__ == '__main__':
    with tempfile.TemporaryDirectory() as tmpdir:
        launch_config = get_launch_config(2, os.path.join(tmpdir, "rdzv"))
        pet.elastic_launch(launch_config, entrypoint=entrypoint)()

Versions

[conda] numpy 1.26.2 pypi_0 pypi
[conda] torch 2.3.0.dev20231218+cu121 pypi_0 pypi
[conda] torchmetrics 1.0.3 pypi_0 pypi
[conda] torchrec 0.5.0.dev20231218+cu121 pypi_0 pypi

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225 @dzhulgakov

[malfet]

Can you attach a backtrace please?

[PaulZhang12]

Running the script does not give a stack trace. This simple repro is based off of torchrec metric tests that fail with Python 3.12, to which there is a stack trace with the seg fault:

Also attaching the backtrace obtained when running the script with gdb:

[XilunWu]

@kurman The stacktrace shows segfault is raised when calling store.compare_set(). Not sure if it has anything to do with my libuv PR #116141 but our tests showed no failure IIRC. Not necessarily issue on our side but it's a heads-up.

[kurman]

Not sure if it has anything to do with my libuv PR

@XilunWu I believe useLibUV false by default so this should not be a root cause. I am testing this now but I take that we should triage this under general python 3.12 support.

[XilunWu]

Can confirm that TCPStore has issue under Python 3.12.

Test steps:

1. checkout the latest pytorch main branch
2. build from source with a python 3.12 conda environment
3. run pytest test/distributed/test_store.py

Output:

================================================================================================================= test session starts =================================================================================================================
platform linux -- Python 3.12.2, pytest-7.4.0, pluggy-1.0.0
rootdir: /data/users/xilunwu/oss/pytorch
configfile: pytest.ini
plugins: hypothesis-6.99.11
collected 85 items                                                                                                                                                                                                                                    
Running 85 items in this shard

test/distributed/test_store.py sFatal Python error: Segmentation fault

Current thread 0x00007f741b395740 (most recent call first):
  File "/data/users/xilunwu/oss/pytorch/test/distributed/test_store.py", line 105 in _test_compare_set
  File "/data/users/xilunwu/oss/pytorch/test/distributed/test_store.py", line 121 in test_compare_set
  File "/data/users/xilunwu/oss/pytorch/torch/testing/_internal/common_utils.py", line 2739 in wrapper
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/unittest/case.py", line 589 in _callTestMethod
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/unittest/case.py", line 634 in run
  File "/data/users/xilunwu/oss/pytorch/torch/testing/_internal/common_utils.py", line 2838 in _run_custom
  File "/data/users/xilunwu/oss/pytorch/torch/testing/_internal/common_utils.py", line 2866 in run
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/unittest/case.py", line 690 in __call__
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/_pytest/unittest.py", line 333 in runtest
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/_pytest/runner.py", line 169 in pytest_runtest_call
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/pluggy/_callers.py", line 39 in _multicall
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/pluggy/_manager.py", line 80 in _hookexec
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/pluggy/_hooks.py", line 265 in __call__
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/_pytest/runner.py", line 262 in <lambda>
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/_pytest/runner.py", line 341 in from_call
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/_pytest/runner.py", line 261 in call_runtest_hook
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/_pytest/runner.py", line 222 in call_and_report
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/_pytest/runner.py", line 133 in runtestprotocol
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/_pytest/runner.py", line 114 in pytest_runtest_protocol
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/pluggy/_callers.py", line 39 in _multicall
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/pluggy/_manager.py", line 80 in _hookexec
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/pluggy/_hooks.py", line 265 in __call__
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/_pytest/main.py", line 349 in pytest_runtestloop
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/pluggy/_callers.py", line 39 in _multicall
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/pluggy/_manager.py", line 80 in _hookexec
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/pluggy/_hooks.py", line 265 in __call__
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/_pytest/main.py", line 324 in _main
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/_pytest/main.py", line 270 in wrap_session
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/_pytest/main.py", line 317 in pytest_cmdline_main
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/pluggy/_callers.py", line 39 in _multicall
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/pluggy/_manager.py", line 80 in _hookexec
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/pluggy/_hooks.py", line 265 in __call__
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/_pytest/config/__init__.py", line 166 in main
  File "/home/xilunwu/local/miniconda3/envs/py312/lib/python3.12/site-packages/_pytest/config/__init__.py", line 189 in console_main
  File "/home/xilunwu/local/miniconda3/envs/py312/bin/pytest", line 11 in <module>

Extension modules: numpy.core._multiarray_umath, numpy.core._multiarray_tests, numpy.linalg._umath_linalg, numpy.fft._pocketfft_internal, numpy.random._common, numpy.random.bit_generator, numpy.random._bounded_integers, numpy.random._mt19937, numpy.random.mtrand, numpy.random._philox, numpy.random._pcg64, numpy.random._sfc64, numpy.random._generator, torch._C, torch._C._fft, torch._C._linalg, torch._C._nested, torch._C._nn, torch._C._sparse, torch._C._special, psutil._psutil_linux, psutil._psutil_posix (total: 22)
Segmentation fault (core dumped)

cc @wconstab @kwen2501 @kurman

[wconstab]

Did you narrow this down to a specific root cause? @XilunWu

[wconstab]

Any update on this?

It would be nice to know if the segfault is happening on the client or server side, and also whether it persists when we use the LibUV backend. We plan to roll out libuv as default anyway, so that would be an easy fix if it were the case.

[c-p-i-o]

whether it persists when we use the LibUV backend.

Issue persists even with LibUV. See #125990 for additional easier repro steps.