fix: Move sagemaker pysdk version check after bootstrap in remote job

Author: qidewenwhen

src/sagemaker/remote_function/runtime_environment/bootstrap_runtime_environment.py

@@ -65,9 +65,6 @@ def main(sys_args=None):
         conda_env = job_conda_env or os.getenv("SAGEMAKER_JOB_CONDA_ENV")
 
         RuntimeEnvironmentManager()._validate_python_version(client_python_version, conda_env)
-        RuntimeEnvironmentManager()._validate_sagemaker_pysdk_version(
-            client_sagemaker_pysdk_version
-        )
 
         user = getpass.getuser()
         if user != "root":
@@ -89,6 +86,10 @@ def main(sys_args=None):
                 client_python_version, conda_env, dependency_settings
             )
 
+        RuntimeEnvironmentManager()._validate_sagemaker_pysdk_version(
+            client_sagemaker_pysdk_version
+        )
+
         exit_code = SUCCESS_EXIT_CODE
     except Exception as e:  # pylint: disable=broad-except
         logger.exception("Error encountered while bootstrapping runtime environment: %s", e)

src/sagemaker/remote_function/runtime_environment/runtime_environment_manager.py

@@ -24,8 +24,6 @@
 import dataclasses
 import json
 
-import sagemaker
-
 
 class _UTCFormatter(logging.Formatter):
     """Class that overrides the default local time provider in log formatter."""
@@ -330,6 +328,7 @@ def _current_python_version(self):
 
     def _current_sagemaker_pysdk_version(self):
         """Returns the current sagemaker python sdk version where program is running"""
+        import sagemaker
 
         return sagemaker.__version__
 
@@ -368,8 +367,8 @@ def _validate_sagemaker_pysdk_version(self, client_sagemaker_pysdk_version):
                 "Inconsistent sagemaker versions found: "
                 "sagemaker pysdk version found in the container is "
                 "'%s' which does not match the '%s' on the local client. "
-                "Please make sure that the python version used in the training container "
-                "is the same as the local python version in case of unexpected behaviors.",
+                "Please make sure that the pysdk version used in the training container "
+                "is the same as the local pysdk version in case of unexpected behaviors.",
                 job_sagemaker_pysdk_version,
                 client_sagemaker_pysdk_version,
             )