Use sagemaker core Session (#1607)

* Use sagemaker core Session

* update tests with session

* update

* flake8

* update docs

---------

Co-authored-by: Gokul Anantha Narayanan <[email protected]>

Author: 2 people

src/sagemaker/modules/__init__.py

@@ -14,5 +14,6 @@
 from __future__ import absolute_import
 
 from sagemaker_core.main.utils import logger as sagemaker_core_logger
+from sagemaker_core.helper.session_helper import Session, get_execution_role  # noqa: F401
 
 logger = sagemaker_core_logger

src/sagemaker/modules/train/model_trainer.py

@@ -26,7 +26,7 @@
 
 from pydantic import BaseModel, ConfigDict, PrivateAttr, validate_call
 
-from sagemaker import get_execution_role, Session
+from sagemaker.modules import Session, get_execution_role
 from sagemaker.modules.configs import (
     Compute,
     StoppingCondition,
@@ -119,8 +119,9 @@ class ModelTrainer(BaseModel):
     ```
 
     Attributes:
-        sagemaker_session (Optiona(Session)):
-            The SageMaker session.
+        session (Optiona(Session)):
+            The SageMakerCore session. For convinience, can be imported like:
+            `from sagemaker.modules import Session`.
             If not specified, a new session will be created.
         role (Optional(str)):
             The IAM role ARN for the training job.

tests/data/modules/params-script/train.py renamed to tests/data/modules/params_script/train.py

@@ -17,7 +17,7 @@
 
 import os
 import boto3
-from sagemaker import Session
+from sagemaker.modules import Session
 
 DEFAULT_REGION = "us-west-2"
 

tests/data/modules/params-script/train.sh renamed to tests/data/modules/params_script/train.sh

@@ -39,7 +39,7 @@
 
 def test_hp_contract_basic_py_script(modules_sagemaker_session):
     source_code = SourceCode(
-        source_dir=f"{DATA_DIR}/modules/params-script",
+        source_dir=f"{DATA_DIR}/modules/params_script",
         entry_script="train.py",
     )
 
@@ -56,7 +56,7 @@ def test_hp_contract_basic_py_script(modules_sagemaker_session):
 
 def test_hp_contract_basic_sh_script(modules_sagemaker_session):
     source_code = SourceCode(
-        source_dir=f"{DATA_DIR}/modules/params-script",
+        source_dir=f"{DATA_DIR}/modules/params_script",
         entry_script="train.sh",
     )
     model_trainer = ModelTrainer(
@@ -72,7 +72,7 @@ def test_hp_contract_basic_sh_script(modules_sagemaker_session):
 
 def test_hp_contract_mpi_script(modules_sagemaker_session):
     source_code = SourceCode(
-        source_dir=f"{DATA_DIR}/modules/params-script",
+        source_dir=f"{DATA_DIR}/modules/params_script",
         entry_script="train.py",
     )
     compute = Compute(instance_type="ml.m5.xlarge", instance_count=2)
@@ -91,7 +91,7 @@ def test_hp_contract_mpi_script(modules_sagemaker_session):
 
 def test_hp_contract_torchrun_script(modules_sagemaker_session):
     source_code = SourceCode(
-        source_dir=f"{DATA_DIR}/modules/params-script",
+        source_dir=f"{DATA_DIR}/modules/params_script",
         entry_script="train.py",
     )
     compute = Compute(instance_type="ml.m5.xlarge", instance_count=2)

tests/integ/sagemaker/modules/conftest.py

@@ -22,7 +22,7 @@
 
 from sagemaker_core.main.resources import TrainingJob
 
-from sagemaker.session import Session
+from sagemaker.modules import Session
 from sagemaker.modules.train.model_trainer import ModelTrainer
 from sagemaker.modules.constants import (
     DEFAULT_INSTANCE_TYPE,
@@ -80,7 +80,7 @@
 
 @pytest.fixture(scope="module", autouse=True)
 def modules_session():
-    with patch("sagemaker.session.Session", spec=Session) as session_mock:
+    with patch("sagemaker.modules.Session", spec=Session) as session_mock:
         session_instance = session_mock.return_value
         session_instance.default_bucket.return_value = DEFAULT_BUCKET
         session_instance.get_caller_identity_arn.return_value = DEFAULT_ROLE