From d8496ac35cf3c2a6e9bb6fd2167ea33fbe3cf889 Mon Sep 17 00:00:00 2001
From: Justin Mahlik <38999128+jmahlik@users.noreply.github.com>
Date: Mon, 25 Mar 2024 15:39:25 -0500
Subject: [PATCH] fix: remove deprecated distutils api calls in workflows

Also remove from tests.

Closes #4534
---
 .pylintrc                               | 2 +-
 src/sagemaker/workflow/_repack_model.py | 9 +--------
 tests/data/_repack_model.py             | 9 +--------
 tests/unit/test_chainer.py              | 1 -
 4 files changed, 3 insertions(+), 18 deletions(-)

diff --git a/.pylintrc b/.pylintrc
index 11e2ababa9..9ad209dd39 100644
--- a/.pylintrc
+++ b/.pylintrc
@@ -310,7 +310,7 @@ ignore-mixin-members=yes
 # (useful for modules/projects where namespaces are manipulated during runtime
 # and thus existing member attributes cannot be deduced by static analysis. It
 # supports qualified module names, as well as Unix pattern matching.
-ignored-modules=distutils
+ignored-modules=
 
 # List of class names for which member attributes should not be checked (useful
 # for classes with dynamically set attributes). This supports the use of
diff --git a/src/sagemaker/workflow/_repack_model.py b/src/sagemaker/workflow/_repack_model.py
index 84b3a426f6..2a9129da2f 100644
--- a/src/sagemaker/workflow/_repack_model.py
+++ b/src/sagemaker/workflow/_repack_model.py
@@ -27,13 +27,6 @@
 # is unpacked for inference, the custom entry point will be used.
 # Reference: https://docs.aws.amazon.com/sagemaker/latest/dg/amazon-sagemaker-toolkits.html
 
-# distutils.dir_util.copy_tree works way better than the half-baked
-# shutil.copytree which bombs on previously existing target dirs...
-# alas ... https://bugs.python.org/issue10948
-# we'll go ahead and use the copy_tree function anyways because this
-# repacking is some short-lived hackery, right??
-from distutils.dir_util import copy_tree
-
 from os.path import abspath, realpath, dirname, normpath, join as joinpath
 
 logger = logging.getLogger(__name__)
@@ -188,7 +181,7 @@ def repack(inference_script, model_archive, dependencies=None, source_dir=None):
 
         # 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")
+        shutil.copytree(src_dir, "/opt/ml/model", dirs_exist_ok=True)
 
 
 if __name__ == "__main__":  # pragma: no cover
diff --git a/tests/data/_repack_model.py b/tests/data/_repack_model.py
index 3cfa6760b3..b370db5dbf 100644
--- a/tests/data/_repack_model.py
+++ b/tests/data/_repack_model.py
@@ -26,13 +26,6 @@
 # is unpacked for inference, the custom entry point will be used.
 # Reference: https://docs.aws.amazon.com/sagemaker/latest/dg/amazon-sagemaker-toolkits.html
 
-# distutils.dir_util.copy_tree works way better than the half-baked
-# shutil.copytree which bombs on previously existing target dirs...
-# alas ... https://bugs.python.org/issue10948
-# we'll go ahead and use the copy_tree function anyways because this
-# 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):  # pragma: no cover
     """Repack custom dependencies and code into an existing model TAR archive
@@ -92,7 +85,7 @@ def repack(inference_script, model_archive, dependencies=None, source_dir=None):
 
         # 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")
+        shutil.copytree(src_dir, "/opt/ml/model", dirs_exist_ok=True)
 
 
 if __name__ == "__main__":  # pragma: no cover
diff --git a/tests/unit/test_chainer.py b/tests/unit/test_chainer.py
index 2c82e2d6fc..6a27115ca6 100644
--- a/tests/unit/test_chainer.py
+++ b/tests/unit/test_chainer.py
@@ -15,7 +15,6 @@
 import logging
 import json
 import os
-from distutils.util import strtobool
 
 import pytest
 from mock import MagicMock, Mock, ANY