diff --git a/nemo_deploy/llm/inference/inference_base.py b/nemo_deploy/llm/inference/inference_base.py index 0a05f1631..129344179 100644 --- a/nemo_deploy/llm/inference/inference_base.py +++ b/nemo_deploy/llm/inference/inference_base.py @@ -16,6 +16,7 @@ import atexit import logging import os +import socket from pathlib import Path from typing import List, Optional, Tuple, Union @@ -414,6 +415,22 @@ def __getattr__(self, name): return getattr(self.mcore_engine, name) +def _default_coordinator_host() -> str: + """Return a bindable address for the data-parallel inference coordinator. + + MCore defaults the coordinator host to ``socket.gethostname()``, which is a + *name*, not an address. Inside Docker it resolves to the container ID and ZMQ + cannot bind it (``ZMQError: No such device``). Resolving that name to its IP + yields an address that is both bindable and reachable from other containers, + which is what the Ray path already gets by exporting the node IP as + ``MASTER_ADDR``. Fall back to loopback only if resolution fails. + """ + try: + return socket.gethostbyname(socket.gethostname()) + except OSError: + return "127.0.0.1" + + def create_mcore_engine( path: Path, num_devices: Optional[int] = None, @@ -522,7 +539,7 @@ def create_mcore_engine( materialize_only_last_token_logits=True, ) - coordinator_host = os.environ.get("MASTER_ADDR") + coordinator_host = os.environ.get("MASTER_ADDR") or _default_coordinator_host() llm = MegatronLLM( model=model, diff --git a/tests/unit_tests/deploy/test_inference_base.py b/tests/unit_tests/deploy/test_inference_base.py index c35f4edea..da84e56ee 100644 --- a/tests/unit_tests/deploy/test_inference_base.py +++ b/tests/unit_tests/deploy/test_inference_base.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os import types import unittest from pathlib import Path @@ -826,6 +827,85 @@ def test_create_mcore_engine_megatron_format( mock_megatron_llm.assert_called_once() self.assertIsNotNone(engine) + @patch("nemo_deploy.llm.inference.inference_base.setup_megatron_model_and_tokenizer_for_inference") + @patch("nemo_deploy.llm.inference.inference_base.MegatronLLM") + def test_create_mcore_engine_coordinator_host_resolves_to_address( + self, + mock_megatron_llm, + mock_setup, + ): + """With MASTER_ADDR unset, coordinator_host must be a bindable address. + + MCore defaults the coordinator to socket.gethostname(), which is a name and + not an address; inside Docker it is the container ID and ZMQ cannot bind it + (ZMQError: No such device). Only the Ray deployable exports MASTER_ADDR, so + on the PyTriton path the resolved address has to come from here. + """ + mock_setup.return_value = ([MagicMock()], MagicMock(), MagicMock()) + mock_megatron_llm.return_value = MagicMock() + + with patch.dict(os.environ, {}, clear=True): + os.environ.pop("MASTER_ADDR", None) + with patch("nemo_deploy.llm.inference.inference_base.socket") as mock_socket: + mock_socket.gethostname.return_value = "test-container-id" # a name, not an address + mock_socket.gethostbyname.return_value = "172.17.0.2" + create_mcore_engine( + path=self.mock_path, + model_format="megatron", + inference_max_seq_length=2048, + max_batch_size=4, + ) + + host = mock_megatron_llm.call_args.kwargs["coordinator_host"] + self.assertEqual(host, "172.17.0.2") + self.assertNotEqual(host, "test-container-id", "must not pass the bare hostname through") + + @patch("nemo_deploy.llm.inference.inference_base.setup_megatron_model_and_tokenizer_for_inference") + @patch("nemo_deploy.llm.inference.inference_base.MegatronLLM") + def test_create_mcore_engine_coordinator_host_falls_back_when_unresolvable( + self, + mock_megatron_llm, + mock_setup, + ): + """If the hostname cannot be resolved, fall back to loopback rather than None.""" + mock_setup.return_value = ([MagicMock()], MagicMock(), MagicMock()) + mock_megatron_llm.return_value = MagicMock() + + with patch.dict(os.environ, {}, clear=True): + os.environ.pop("MASTER_ADDR", None) + with patch("nemo_deploy.llm.inference.inference_base.socket") as mock_socket: + mock_socket.gethostname.return_value = "test-container-id" + mock_socket.gethostbyname.side_effect = OSError("name resolution failed") + create_mcore_engine( + path=self.mock_path, + model_format="megatron", + inference_max_seq_length=2048, + max_batch_size=4, + ) + + self.assertEqual(mock_megatron_llm.call_args.kwargs["coordinator_host"], "127.0.0.1") + + @patch("nemo_deploy.llm.inference.inference_base.setup_megatron_model_and_tokenizer_for_inference") + @patch("nemo_deploy.llm.inference.inference_base.MegatronLLM") + def test_create_mcore_engine_coordinator_host_prefers_master_addr( + self, + mock_megatron_llm, + mock_setup, + ): + """An explicit MASTER_ADDR still wins over the loopback default.""" + mock_setup.return_value = ([MagicMock()], MagicMock(), MagicMock()) + mock_megatron_llm.return_value = MagicMock() + + with patch.dict(os.environ, {"MASTER_ADDR": "10.0.0.5"}): + create_mcore_engine( + path=self.mock_path, + model_format="megatron", + inference_max_seq_length=2048, + max_batch_size=4, + ) + + self.assertEqual(mock_megatron_llm.call_args.kwargs["coordinator_host"], "10.0.0.5") + @patch("nemo_deploy.llm.inference.inference_base.torch_distributed_init") @patch("nemo_deploy.llm.inference.inference_base.load_model_config") @patch("nemo_deploy.llm.inference.inference_base.initialize_megatron_for_inference")