Test retries and error handling, clean some lines

Signed-off-by: Frankie Siino <fsiino@nvidia.com>

Author: fsiino-nvidia

nemo_gym/global_config.py

@@ -193,13 +193,9 @@ def parse(self, parse_config: Optional[GlobalConfigDictParserConfig] = None) ->
         # Do one pass through all the configs validate and populate various configs for our servers.
         default_host = global_config_dict.get(DEFAULT_HOST_KEY_NAME) or "127.0.0.1"
 
-        head_server_config = global_config_dict.get(HEAD_SERVER_KEY_NAME)
-        if head_server_config:
-            head_server_host = head_server_config.get("host") or default_host
-            head_server_port = head_server_config.get("port") or DEFAULT_HEAD_SERVER_PORT
-        else:
-            head_server_host = default_host
-            head_server_port = DEFAULT_HEAD_SERVER_PORT
+        head_server_config = global_config_dict.get(HEAD_SERVER_KEY_NAME, {})
+        head_server_host = head_server_config.get("host", default_host)
+        head_server_port = head_server_config.get("port", DEFAULT_HEAD_SERVER_PORT)
 
         self.validate_and_populate_defaults(server_instance_configs, default_host, head_server_host, head_server_port)
 

tests/unit_tests/test_global_config.py

@@ -18,7 +18,9 @@
 import nemo_gym.global_config
 import nemo_gym.server_utils
 from nemo_gym.global_config import (
+    DEFAULT_HEAD_SERVER_PORT,
     NEMO_GYM_CONFIG_DICT_ENV_VAR_NAME,
+    find_open_port,
     get_first_server_config_dict,
     get_global_config_dict,
 )
@@ -371,3 +373,40 @@ def test_get_first_server_config_dict(self) -> None:
             }
         )
         assert {"my_key": "my_value"} == get_first_server_config_dict(global_config_dict, "a")
+
+    def test_find_open_port_avoids_head_server_port(self, monkeypatch: MonkeyPatch) -> None:
+        """Test that find_open_port retries when the head server port is returned."""
+        socket_mock = MagicMock()
+        socket_instance = MagicMock()
+        socket_mock.return_value.__enter__ = MagicMock(return_value=socket_instance)
+        socket_mock.return_value.__exit__ = MagicMock(return_value=False)
+
+        socket_instance.getsockname.side_effect = [
+            ("", DEFAULT_HEAD_SERVER_PORT),  # first attempt: 11000 (conflict)
+            ("", 12345),  # second attempt (safe)
+        ]
+
+        monkeypatch.setattr(nemo_gym.global_config, "socket", socket_mock)
+
+        port = find_open_port(head_server_host="127.0.0.1", head_server_port=DEFAULT_HEAD_SERVER_PORT)
+
+        assert port == 12345
+        assert socket_instance.getsockname.call_count == 2  # first: conflict, second: success
+
+    def test_find_open_port_raises_after_max_retries(self, monkeypatch: MonkeyPatch) -> None:
+        """Test that find_open_port raises RuntimeError after exhausting retries."""
+        socket_mock = MagicMock()
+        socket_instance = MagicMock()
+        socket_mock.return_value.__enter__ = MagicMock(return_value=socket_instance)
+        socket_mock.return_value.__exit__ = MagicMock(return_value=False)
+
+        socket_instance.getsockname.return_value = ("", DEFAULT_HEAD_SERVER_PORT)  # force conflict
+
+        monkeypatch.setattr(nemo_gym.global_config, "socket", socket_mock)
+
+        with raises(RuntimeError) as exc_info:
+            find_open_port(head_server_host="127.0.0.1", head_server_port=DEFAULT_HEAD_SERVER_PORT, max_retries=5)
+
+        assert "Unable to find an open port" in str(exc_info.value)
+        assert "after 5 attempts" in str(exc_info.value)
+        assert socket_instance.getsockname.call_count == 5