Fix pylint

Author: Payton Staub

src/sagemaker/workflow/_repack_model.py

@@ -33,7 +33,16 @@
 # repacking is some short-lived hackery, right??
 from distutils.dir_util import copy_tree
 
+
 def repack(inference_script, model_archive, dependencies=None, source_dir=None):
+    """Repack custom dependencies and code into an existing model TAR archive
+
+    Args:
+        inference_script (str): The path to the custom entry point.
+        model_archive (str): The name of the model TAR archive.
+        dependencies (str): A space-delimited string of paths to custom dependencies.
+        source_dir (str): The path to a custom source directory.
+    """
 
     # the data directory contains a model archive generated by a previous training job
     data_directory = "/opt/ml/input/data/training"
@@ -53,7 +62,7 @@ def repack(inference_script, model_archive, dependencies=None, source_dir=None):
         with tarfile.open(name=local_path, mode="r:gz") as tf:
             tf.extractall(path=src_dir)
 
-        # copy the custom inference script to code/ 
+        # copy the custom inference script to code/
         entry_point = os.path.join("/opt/ml/code", inference_script)
         shutil.copy2(entry_point, os.path.join(src_dir, "code", inference_script))
 
@@ -71,19 +80,23 @@ def repack(inference_script, model_archive, dependencies=None, source_dir=None):
 
         # copy any dependencies to code/lib/
         if dependencies:
-            for dependency in dependencies.split(' '):
+            for dependency in dependencies.split(" "):
                 actual_dependency_path = os.path.join("/opt/ml/code", dependency)
                 lib_dir = os.path.join(code_dir, "lib")
                 if not os.path.exists(lib_dir):
-                        os.mkdir(lib_dir)
+                    os.mkdir(lib_dir)
                 if os.path.isdir(actual_dependency_path):
-                    shutil.copytree(actual_dependency_path, os.path.join(lib_dir, os.path.basename(actual_dependency_path)))
+                    shutil.copytree(
+                        actual_dependency_path,
+                        os.path.join(lib_dir, os.path.basename(actual_dependency_path)),
+                    )
                 else:
                     shutil.copy2(actual_dependency_path, lib_dir)
 
         # copy the "src" dir, which includes the previous training job's model and the
         # custom inference script, to the output of this training job
-        copy_tree(src_dir, "/opt/ml/model")  
+        copy_tree(src_dir, "/opt/ml/model")
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -96,5 +109,5 @@ def repack(inference_script, model_archive, dependencies=None, source_dir=None):
         inference_script=args.inference_script,
         dependencies=args.dependencies,
         source_dir=args.source_dir,
-        model_archive=args.model_archive
-        )
+        model_archive=args.model_archive,
+    )

src/sagemaker/workflow/_utils.py

@@ -148,7 +148,7 @@ def __init__(
         # convert dependencies array into space-delimited string
         dependencies_hyperparameter = None
         if self._dependencies:
-            dependencies_hyperparameter = ' '.join(self._dependencies)
+            dependencies_hyperparameter = " ".join(self._dependencies)
 
         # the real estimator and inputs
         repacker = SKLearn(
@@ -163,7 +163,7 @@ def __init__(
                 "inference_script": self._entry_point_basename,
                 "model_archive": self._model_archive,
                 "dependencies": dependencies_hyperparameter,
-                "source_dir": self._source_dir
+                "source_dir": self._source_dir,
             },
             subnets=subnets,
             security_group_ids=security_group_ids,

tests/unit/sagemaker/workflow/test_repack_model_script.py

@@ -21,10 +21,13 @@
 import pytest
 import time
 
-@pytest.mark.skip(reason="""This test operates on the root file system 
+
+@pytest.mark.skip(
+    reason="""This test operates on the root file system
                             and will likely fail due to permission errors.
                             Temporarily remove this skip decorator and run
-                            the test after making changes to _repack_model.py""")
+                            the test after making changes to _repack_model.py"""
+)
 def test_repack_entry_point_only(tmp):
     model_name = "xg-boost-model"
     fake_model_path = os.path.join(tmp, model_name)
@@ -51,19 +54,19 @@ def test_repack_entry_point_only(tmp):
     )
 
     # repack
-    _repack_model.repack(
-        inference_script="inference.py",
-        model_archive=model_tar_name
-    )
+    _repack_model.repack(inference_script="inference.py", model_archive=model_tar_name)
 
     # /opt/ml/model should now have the original model and the inference script
     assert os.path.exists(os.path.join("/opt/ml/model", model_name))
     assert os.path.exists(os.path.join("/opt/ml/model/code", "inference.py"))
 
-@pytest.mark.skip(reason="""This test operates on the root file system 
+
+@pytest.mark.skip(
+    reason="""This test operates on the root file system
                             and will likely fail due to permission errors.
                             Temporarily remove this skip decorator and run
-                            the test after making changes to _repack_model.py""")
+                            the test after making changes to _repack_model.py"""
+)
 def test_repack_with_dependencies(tmp):
     model_name = "xg-boost-model"
     fake_model_path = os.path.join(tmp, model_name)
@@ -84,19 +87,14 @@ def test_repack_with_dependencies(tmp):
     # create files that will be added to model.tar.gz
     create_file_tree(
         "/opt/ml/code",
-        [
-            "inference.py",
-            "dependencies/a",
-            "bb",
-            "dependencies/some/dir/b"
-        ],
+        ["inference.py", "dependencies/a", "bb", "dependencies/some/dir/b"],
     )
 
     # repack
     _repack_model.repack(
         inference_script="inference.py",
         model_archive=model_tar_name,
-        dependencies=["dependencies/a", "bb", "dependencies/some/dir"]
+        dependencies=["dependencies/a", "bb", "dependencies/some/dir"],
     )
 
     # /opt/ml/model should now have the original model and the inference script
@@ -106,10 +104,13 @@ def test_repack_with_dependencies(tmp):
     assert os.path.exists(os.path.join("/opt/ml/model/code/lib", "bb"))
     assert os.path.exists(os.path.join("/opt/ml/model/code/lib/dir", "b"))
 
-@pytest.mark.skip(reason="""This test operates on the root file system 
+
+@pytest.mark.skip(
+    reason="""This test operates on the root file system
                             and will likely fail due to permission errors.
                             Temporarily remove this skip decorator and run
-                            the test after making changes to _repack_model.py""")
+                            the test after making changes to _repack_model.py"""
+)
 def test_repack_with_source_dir_and_dependencies(tmp):
     model_name = "xg-boost-model"
     fake_model_path = os.path.join(tmp, model_name)
@@ -136,7 +137,7 @@ def test_repack_with_source_dir_and_dependencies(tmp):
             "bb",
             "dependencies/some/dir/b",
             "sourcedir/foo.py",
-            "sourcedir/some/dir/a"
+            "sourcedir/some/dir/a",
         ],
     )
 
@@ -145,7 +146,7 @@ def test_repack_with_source_dir_and_dependencies(tmp):
         inference_script="inference.py",
         model_archive=model_tar_name,
         dependencies=["dependencies/a", "bb", "dependencies/some/dir"],
-        source_dir="sourcedir"
+        source_dir="sourcedir",
     )
 
     # /opt/ml/model should now have the original model and the inference script