diff --git a/src/sagemaker/workflow/_repack_model.py b/src/sagemaker/workflow/_repack_model.py index 60b74d66c7..6ce7e41831 100644 --- a/src/sagemaker/workflow/_repack_model.py +++ b/src/sagemaker/workflow/_repack_model.py @@ -62,15 +62,15 @@ 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/ - entry_point = os.path.join("/opt/ml/code", inference_script) - shutil.copy2(entry_point, os.path.join(src_dir, "code", inference_script)) - - # copy source_dir to code/ if source_dir: + # copy /opt/ml/code to code/ if os.path.exists(code_dir): shutil.rmtree(code_dir) - shutil.copytree(source_dir, code_dir) + shutil.copytree("/opt/ml/code", code_dir) + else: + # copy the custom inference script to code/ + entry_point = os.path.join("/opt/ml/code", inference_script) + shutil.copy2(entry_point, os.path.join(code_dir, inference_script)) # copy any dependencies to code/lib/ if dependencies: @@ -79,13 +79,16 @@ def repack(inference_script, model_archive, dependencies=None, source_dir=None): lib_dir = os.path.join(code_dir, "lib") if not os.path.exists(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)), - ) - else: + if os.path.isfile(actual_dependency_path): shutil.copy2(actual_dependency_path, lib_dir) + else: + if os.path.exists(lib_dir): + shutil.rmtree(lib_dir) + # a directory is in the dependencies. we have to copy + # all of /opt/ml/code into the lib dir because the original directory + # was flattened by the SDK training job upload.. + shutil.copytree("/opt/ml/code", lib_dir) + break # copy the "src" dir, which includes the previous training job's model and the # custom inference script, to the output of this training job diff --git a/tests/unit/sagemaker/workflow/test_repack_model_script.py b/tests/unit/sagemaker/workflow/test_repack_model_script.py index 67c8231dcc..69c9e7b740 100644 --- a/tests/unit/sagemaker/workflow/test_repack_model_script.py +++ b/tests/unit/sagemaker/workflow/test_repack_model_script.py @@ -94,7 +94,7 @@ def test_repack_with_dependencies(tmp): _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 @@ -145,7 +145,7 @@ def test_repack_with_source_dir_and_dependencies(tmp): _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", source_dir="sourcedir", )