Prevent local command deadlock on unconsumed pipes

src/torchrunx/utils/environment.py

@@ -178,12 +178,13 @@ def execute_command(
     if is_localhost:
         # S602: subprocess.Popen is called with shell=True (https://docs.python.org/3.10/library/subprocess.html#security-considerations)
         # Made sure to shlex.quote arguments in build_command to prevent shell injection
+        output_target = subprocess.PIPE if return_stdout_stderr else None
         process = subprocess.Popen(  # noqa: S602
             command,
             shell=True,
             text=True,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
+            stdout=output_target,
+            stderr=output_target,
         )
 
         if return_stdout_stderr:

tests/test_ci.py

@@ -1,15 +1,18 @@
 import datetime
 import os
+import subprocess
 import tempfile
 import time
 from pathlib import Path
 from typing import NoReturn
+from unittest import mock
 
 import pytest
 import torch
 import torch.distributed as dist
 
 import torchrunx as trx
+from torchrunx.utils.environment import execute_command
 
 
 def test_simple_localhost() -> None:
@@ -105,5 +108,34 @@ def error_func() -> NoReturn:
     assert "abcdefg" in str(excinfo.value)
 
 
+def test_execute_command_inherits_output_when_not_returned(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    process = mock.Mock()
+    popen = mock.Mock(return_value=process)
+    monkeypatch.setattr(subprocess, "Popen", popen)
+
+    result = execute_command("echo output", "localhost")
+
+    assert result == ("", "")
+    assert popen.call_args.kwargs["stdout"] is None
+    assert popen.call_args.kwargs["stderr"] is None
+    process.communicate.assert_not_called()
+
+
+def test_execute_command_captures_requested_output(monkeypatch: pytest.MonkeyPatch) -> None:
+    process = mock.Mock()
+    process.communicate.return_value = ("stdout", "stderr")
+    popen = mock.Mock(return_value=process)
+    monkeypatch.setattr(subprocess, "Popen", popen)
+
+    result = execute_command("echo output", "localhost", return_stdout_stderr=True)
+
+    assert result == ("stdout", "stderr")
+    assert popen.call_args.kwargs["stdout"] == subprocess.PIPE
+    assert popen.call_args.kwargs["stderr"] == subprocess.PIPE
+    process.communicate.assert_called_once_with()
+
+
 if __name__ == "__main__":
     test_simple_localhost()