fix: patch socket call and update flake8 violations (#15)

Author: metrizable

test/integration/local/test_smdataparallel.py

@@ -12,16 +12,15 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
-import json
 import os
-import tarfile
 
 import pytest
 from sagemaker.pytorch import PyTorch
 
 from integration import resources_path
 from utils.local_mode_utils import assert_files_exist
 
+
 # TODO: Enable the test once SMDataParallel DLC is publicly accessibly
 @pytest.mark.skip(reason="SMDataParallel DLC is not publicly accessible")
 @pytest.mark.skip_cpu

test/integration/sagemaker/test_smdataparallel.py

@@ -12,9 +12,7 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
-import json
 import os
-import tarfile
 
 import pytest
 from sagemaker.pytorch import PyTorch
@@ -30,7 +28,9 @@
     "instances, train_instance_type",
     [(1, "ml.p3.16xlarge"), (2, "ml.p3.16xlarge"), (1, "ml.p3dn.24xlarge"), (2, "ml.p3dn.24xlarge")],
 )
-def test_smdataparallel_training(instances, train_instance_type, sagemaker_session, image_uri, framework_version, tmpdir):
+def test_smdataparallel_training(
+    instances, train_instance_type, sagemaker_session, image_uri, framework_version, tmpdir
+):
     default_bucket = sagemaker_session.default_bucket()
     output_path = "s3://" + os.path.join(default_bucket, "pytorch/smdataparallel")
 

test/unit/test_train.py

@@ -72,6 +72,7 @@ def test_train(run_entry_point, training_env):
 
 
 @patch("sagemaker_training.entry_point.run")
+@patch('socket.gethostbyname', MagicMock())
 def test_train_smdataparallel(run_module, training_env):
     training_env.additional_framework_parameters["sagemaker_distributed_dataparallel_enabled"] = True