diff --git a/tests/integ/sagemaker/experiments/helpers.py b/tests/integ/sagemaker/experiments/helpers.py index 9a22c3a30c..656cccd8dc 100644 --- a/tests/integ/sagemaker/experiments/helpers.py +++ b/tests/integ/sagemaker/experiments/helpers.py @@ -13,9 +13,12 @@ from __future__ import absolute_import from contextlib import contextmanager +import pytest +import logging from sagemaker import utils from sagemaker.experiments.experiment import Experiment +from sagemaker.experiments._run_context import _RunContext EXP_INTEG_TEST_NAME_PREFIX = "experiments-integ" @@ -40,3 +43,16 @@ def cleanup_exp_resources(exp_names, sagemaker_session): for exp_name in exp_names: exp = Experiment.load(experiment_name=exp_name, sagemaker_session=sagemaker_session) exp._delete_all(action="--force") + +@pytest.fixture +def clear_run_context(): + current_run = _RunContext.get_current_run() + if current_run == None: + return + + logging.info( + f"RunContext already populated by run {current_run.run_name}" + f" in experiment {current_run.experiment_name}." + " Clearing context manually" + ) + _RunContext.drop_current_run() \ No newline at end of file diff --git a/tests/integ/sagemaker/experiments/test_run.py b/tests/integ/sagemaker/experiments/test_run.py index 4f59d11c54..57d3ef41d4 100644 --- a/tests/integ/sagemaker/experiments/test_run.py +++ b/tests/integ/sagemaker/experiments/test_run.py @@ -32,7 +32,7 @@ from sagemaker.experiments.trial_component import _TrialComponent from sagemaker.sklearn import SKLearn from sagemaker.utils import retry_with_backoff, unique_name_from_base -from tests.integ.sagemaker.experiments.helpers import name, cleanup_exp_resources +from tests.integ.sagemaker.experiments.helpers import name, cleanup_exp_resources, clear_run_context from sagemaker.experiments.run import ( RUN_NAME_BASE, DELIMITER, @@ -55,7 +55,7 @@ def artifact_file_path(tempdir): metric_name = "Test-Local-Init-Log-Metric" -def test_local_run_with_load(sagemaker_session, artifact_file_path): +def test_local_run_with_load(sagemaker_session, artifact_file_path, clear_run_context): exp_name = f"My-Local-Exp-{name()}" with cleanup_exp_resources(exp_names=[exp_name], sagemaker_session=sagemaker_session): # Run name is not provided, will create a new TC @@ -86,7 +86,9 @@ def verify_load_run(): retry_with_backoff(verify_load_run, 4) -def test_two_local_run_init_with_same_run_name_and_different_exp_names(sagemaker_session): +def test_two_local_run_init_with_same_run_name_and_different_exp_names( + sagemaker_session, clear_run_context +): exp_name1 = f"my-two-local-exp1-{name()}" exp_name2 = f"my-two-local-exp2-{name()}" run_name = "test-run" @@ -124,7 +126,9 @@ def test_two_local_run_init_with_same_run_name_and_different_exp_names(sagemaker ("my-test4", "test-run", "run-display-name-test"), # with supplied display name ], ) -def test_run_name_vs_trial_component_name_edge_cases(sagemaker_session, input_names): +def test_run_name_vs_trial_component_name_edge_cases( + sagemaker_session, input_names, clear_run_context +): exp_name, run_name, run_display_name = input_names with cleanup_exp_resources(exp_names=[exp_name], sagemaker_session=sagemaker_session): with Run( @@ -177,6 +181,7 @@ def test_run_from_local_and_train_job_and_all_exp_cfg_match( execution_role, sagemaker_client_config, sagemaker_metrics_config, + clear_run_context, ): # Notes: # 1. The 1st Run created locally and its exp config was auto passed to the job @@ -277,6 +282,7 @@ def test_run_from_local_and_train_job_and_exp_cfg_not_match( execution_role, sagemaker_client_config, sagemaker_metrics_config, + clear_run_context, ): # Notes: # 1. The 1st Run created locally and its exp config was auto passed to the job @@ -363,6 +369,7 @@ def test_run_from_train_job_only( execution_role, sagemaker_client_config, sagemaker_metrics_config, + clear_run_context, ): # Notes: # 1. No Run created locally or specified in experiment config @@ -413,6 +420,7 @@ def test_run_from_processing_job_and_override_default_exp_config( execution_role, sagemaker_client_config, sagemaker_metrics_config, + clear_run_context, ): # Notes: # 1. The 1st Run (run) created locally @@ -492,6 +500,7 @@ def test_run_from_transform_job( execution_role, sagemaker_client_config, sagemaker_metrics_config, + clear_run_context, ): # Notes: # 1. The 1st Run (run) created locally @@ -573,6 +582,7 @@ def test_load_run_auto_pass_in_exp_config_to_job( execution_role, sagemaker_client_config, sagemaker_metrics_config, + clear_run_context, ): # Notes: # 1. In local side, load the Run created previously and invoke a job under the load context @@ -621,7 +631,7 @@ def test_load_run_auto_pass_in_exp_config_to_job( ) -def test_list(run_obj, sagemaker_session): +def test_list(run_obj, sagemaker_session, clear_run_context): tc1 = _TrialComponent.create( trial_component_name=f"non-run-tc1-{name()}", sagemaker_session=sagemaker_session, @@ -643,7 +653,7 @@ def test_list(run_obj, sagemaker_session): assert run_tcs[0].experiment_config == run_obj.experiment_config -def test_list_twice(run_obj, sagemaker_session): +def test_list_twice(run_obj, sagemaker_session, clear_run_context): tc1 = _TrialComponent.create( trial_component_name=f"non-run-tc1-{name()}", sagemaker_session=sagemaker_session,