Use exact python path in trainer template (#1584)

Author: benieric

src/sagemaker/modules/templates.py

@@ -50,18 +50,16 @@
 }}
 
 check_python() {{
-    if command -v python3 &>/dev/null; then
-        SM_PYTHON_CMD="python3"
-        SM_PIP_CMD="pip3"
-        echo "Found python3"
-    elif command -v python &>/dev/null; then
-        SM_PYTHON_CMD="python"
-        SM_PIP_CMD="pip"
-        echo "Found python"
-    else
-        echo "Python may not be installed"
+    SM_PYTHON_CMD=$(command -v python3 || command -v python)
+    SM_PIP_CMD=$(command -v pip3 || command -v pip)
+
+    # Check if Python is found
+    if [[ -z "$SM_PYTHON_CMD" || -z "$SM_PIP_CMD" ]]; then
+        echo "Error: The Python executable was not found in the system path."
         return 1
     fi
+
+    return 0
 }}
 
 trap 'handle_error' ERR

src/sagemaker/modules/testing_notebooks/base_model_trainer.ipynb

@@ -46,7 +46,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model_trainer.train()"
+    "model_trainer.train(wait=False)"
    ]
   },
   {
@@ -94,7 +94,7 @@
     "    environment=env_vars,\n",
     ")\n",
     "\n",
-    "model_trainer.train()"
+    "model_trainer.train(wait=False)"
    ]
   },
   {
@@ -137,7 +137,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model_trainer.train()"
+    "model_trainer.train(wait=False)"
    ]
   },
   {

src/sagemaker/modules/train/model_trainer.py

@@ -293,7 +293,7 @@ def model_post_init(self, __context: Any):
             logger.warning("Session not provided. Using default Session.")
 
         if self.role is None:
-            self.role = get_execution_role()
+            self.role = get_execution_role(sagemaker_session=self.session)
             logger.warning(f"Role not provided. Using default role:\n{self.role}")
 
         if self.base_job_name is None:

src/sagemaker/serve/builder/model_builder.py

@@ -112,6 +112,7 @@
     download_huggingface_model_metadata,
 )
 from sagemaker.serve.validations.optimization import _validate_optimization_configuration
+from sagemaker.modules.train import ModelTrainer
 
 logger = logging.getLogger(__name__)
 
@@ -272,9 +273,9 @@ class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers, TensorflowServing,
     schema_builder: Optional[SchemaBuilder] = field(
         default=None, metadata={"help": "Defines the i/o schema of the model"}
     )
-    model: Optional[Union[object, str, "ModelTrainer", TrainingJob, Estimator]] = field(
+    model: Optional[Union[object, str, ModelTrainer, TrainingJob, Estimator]] = field(
         default=None,
-        metadata={"help": "Define object from which training artifacts can be extracted"}
+        metadata={"help": "Define object from which training artifacts can be extracted"},
     )
     inference_spec: InferenceSpec = field(
         default=None,
@@ -851,7 +852,7 @@ def build(  # pylint: disable=R0911
         Returns:
             Type[Model]: A deployable ``Model`` object.
         """
-        from sagemaker.modules.train.model_trainer import ModelTrainer
+
         self.modes = dict()
 
         if mode:

tests/integ/sagemaker/modules/conftest.py

@@ -14,28 +14,27 @@
 from __future__ import absolute_import
 
 import pytest
-import json
 
+import os
 import boto3
-from botocore.config import Config
 from sagemaker import Session
 
 DEFAULT_REGION = "us-west-2"
 
+
 @pytest.fixture(scope="module")
-def modules_boto_session(request):
-    config = request.config.getoption("--boto-config")
-    if config:
-        return boto3.Session(**json.loads(config))
+def modules_sagemaker_session():
+    region = os.environ.get("AWS_DEFAULT_REGION")
+    if not region:
+        os.environ["AWS_DEFAULT_REGION"] = DEFAULT_REGION
+        region_manual_set = True
     else:
-        return boto3.Session(region_name=DEFAULT_REGION)
+        region_manual_set = True
 
-@pytest.fixture(scope="module")
-def modules_sagemaker_session(request, modules_boto_session):
-    sagemaker_client = (
-        modules_boto_session.client(
-            "sagemaker", 
-            config=Config(retries={"max_attempts": 10, "mode": "standard"})
-        )
-    )
-    return Session(boto_session=modules_boto_session)
+    boto_session = boto3.Session(region_name=os.environ["AWS_DEFAULT_REGION"])
+    sagemaker_session = Session(boto_session=boto_session)
+
+    yield sagemaker_session
+
+    if region_manual_set:
+        del os.environ["AWS_DEFAULT_REGION"]