change: bypass DNS check for studio local exec (#252)

* change: bypass DNS check for studio local exec

* fix: unit test

* fix: unit test

---------

Co-authored-by: Mufaddal Rohawala <[email protected]>

Author: mufaddal-rohawala

src/sagemaker_pytorch_container/training.py

@@ -44,10 +44,15 @@ def train(training_environment):
         training_environment: training environment object containing environment
             variables, training arguments and hyperparameters.
     """
-    # Block until all host DNS lookups succeed. Relies on retrying dns_lookup.
-    logger.info('Block until all host DNS lookups succeed.')
-    for host in training_environment.hosts:
-        _dns_lookup(host)
+    _sm_studio_local_mode = os.environ.get("SM_STUDIO_LOCAL_MODE", "False").lower() == "true"
+
+    if not _sm_studio_local_mode:
+        # Block until all host DNS lookups succeed. Relies on retrying dns_lookup.
+        logger.info('Block until all host DNS lookups succeed.')
+        for host in training_environment.hosts:
+            _dns_lookup(host)
+    else:
+        logger.info('Bypass DNS check in case of Studio Local Mode execution.')
 
     _set_nccl_environment(training_environment.network_interface_name)
 

test/unit/test_train.py

@@ -69,17 +69,34 @@ def fixture_user_module_with_save():
     return MagicMock(spec=['train', 'save'])
 
 
+@patch('sagemaker_pytorch_container.training._dns_lookup')
 @patch('sagemaker_training.entry_point.run')
 @patch('socket.gethostbyname', MagicMock())
-def test_train(run_entry_point, training_env):
+def test_train(run_entry_point, dns_lookup, training_env):
     train(training_env)
+    dns_lookup.assert_called_once_with('algo-1')
+    run_entry_point.assert_called_with(uri=training_env.module_dir,
+                                       user_entry_point=training_env.user_entry_point,
+                                       args=training_env.to_cmd_args(),
+                                       env_vars=training_env.to_env_vars(),
+                                       capture_error=True,
+                                       runner_type=runner.ProcessRunnerType)
+
 
+@patch('sagemaker_pytorch_container.training._dns_lookup')
+@patch('sagemaker_training.entry_point.run')
+@patch('socket.gethostbyname', MagicMock())
+def test_train_with_sm_studio_local_mode_enabled(run_entry_point, dns_lookup, training_env):
+    os.environ['SM_STUDIO_LOCAL_MODE'] = 'True'
+    train(training_env)
+    dns_lookup.assert_not_called()
     run_entry_point.assert_called_with(uri=training_env.module_dir,
                                        user_entry_point=training_env.user_entry_point,
                                        args=training_env.to_cmd_args(),
                                        env_vars=training_env.to_env_vars(),
                                        capture_error=True,
                                        runner_type=runner.ProcessRunnerType)
+    del os.environ['SM_STUDIO_LOCAL_MODE']
 
 
 @patch('sagemaker_training.entry_point.run')