feat: Enable customizing artifact output path (#3965)

Author: rd-pong

src/sagemaker/experiments/run.py

@@ -29,6 +29,7 @@
 from sagemaker.experiments._helper import (
     _ArtifactUploader,
     _LineageArtifactTracker,
+    _DEFAULT_ARTIFACT_PREFIX,
 )
 from sagemaker.experiments._environment import _RunEnvironment
 from sagemaker.experiments._run_context import _RunContext
@@ -95,6 +96,8 @@ def __init__(
         run_display_name: Optional[str] = None,
         tags: Optional[List[Dict[str, str]]] = None,
         sagemaker_session: Optional["Session"] = None,
+        artifact_bucket: Optional[str] = None,
+        artifact_prefix: Optional[str] = None,
     ):
         """Construct a `Run` instance.
 
@@ -152,6 +155,11 @@ def __init__(
                 manages interactions with Amazon SageMaker APIs and any other
                 AWS services needed. If not specified, one is created using the
                 default AWS configuration chain.
+            artifact_bucket (str): The S3 bucket to upload the artifact to.
+                If not specified, the default bucket defined in `sagemaker_session`
+                will be used.
+            artifact_prefix (str): The S3 key prefix used to generate the S3 path
+                to upload the artifact to (default: "trial-component-artifacts").
         """
         # TODO: we should revert the lower casting once backend fix reaches prod
         self.experiment_name = experiment_name.lower()
@@ -197,6 +205,10 @@ def __init__(
         self._artifact_uploader = _ArtifactUploader(
             trial_component_name=self._trial_component.trial_component_name,
             sagemaker_session=sagemaker_session,
+            artifact_bucket=artifact_bucket,
+            artifact_prefix=_DEFAULT_ARTIFACT_PREFIX
+            if artifact_prefix is None
+            else artifact_prefix,
         )
         self._lineage_artifact_tracker = _LineageArtifactTracker(
             trial_component_arn=self._trial_component.trial_component_arn,
@@ -729,6 +741,8 @@ def load_run(
     run_name: Optional[str] = None,
     experiment_name: Optional[str] = None,
     sagemaker_session: Optional["Session"] = None,
+    artifact_bucket: Optional[str] = None,
+    artifact_prefix: Optional[str] = None,
 ) -> Run:
     """Load an existing run.
 
@@ -792,6 +806,11 @@ def load_run(
             manages interactions with Amazon SageMaker APIs and any other
             AWS services needed. If not specified, one is created using the
             default AWS configuration chain.
+        artifact_bucket (str): The S3 bucket to upload the artifact to.
+                If not specified, the default bucket defined in `sagemaker_session`
+                will be used.
+        artifact_prefix (str): The S3 key prefix used to generate the S3 path
+            to upload the artifact to (default: "trial-component-artifacts").
 
     Returns:
         Run: The loaded Run object.
@@ -811,6 +830,8 @@ def load_run(
             experiment_name=experiment_name,
             run_name=run_name,
             sagemaker_session=sagemaker_session or _utils.default_session(),
+            artifact_bucket=artifact_bucket,
+            artifact_prefix=artifact_prefix,
         )
     elif _RunContext.get_current_run():
         run_instance = _RunContext.get_current_run()
@@ -827,6 +848,8 @@ def load_run(
             experiment_name=experiment_name,
             run_name=run_name,
             sagemaker_session=sagemaker_session or _utils.default_session(),
+            artifact_bucket=artifact_bucket,
+            artifact_prefix=artifact_prefix,
         )
     else:
         raise RuntimeError(

tests/unit/sagemaker/experiments/helpers.py

@@ -22,6 +22,8 @@
 TEST_EXP_DISPLAY_NAME = "my-experiment-display-name"
 TEST_RUN_DISPLAY_NAME = "my-run-display-name"
 TEST_TAGS = [{"Key": "some-key", "Value": "some-value"}]
+TEST_ARTIFACT_BUCKET = "my-artifact-bucket"
+TEST_ARTIFACT_PREFIX = "my-artifact-prefix"
 
 
 def mock_tc_load_or_create_func(

tests/unit/sagemaker/experiments/test_run.py

@@ -44,6 +44,7 @@
 from sagemaker.experiments import Run, load_run, list_runs
 from sagemaker.experiments.trial import _Trial
 from sagemaker.experiments.trial_component import _TrialComponent
+from sagemaker.experiments._helper import _DEFAULT_ARTIFACT_PREFIX
 from tests.unit.sagemaker.experiments.helpers import (
     mock_trial_load_or_create_func,
     mock_tc_load_or_create_func,
@@ -52,9 +53,25 @@
     TEST_RUN_NAME,
     TEST_EXP_DISPLAY_NAME,
     TEST_RUN_DISPLAY_NAME,
+    TEST_ARTIFACT_BUCKET,
+    TEST_ARTIFACT_PREFIX,
 )
 
 
+@pytest.mark.parametrize(
+    ("kwargs", "expected_artifact_bucket", "expected_artifact_prefix"),
+    [
+        ({}, None, _DEFAULT_ARTIFACT_PREFIX),
+        (
+            {
+                "artifact_bucket": TEST_ARTIFACT_BUCKET,
+                "artifact_prefix": TEST_ARTIFACT_PREFIX,
+            },
+            TEST_ARTIFACT_BUCKET,
+            TEST_ARTIFACT_PREFIX,
+        ),
+    ],
+)
 @patch(
     "sagemaker.experiments.run.Experiment._load_or_create",
     MagicMock(return_value=Experiment(experiment_name=TEST_EXP_NAME)),
@@ -69,9 +86,18 @@
     MagicMock(side_effect=mock_tc_load_or_create_func),
 )
 @patch.object(_TrialComponent, "save")
-def test_run_init(mock_tc_save, sagemaker_session):
+def test_run_init(
+    mock_tc_save,
+    sagemaker_session,
+    kwargs,
+    expected_artifact_bucket,
+    expected_artifact_prefix,
+):
     with Run(
-        experiment_name=TEST_EXP_NAME, run_name=TEST_RUN_NAME, sagemaker_session=sagemaker_session
+        experiment_name=TEST_EXP_NAME,
+        run_name=TEST_RUN_NAME,
+        sagemaker_session=sagemaker_session,
+        **kwargs,
     ) as run_obj:
         assert not run_obj._in_load
         assert not run_obj._inside_load_context
@@ -90,6 +116,8 @@ def test_run_init(mock_tc_save, sagemaker_session):
             TRIAL_NAME: run_obj.run_group_name,
             RUN_NAME: expected_tc_name,
         }
+        assert run_obj._artifact_uploader.artifact_bucket == expected_artifact_bucket
+        assert run_obj._artifact_uploader.artifact_prefix == expected_artifact_prefix
 
     # trail_component.save is called when entering/ exiting the with block
     mock_tc_save.assert_called()
@@ -124,6 +152,20 @@ def test_run_init_name_length_exceed_limit(sagemaker_session):
     )
 
 
+@pytest.mark.parametrize(
+    ("kwargs", "expected_artifact_bucket", "expected_artifact_prefix"),
+    [
+        ({}, None, _DEFAULT_ARTIFACT_PREFIX),
+        (
+            {
+                "artifact_bucket": TEST_ARTIFACT_BUCKET,
+                "artifact_prefix": TEST_ARTIFACT_PREFIX,
+            },
+            TEST_ARTIFACT_BUCKET,
+            TEST_ARTIFACT_PREFIX,
+        ),
+    ],
+)
 @patch.object(_TrialComponent, "save", MagicMock(return_value=None))
 @patch(
     "sagemaker.experiments.run.Experiment._load_or_create",
@@ -139,7 +181,13 @@ def test_run_init_name_length_exceed_limit(sagemaker_session):
     MagicMock(side_effect=mock_tc_load_or_create_func),
 )
 @patch("sagemaker.experiments.run._RunEnvironment")
-def test_run_load_no_run_name_and_in_train_job(mock_run_env, sagemaker_session):
+def test_run_load_no_run_name_and_in_train_job(
+    mock_run_env,
+    sagemaker_session,
+    kwargs,
+    expected_artifact_bucket,
+    expected_artifact_prefix,
+):
     client = sagemaker_session.sagemaker_client
     job_name = "my-train-job"
     rv = Mock()
@@ -158,7 +206,7 @@ def test_run_load_no_run_name_and_in_train_job(mock_run_env, sagemaker_session):
         # The Run object has been created else where
         "ExperimentConfig": exp_config,
     }
-    with load_run(sagemaker_session=sagemaker_session) as run_obj:
+    with load_run(sagemaker_session=sagemaker_session, **kwargs) as run_obj:
         assert run_obj._in_load
         assert not run_obj._inside_init_context
         assert run_obj._inside_load_context
@@ -169,6 +217,8 @@ def test_run_load_no_run_name_and_in_train_job(mock_run_env, sagemaker_session):
         assert run_obj.experiment_name == TEST_EXP_NAME
         assert run_obj._experiment
         assert run_obj.experiment_config == exp_config
+        assert run_obj._artifact_uploader.artifact_bucket == expected_artifact_bucket
+        assert run_obj._artifact_uploader.artifact_prefix == expected_artifact_prefix
 
     client.describe_training_job.assert_called_once_with(TrainingJobName=job_name)
 
@@ -215,6 +265,20 @@ def test_run_load_no_run_name_and_not_in_train_job_but_no_obj_in_context(sagemak
     assert "Failed to load a Run object" in str(err)
 
 
+@pytest.mark.parametrize(
+    ("kwargs", "expected_artifact_bucket", "expected_artifact_prefix"),
+    [
+        ({}, None, _DEFAULT_ARTIFACT_PREFIX),
+        (
+            {
+                "artifact_bucket": TEST_ARTIFACT_BUCKET,
+                "artifact_prefix": TEST_ARTIFACT_PREFIX,
+            },
+            TEST_ARTIFACT_BUCKET,
+            TEST_ARTIFACT_PREFIX,
+        ),
+    ],
+)
 @patch.object(_TrialComponent, "save", MagicMock(return_value=None))
 @patch(
     "sagemaker.experiments.run.Experiment._load_or_create",
@@ -229,11 +293,14 @@ def test_run_load_no_run_name_and_not_in_train_job_but_no_obj_in_context(sagemak
     "sagemaker.experiments.run._TrialComponent._load_or_create",
     MagicMock(side_effect=mock_tc_load_or_create_func),
 )
-def test_run_load_with_run_name_and_exp_name(sagemaker_session):
+def test_run_load_with_run_name_and_exp_name(
+    sagemaker_session, kwargs, expected_artifact_bucket, expected_artifact_prefix
+):
     with load_run(
         run_name=TEST_RUN_NAME,
         experiment_name=TEST_EXP_NAME,
         sagemaker_session=sagemaker_session,
+        **kwargs,
     ) as run_obj:
         expected_tc_name = f"{TEST_EXP_NAME}{DELIMITER}{TEST_RUN_NAME}"
         expected_exp_config = {
@@ -249,6 +316,8 @@ def test_run_load_with_run_name_and_exp_name(sagemaker_session):
         assert run_obj._trial
         assert run_obj._experiment
         assert run_obj.experiment_config == expected_exp_config
+        assert run_obj._artifact_uploader.artifact_bucket == expected_artifact_bucket
+        assert run_obj._artifact_uploader.artifact_prefix == expected_artifact_prefix
 
 
 def test_run_load_with_run_name_but_no_exp_name(sagemaker_session):