fix session creation in load_run() (aws#884)

Co-authored-by: Ao Guo <[email protected]>

Author: 2 people

src/sagemaker/experiments/run.py

@@ -795,7 +795,6 @@ def load_run(
     Returns:
         Run: The loaded Run object.
     """
-    sagemaker_session = sagemaker_session or _utils.default_session()
     environment = _RunEnvironment.load()
 
     verify_load_input_names(run_name=run_name, experiment_name=experiment_name)
@@ -810,13 +809,13 @@ def load_run(
         run_instance = Run(
             experiment_name=experiment_name,
             run_name=run_name,
-            sagemaker_session=sagemaker_session,
+            sagemaker_session=sagemaker_session or _utils.default_session(),
         )
     elif _RunContext.get_current_run():
         run_instance = _RunContext.get_current_run()
     elif environment:
         exp_config = get_tc_and_exp_config_from_job_env(
-            environment=environment, sagemaker_session=sagemaker_session
+            environment=environment, sagemaker_session=sagemaker_session or _utils.default_session()
         )
         run_name = Run._extract_run_name_from_tc_name(
             trial_component_name=exp_config[RUN_NAME],
@@ -826,7 +825,7 @@ def load_run(
         run_instance = Run(
             experiment_name=experiment_name,
             run_name=run_name,
-            sagemaker_session=sagemaker_session,
+            sagemaker_session=sagemaker_session or _utils.default_session(),
         )
     else:
         raise RuntimeError(

tests/integ/sagemaker/remote_function/test_decorator.py

@@ -382,9 +382,7 @@ def test_decorator_load_run_inside_remote_function(
         sagemaker_session=sagemaker_session,
     )
     def train():
-        boto_session = boto3.Session(region_name=os.environ["AWS_REGION"])
-        sagemaker_session = Session(boto_session=boto_session)
-        with load_run(sagemaker_session=sagemaker_session) as run:
+        with load_run() as run:
             run.log_parameters({"p3": 3.0, "p4": 4})
             run.log_metric("test-job-load-log-metric", 0.1)
 

tests/integ/sagemaker/remote_function/test_executor.py

@@ -12,10 +12,6 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
-import boto3
-import os
-
-from sagemaker import Session
 from sagemaker.experiments.trial_component import _TrialComponent
 from sagemaker.remote_function import RemoteExecutor
 from sagemaker.remote_function.client import get_future, list_futures
@@ -85,17 +81,13 @@ def test_executor_submit_with_run_inside(
     sagemaker_session, dummy_container_without_error, cpu_instance_type
 ):
     def square(x):
-        boto_session = boto3.Session(region_name=os.environ["AWS_REGION"])
-        sm_session = Session(boto_session=boto_session)
-        with load_run(sagemaker_session=sm_session) as run:
+        with load_run() as run:
             result = x * x
             run.log_metric("x", result)
         return result
 
     def cube(x):
-        boto_session = boto3.Session(region_name=os.environ["AWS_REGION"])
-        sm_session = Session(boto_session=boto_session)
-        with load_run(sagemaker_session=sm_session) as run:
+        with load_run() as run:
             result = x * x * x
             run.log_metric("x", result)
         return result
@@ -148,17 +140,13 @@ def test_executor_submit_with_run_outside(
     sagemaker_session, dummy_container_without_error, cpu_instance_type
 ):
     def square(x):
-        boto_session = boto3.Session(region_name=os.environ["AWS_REGION"])
-        sm_session = Session(boto_session=boto_session)
-        with load_run(sagemaker_session=sm_session) as run:
+        with load_run() as run:
             result = x * x
             run.log_metric("x", result)
         return result
 
     def cube(x):
-        boto_session = boto3.Session(region_name=os.environ["AWS_REGION"])
-        sm_session = Session(boto_session=boto_session)
-        with load_run(sagemaker_session=sm_session) as run:
+        with load_run() as run:
             result = x * x * x
             run.log_metric("x", result)
         return result
@@ -209,9 +197,7 @@ def cube(x):
 
 def test_executor_map_with_run(sagemaker_session, dummy_container_without_error, cpu_instance_type):
     def square(x):
-        boto_session = boto3.Session(region_name=os.environ["AWS_REGION"])
-        sm_session = Session(boto_session=boto_session)
-        with load_run(sagemaker_session=sm_session) as run:
+        with load_run() as run:
             result = x * x
             run.log_metric("x", result)
         return result