update: validate python version (aws#888)

Co-authored-by: Rohan Gujarathi <[email protected]>

Author: 2 people

src/sagemaker/remote_function/client.py

@@ -170,8 +170,12 @@ def wrapper(*args, **kwargs):
                             if (
                                 "FailureReason" in describe_result
                                 and describe_result["FailureReason"]
+                                and "RuntimeEnvironmentError: " in describe_result["FailureReason"]
                             ):
-                                raise RuntimeEnvironmentError(describe_result["FailureReason"])
+                                failure_msg = describe_result["FailureReason"].replace(
+                                    "RuntimeEnvironmentError: ", ""
+                                )
+                                raise RuntimeEnvironmentError(failure_msg)
                             raise RemoteFunctionError(
                                 "Failed to execute remote function. "
                                 + "Check corresponding job for details."
@@ -594,10 +598,13 @@ def from_describe_response(describe_training_job_response, sagemaker_session):
                     if (
                         "FailureReason" in describe_training_job_response
                         and describe_training_job_response["FailureReason"]
+                        and "RuntimeEnvironmentError: "
+                        in describe_training_job_response["FailureReason"]
                     ):
-                        job_exception = RuntimeEnvironmentError(
-                            describe_training_job_response["FailureReason"]
+                        failure_msg = describe_training_job_response["FailureReason"].replace(
+                            "RuntimeEnvironmentError: ", ""
                         )
+                        job_exception = RuntimeEnvironmentError(failure_msg)
                     else:
                         job_exception = RemoteFunctionError(
                             "Failed to execute remote function. "
@@ -687,10 +694,13 @@ def result(self, timeout: float = None) -> Any:
                             if (
                                 "FailureReason" in self._job.describe()
                                 and self._job.describe()["FailureReason"]
+                                and "RuntimeEnvironmentError: "
+                                in self._job.describe()["FailureReason"]
                             ):
-                                self._exception = RuntimeEnvironmentError(
-                                    self._job.describe()["FailureReason"]
+                                failure_msg = self._job.describe()["FailureReason"].replace(
+                                    "RuntimeEnvironmentError: ", ""
                                 )
+                                self._exception = RuntimeEnvironmentError(failure_msg)
                             else:
                                 self._exception = RemoteFunctionError(
                                     "Failed to execute remote function. "

src/sagemaker/remote_function/job.py

@@ -354,6 +354,9 @@ def start(job_settings: _JobSettings, func, func_args, func_kwargs, run_info=Non
 
         container_args = ["--s3_base_uri", s3_base_uri]
         container_args.extend(["--region", job_settings.sagemaker_session.boto_region_name])
+        container_args.extend(
+            ["--client_python_version", RuntimeEnvironmentManager()._current_python_version()]
+        )
         if job_settings.s3_kms_key:
             container_args.extend(["--s3_kms_key", job_settings.s3_kms_key])
 

src/sagemaker/remote_function/runtime_environment/bootstrap_runtime_environment.py

@@ -50,15 +50,22 @@ def main():
 
     try:
         args = _parse_agrs()
+        client_python_version = args.client_python_version
+        job_conda_env = args.job_conda_env
+
+        conda_env = job_conda_env or os.getenv("SAGEMAKER_JOB_CONDA_ENV")
+
+        RuntimeEnvironmentManager()._validate_python_version(client_python_version, conda_env)
 
         _execute_pre_exec_cmds()
 
-        _bootstrap_runtime_environment(job_conda_env=args.job_conda_env)
+        _bootstrap_runtime_environment(client_python_version, conda_env)
 
         exit_code = SUCCESS_EXIT_CODE
     except Exception as e:  # pylint: disable=broad-except
         logger.exception("Error encountered while bootstrapping runtime environment: %s", e)
-        _write_failure_reason_file(repr(e))
+
+        _write_failure_reason_file(str(e))
     finally:
         sys.exit(exit_code)
 
@@ -70,12 +77,13 @@ def _execute_pre_exec_cmds():
 
 
 def _bootstrap_runtime_environment(
-    job_conda_env: str = None,
+    client_python_version: str,
+    conda_env: str = None,
 ):
     """Bootstrap runtime environment for remote function invocation
 
     Args:
-        job_conda_env (str): conda environment to be activated. Default is None.
+        conda_env (str): conda environment to be activated. Default is None.
     """
     workspace_archive_dir_path = os.path.join(BASE_CHANNEL_PATH, REMOTE_FUNCTION_WORKSPACE)
 
@@ -109,7 +117,9 @@ def _bootstrap_runtime_environment(
     if dependencies_file:
         _create_sm_installer_symlinks()
         RuntimeEnvironmentManager().bootstrap(
-            local_dependencies_file=dependencies_file, job_conda_env=job_conda_env
+            local_dependencies_file=dependencies_file,
+            conda_env=conda_env,
+            client_python_version=client_python_version,
         )
     else:
         logger.info(
@@ -151,13 +161,14 @@ def _write_failure_reason_file(failure_msg):
     """
     if not os.path.exists(FAILURE_REASON_PATH):
         with open(FAILURE_REASON_PATH, "w") as f:
-            f.write(failure_msg)
+            f.write("RuntimeEnvironmentError: " + failure_msg)
 
 
 def _parse_agrs():
     """Parses CLI arguments."""
     parser = argparse.ArgumentParser()
     parser.add_argument("--job_conda_env", type=str)
+    parser.add_argument("--client_python_version")
     args, _ = parser.parse_known_args()
     return args
 

src/sagemaker/remote_function/runtime_environment/runtime_environment_manager.py

@@ -16,6 +16,7 @@
 
 import logging
 import sys
+import shlex
 import os
 import subprocess
 import tempfile
@@ -76,18 +77,18 @@ def snapshot(self, dependencies: str = None) -> str:
 
         raise ValueError(f'Invalid dependencies provided: "{dependencies}"')
 
-    def bootstrap(self, local_dependencies_file: str, job_conda_env: str = None):
+    def bootstrap(
+        self, local_dependencies_file: str, client_python_version: str, conda_env: str = None
+    ):
         """Bootstraps the runtime environment by installing the additional dependencies if any.
 
         Args:
             dependencies_s3_uri (str): S3 URI where dependencies file exists.
-            job_conda_env (str): conda environment to be activated. Default is None.
+            conda_env (str): conda environment to be activated. Default is None.
 
         Returns: None
         """
 
-        conda_env = job_conda_env or os.getenv("JOB_CONDA_ENV")
-
         if local_dependencies_file.endswith(".txt"):
             if conda_env:
                 self._install_req_txt_in_conda_env(conda_env, local_dependencies_file)
@@ -102,6 +103,7 @@ def bootstrap(self, local_dependencies_file: str, job_conda_env: str = None):
             else:
                 conda_env = "sagemaker-runtime-env"
                 self._create_conda_env(conda_env, local_dependencies_file)
+                self._validate_python_version(client_python_version, conda_env)
             self._write_conda_env_to_file(conda_env)
 
     def _is_file_exists(self, dependencies):
@@ -176,6 +178,44 @@ def _get_conda_exe(self):
             return "conda"
         raise ValueError("Neither conda nor mamba is installed on the image")
 
+    def _python_version_in_conda_env(self, env_name):
+        """Returns python version inside a conda environment"""
+        cmd = f"{self._get_conda_exe()} run -n {env_name} python --version"
+        try:
+            output = (
+                subprocess.check_output(shlex.split(cmd), stderr=subprocess.STDOUT)
+                .decode("utf-8")
+                .strip()
+            )
+            # convert 'Python 3.7.16' to [3, 7, 16]
+            version = output.split("Python ")[1].split(".")
+            return version[0] + "." + version[1]
+        except subprocess.CalledProcessError as e:
+            raise RuntimeEnvironmentError(e.output)
+
+    def _current_python_version(self):
+        """Returns the current python version where program is running"""
+
+        return f"{sys.version_info.major}.{sys.version_info.minor}"
+
+    def _validate_python_version(self, client_python_version: str, conda_env: str = None):
+        """Validate the python version
+
+        Validates if the python version where remote function runs
+        matches the one used on client side.
+        """
+        if conda_env:
+            job_python_version = self._python_version_in_conda_env(conda_env)
+        else:
+            job_python_version = self._current_python_version()
+        if client_python_version != job_python_version:
+            raise RuntimeEnvironmentError(
+                f"Python version found in the container is {job_python_version} which "
+                f"does not match python version {client_python_version} on the local client . "
+                f"Please make sure that the python version used in the training container "
+                f"is same as the local python version."
+            )
+
 
 def _run_shell_cmd(cmd: str):
     """This method runs a given shell command using subprocess

tests/integ/sagemaker/remote_function/conftest.py

@@ -51,7 +51,7 @@
         && mamba run -n default_env pip install '{source_archive}[remote_function]' \
         && mamba run -n integ_test_env pip install '{source_archive}[remote_function]'\n"
     "ENV SHELL=/bin/bash\n"
-    "ENV JOB_CONDA_ENV=default_env\n"
+    "ENV SAGEMAKER_JOB_CONDA_ENV=default_env\n"
 )
 
 CONDA_YML_FILE_TEMPLATE = (

tests/integ/sagemaker/remote_function/test_decorator.py

@@ -74,7 +74,6 @@ def divide(x, y):
         divide(10, 0)
 
 
-@pytest.mark.skip
 def test_remote_python_runtime_is_incompatible(
     sagemaker_session, dummy_container_incompatible_python_runtime, cpu_instance_type
 ):
@@ -87,8 +86,13 @@ def test_remote_python_runtime_is_incompatible(
     def divide(x, y):
         return x / y
 
-    # TODO: should raise serialization error
-    with pytest.raises(RuntimeError):
+    with pytest.raises(
+        RuntimeEnvironmentError,
+        match=(
+            "Please make sure that the python version used in the training container "
+            "is same as the local python version."
+        ),
+    ):
         divide(10, 2)
 
 

tests/unit/sagemaker/remote_function/runtime_environment/test_bootstrap_runtime_environment.py

@@ -22,20 +22,26 @@
 
 TEST_JOB_CONDA_ENV = "conda_env"
 TEST_DEPENDENCIES_PATH = "/user/set/workdir"
+TEST_PYTHON_VERSION = "3.10"
 TEST_WORKSPACE_ARCHIVE_DIR_PATH = "/opt/ml/input/data/sm_rf_user_ws"
 TEST_WORKSPACE_ARCHIVE_PATH = "/opt/ml/input/data/sm_rf_user_ws/workspace.zip"
 
 
 def mock_args():
     args = Mock()
     args.job_conda_env = TEST_JOB_CONDA_ENV
+    args.client_python_version = TEST_PYTHON_VERSION
     return args
 
 
 @patch(
     "sagemaker.remote_function.runtime_environment.bootstrap_runtime_environment._parse_agrs",
     new=mock_args,
 )
+@patch(
+    "sagemaker.remote_function.runtime_environment.runtime_environment_manager."
+    "RuntimeEnvironmentManager._validate_python_version"
+)
 @patch(
     "sagemaker.remote_function.runtime_environment.bootstrap_runtime_environment._create_sm_installer_symlinks"
 )
@@ -49,9 +55,17 @@ def mock_args():
     "sagemaker.remote_function.runtime_environment.runtime_environment_manager.RuntimeEnvironmentManager.bootstrap"
 )
 def test_main_success(
-    bootstrap_runtime, list_dir, file_exists, path_exists, getcwd, _exit_process, symlink
+    bootstrap_runtime,
+    list_dir,
+    file_exists,
+    path_exists,
+    getcwd,
+    _exit_process,
+    symlink,
+    validate_python,
 ):
     bootstrap.main()
+    validate_python.assert_called_once_with(TEST_PYTHON_VERSION, TEST_JOB_CONDA_ENV)
     path_exists.assert_called_once_with(TEST_WORKSPACE_ARCHIVE_DIR_PATH)
     file_exists.assert_called_once_with(TEST_WORKSPACE_ARCHIVE_PATH)
     getcwd.assert_called()
@@ -64,6 +78,10 @@ def test_main_success(
     "sagemaker.remote_function.runtime_environment.bootstrap_runtime_environment._parse_agrs",
     new=mock_args,
 )
+@patch(
+    "sagemaker.remote_function.runtime_environment.runtime_environment_manager."
+    "RuntimeEnvironmentManager._validate_python_version"
+)
 @patch(
     "sagemaker.remote_function.runtime_environment.bootstrap_runtime_environment._create_sm_installer_symlinks"
 )
@@ -74,14 +92,15 @@ def test_main_success(
 @patch(
     "sagemaker.remote_function.runtime_environment.bootstrap_runtime_environment._bootstrap_runtime_environment"
 )
-def test_main_failure(bootstrap_runtime, write_failure, _exit_process, symlink):
+def test_main_failure(bootstrap_runtime, write_failure, _exit_process, symlink, validate_python):
     runtime_err = RuntimeEnvironmentError("some failure reason")
     bootstrap_runtime.side_effect = runtime_err
 
     bootstrap.main()
 
+    validate_python.assert_called_once_with(TEST_PYTHON_VERSION, TEST_JOB_CONDA_ENV)
     bootstrap_runtime.assert_called()
-    write_failure.assert_called_with(repr(runtime_err))
+    write_failure.assert_called_with(str(runtime_err))
     _exit_process.assert_called_with(1)
 
 
@@ -90,13 +109,20 @@ def test_main_failure(bootstrap_runtime, write_failure, _exit_process, symlink):
     new=mock_args,
 )
 @patch("sys.exit")
+@patch(
+    "sagemaker.remote_function.runtime_environment.runtime_environment_manager."
+    "RuntimeEnvironmentManager._validate_python_version"
+)
 @patch(
     "sagemaker.remote_function.runtime_environment.bootstrap_runtime_environment._write_failure_reason_file"
 )
 @patch("os.path.exists", return_value=False)
-def test_main_channel_folder_does_not_exist(path_exists, write_failure, _exit_process):
+def test_main_channel_folder_does_not_exist(
+    path_exists, write_failure, validate_python, _exit_process
+):
     bootstrap.main()
     path_exists.assert_called_once_with(TEST_WORKSPACE_ARCHIVE_DIR_PATH)
+    validate_python.assert_called_once_with(TEST_PYTHON_VERSION, TEST_JOB_CONDA_ENV)
     write_failure.assert_not_called()
     _exit_process.assert_called_with(0)
 
@@ -105,14 +131,21 @@ def test_main_channel_folder_does_not_exist(path_exists, write_failure, _exit_pr
     "sagemaker.remote_function.runtime_environment.bootstrap_runtime_environment._parse_agrs",
     new=mock_args,
 )
+@patch(
+    "sagemaker.remote_function.runtime_environment.runtime_environment_manager"
+    ".RuntimeEnvironmentManager._validate_python_version"
+)
 @patch("sys.exit")
 @patch("os.path.exists", return_value=True)
 @patch("os.path.isfile", return_value=False)
 @patch(
     "sagemaker.remote_function.runtime_environment.runtime_environment_manager.RuntimeEnvironmentManager.bootstrap"
 )
-def test_main_no_workspace_archive(bootstrap_runtime, file_exists, path_exists, _exit_process):
+def test_main_no_workspace_archive(
+    bootstrap_runtime, file_exists, path_exists, _exit_process, validate_python
+):
     bootstrap.main()
+    validate_python.assert_called_once_with(TEST_PYTHON_VERSION, TEST_JOB_CONDA_ENV)
     path_exists.assert_called_once_with(TEST_WORKSPACE_ARCHIVE_DIR_PATH)
     file_exists.assert_called_once_with(TEST_WORKSPACE_ARCHIVE_PATH)
     bootstrap_runtime.assert_not_called()
@@ -123,6 +156,10 @@ def test_main_no_workspace_archive(bootstrap_runtime, file_exists, path_exists,
     "sagemaker.remote_function.runtime_environment.bootstrap_runtime_environment._parse_agrs",
     new=mock_args,
 )
+@patch(
+    "sagemaker.remote_function.runtime_environment.runtime_environment_manager."
+    "RuntimeEnvironmentManager._validate_python_version"
+)
 @patch("sys.exit")
 @patch("shutil.unpack_archive", Mock())
 @patch("os.path.exists", return_value=True)
@@ -133,9 +170,10 @@ def test_main_no_workspace_archive(bootstrap_runtime, file_exists, path_exists,
     "sagemaker.remote_function.runtime_environment.runtime_environment_manager.RuntimeEnvironmentManager.bootstrap"
 )
 def test_main_no_dependency_file(
-    bootstrap_runtime, list_dir, get_cwd, file_exists, path_exists, _exit_process
+    bootstrap_runtime, list_dir, get_cwd, file_exists, path_exists, _exit_process, validate_python
 ):
     bootstrap.main()
+    validate_python.assert_called_once_with(TEST_PYTHON_VERSION, TEST_JOB_CONDA_ENV)
     path_exists.assert_called_once_with(TEST_WORKSPACE_ARCHIVE_DIR_PATH)
     file_exists.assert_called_once_with(TEST_WORKSPACE_ARCHIVE_PATH)
     get_cwd.assert_called_once()