feat: moved find_free_port from test to src to use it in dcp to huggingface conversion

Author: BlueCrescent

src/modalities/conversion/gpt2/conversion_model.py

@@ -16,6 +16,7 @@
 from modalities.models.utils import ModelTypeEnum, get_model_from_config
 from modalities.running_env.cuda_env import MultiProcessingCudaEnv
 from modalities.running_env.env_utils import PyTorchDtypes
+from modalities.utils.ports import find_free_port
 
 
 def convert_model_checkpoint(modalities_config: ConfigDictType) -> tuple[GPT2ForCausalLM | LlamaForCausalLM, GPT2LLM]:
@@ -110,7 +111,9 @@ def check_converted_dcp_model(
     vocab_size: int = new_config["model_raw" if "model_raw" in new_config else "model"]["config"]["vocab_size"]
     if isinstance(device_id_modalities, str):
         device_id_modalities = int(device_id_modalities.replace("cuda:", ""))
-    with MultiProcessingCudaEnv(ProcessGroupBackendType.nccl, 0, 0, 1, 24570, device_id=device_id_modalities):
+    with MultiProcessingCudaEnv(
+        ProcessGroupBackendType.nccl, 0, 0, 1, find_free_port(), device_id=device_id_modalities
+    ):
         modalities_model = get_model_from_config(new_config, model_type=ModelTypeEnum.DCP_CHECKPOINTED_MODEL)
         check_converted_model(hf_model, modalities_model, num_testruns=num_testruns, vocab_size=vocab_size)
 

src/modalities/utils/ports.py

@@ -0,0 +1,9 @@
+import socket
+
+
+def find_free_port():
+    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    s.bind(("127.0.0.1", 0))
+    port = s.getsockname()[1]
+    s.close()
+    return port

tests/conversion/gpt2/conftest.py

@@ -23,8 +23,9 @@
 from modalities.registry.registry import Registry
 from modalities.running_env.cuda_env import MultiProcessingCudaEnv
 from modalities.training.training_progress import TrainingProgress
+from modalities.utils.ports import find_free_port
 from tests.conftest import _ROOT_DIR
-from tests.utility import find_free_port, monitor_child_processes
+from tests.utility import monitor_child_processes
 
 
 @pytest.fixture

tests/dataloader/distributed/test_distributed_multidim_dataloader.py

@@ -10,10 +10,11 @@
 from modalities.dataloader.sampler_factory import SamplerFactory
 from modalities.running_env.cuda_env import MultiProcessingCudaEnv
 from modalities.running_env.fsdp.device_mesh import ParallelismDegrees, get_device_mesh, get_mesh_for_parallelism_method
+from modalities.utils.ports import find_free_port
 from tests.dataloader.distributed.mocks import MultiProcessingCudaEnvMock
 from tests.dataloader.dummy_sequential_dataset import TestDataset
 from tests.mocks import MockDeviceMesh
-from tests.utility import find_free_port, tensors_equal_across_mesh, tensors_pairwise_not_equal_across_mesh
+from tests.utility import tensors_equal_across_mesh, tensors_pairwise_not_equal_across_mesh
 
 
 @pytest.mark.parametrize("world_size, dp_degree", [(4, 2)])

tests/fsdp2_parallelization/test_tensor_parallelism.py

@@ -17,7 +17,7 @@
 from modalities.models.gpt2.gpt2_model import TransformerMLP
 from modalities.models.model import SwiGLU
 from modalities.running_env.cuda_env import MultiProcessingCudaEnv
-from tests.utility import find_free_port
+from modalities.utils.ports import find_free_port
 
 
 def patch_config_file(original_config_path: Path, activation_type: str, tmp_dir: Path) -> Path:

tests/test_optimizer_factory.py

@@ -19,8 +19,8 @@
 from modalities.registry.registry import Registry
 from modalities.running_env.cuda_env import MultiProcessingCudaEnv
 from modalities.running_env.env_utils import MixedPrecisionSettings
+from modalities.utils.ports import find_free_port
 from tests.conftest import _ROOT_DIR
-from tests.utility import find_free_port
 
 # number of parameters for each optimizer group
 GPT2_LINEAR = 66130944

tests/test_util.py

@@ -13,8 +13,8 @@
 from modalities.config.pydantic_if_types import PydanticAppStateType, PydanticDeviceMeshIFType
 from modalities.running_env.cuda_env import MultiProcessingCudaEnv
 from modalities.util import get_local_number_of_trainable_parameters, get_total_number_of_trainable_parameters
+from modalities.utils.ports import find_free_port
 from modalities.utils.typing_utils import FSDPX
-from tests.utility import find_free_port
 
 
 def test_get_local_number_of_trainable_parameters():

tests/training/gradient_clipping/test_fsdp_gradient_clipper.py

@@ -15,7 +15,7 @@
     FSDP2LoggingOnlyGradientClipper,
     GradientClippingMode,
 )
-from tests.utility import find_free_port
+from modalities.utils.ports import find_free_port
 
 
 class MockFSDPModel:

tests/utility.py

@@ -1,5 +1,4 @@
 import os
-import socket
 import time
 from multiprocessing import Queue
 from multiprocessing.managers import SyncManager
@@ -13,14 +12,6 @@
 from modalities.batch import DatasetBatch
 
 
-def find_free_port():
-    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-    s.bind(("127.0.0.1", 0))
-    port = s.getsockname()[1]
-    s.close()
-    return port
-
-
 def add_debugger_to_distributed_test():
     """Add a debugger to a distributed test.
     This function should be called at the beginning of the test.