Skip to content

Commit d52d42a

Browse files
yzhu0qidewenwhen
yzhu0
authored and
qidewenwhen
committed
feature: add RunName to expeirment_config (aws#696)
1 parent fd0f4ff commit d52d42a

17 files changed

+61
-15
lines changed

‎src/sagemaker/amazon/amazon_estimator.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -242,8 +242,8 @@ def fit(
242242
generates a default job name, based on the training image name
243243
and current timestamp.
244244
experiment_config (dict[str, str]): Experiment management configuration.
245-
Optionally, the dict can contain three keys:
246-
'ExperimentName', 'TrialName', and 'TrialComponentDisplayName'.
245+
Optionally, the dict can contain four keys:
246+
'ExperimentName', 'TrialName', 'TrialComponentDisplayName' and 'RunName'.
247247
The behavior of setting these keys is as follows:
248248
* If `ExperimentName` is supplied but `TrialName` is not a Trial will be
249249
automatically created and the job's Trial Component associated with the Trial.

‎src/sagemaker/estimator.py

+8-6
Original file line numberDiff line numberDiff line change
@@ -1103,8 +1103,8 @@ def fit(
11031103
job_name (str): Training job name. If not specified, the estimator generates
11041104
a default job name based on the training image name and current timestamp.
11051105
experiment_config (dict[str, str]): Experiment management configuration.
1106-
Optionally, the dict can contain three keys:
1107-
'ExperimentName', 'TrialName', and 'TrialComponentDisplayName'.
1106+
Optionally, the dict can contain four keys:
1107+
'ExperimentName', 'TrialName', 'TrialComponentDisplayName' and 'RunName'..
11081108
The behavior of setting these keys is as follows:
11091109
* If `ExperimentName` is supplied but `TrialName` is not a Trial will be
11101110
automatically created and the job's Trial Component associated with the Trial.
@@ -2023,8 +2023,8 @@ def start_new(cls, estimator, inputs, experiment_config):
20232023
inputs (str): Parameters used when called
20242024
:meth:`~sagemaker.estimator.EstimatorBase.fit`.
20252025
experiment_config (dict[str, str]): Experiment management configuration.
2026-
Optionally, the dict can contain three keys:
2027-
'ExperimentName', 'TrialName', and 'TrialComponentDisplayName'.
2026+
Optionally, the dict can contain four keys:
2027+
'ExperimentName', 'TrialName', 'TrialComponentDisplayName' and 'RunName'.
20282028
The behavior of setting these keys is as follows:
20292029
* If `ExperimentName` is supplied but `TrialName` is not a Trial will be
20302030
automatically created and the job's Trial Component associated with the Trial.
@@ -2033,6 +2033,7 @@ def start_new(cls, estimator, inputs, experiment_config):
20332033
* If both `ExperimentName` and `TrialName` are not supplied the trial component
20342034
will be unassociated.
20352035
* `TrialComponentDisplayName` is used for display in Studio.
2036+
* `RunName` is used to record an experiment run.
20362037
Returns:
20372038
sagemaker.estimator._TrainingJob: Constructed object that captures
20382039
all information about the started training job.
@@ -2053,8 +2054,8 @@ def _get_train_args(cls, estimator, inputs, experiment_config):
20532054
inputs (str): Parameters used when called
20542055
:meth:`~sagemaker.estimator.EstimatorBase.fit`.
20552056
experiment_config (dict[str, str]): Experiment management configuration.
2056-
Optionally, the dict can contain three keys:
2057-
'ExperimentName', 'TrialName', and 'TrialComponentDisplayName'.
2057+
Optionally, the dict can contain four keys:
2058+
'ExperimentName', 'TrialName', 'TrialComponentDisplayName' and 'RunName'.
20582059
The behavior of setting these keys is as follows:
20592060
* If `ExperimentName` is supplied but `TrialName` is not a Trial will be
20602061
automatically created and the job's Trial Component associated with the Trial.
@@ -2063,6 +2064,7 @@ def _get_train_args(cls, estimator, inputs, experiment_config):
20632064
* If both `ExperimentName` and `TrialName` are not supplied the trial component
20642065
will be unassociated.
20652066
* `TrialComponentDisplayName` is used for display in Studio.
2067+
* `RunName` is used to record an experiment run.
20662068
20672069
Returns:
20682070
Dict: dict for `sagemaker.session.Session.train` method

‎src/sagemaker/session.py

+6-4
Original file line numberDiff line numberDiff line change
@@ -548,8 +548,8 @@ def train( # noqa: C901
548548
checkpoints will be provided under `/opt/ml/checkpoints/`.
549549
(default: ``None``).
550550
experiment_config (dict[str, str]): Experiment management configuration.
551-
Optionally, the dict can contain three keys:
552-
'ExperimentName', 'TrialName', and 'TrialComponentDisplayName'.
551+
Optionally, the dict can contain four keys:
552+
'ExperimentName', 'TrialName', 'TrialComponentDisplayName' and 'RunName'.
553553
The behavior of setting these keys is as follows:
554554
* If `ExperimentName` is supplied but `TrialName` is not a Trial will be
555555
automatically created and the job's Trial Component associated with the Trial.
@@ -558,6 +558,7 @@ def train( # noqa: C901
558558
* If both `ExperimentName` and `TrialName` are not supplied the trial component
559559
will be unassociated.
560560
* `TrialComponentDisplayName` is used for display in Studio.
561+
* `RunName` is used to record an experiment run.
561562
enable_sagemaker_metrics (bool): enable SageMaker Metrics Time
562563
Series. For more information see:
563564
https://docs.aws.amazon.com/sagemaker/latest/dg/API_AlgorithmSpecification.html#SageMaker-Type-AlgorithmSpecification-EnableSageMakerMetricsTimeSeries
@@ -703,8 +704,8 @@ def _get_train_request( # noqa: C901
703704
checkpoints will be provided under `/opt/ml/checkpoints/`.
704705
(default: ``None``).
705706
experiment_config (dict[str, str]): Experiment management configuration.
706-
Optionally, the dict can contain three keys:
707-
'ExperimentName', 'TrialName', and 'TrialComponentDisplayName'.
707+
Optionally, the dict can contain four keys:
708+
'ExperimentName', 'TrialName', 'TrialComponentDisplayName' and 'RunName'.
708709
The behavior of setting these keys is as follows:
709710
* If `ExperimentName` is supplied but `TrialName` is not a Trial will be
710711
automatically created and the job's Trial Component associated with the Trial.
@@ -713,6 +714,7 @@ def _get_train_request( # noqa: C901
713714
* If both `ExperimentName` and `TrialName` are not supplied the trial component
714715
will be unassociated.
715716
* `TrialComponentDisplayName` is used for display in Studio.
717+
* `RunName` is used to record an experiment run.
716718
enable_sagemaker_metrics (bool): enable SageMaker Metrics Time
717719
Series. For more information see:
718720
https://docs.aws.amazon.com/sagemaker/latest/dg/API_AlgorithmSpecification.html#SageMaker-Type-AlgorithmSpecification-EnableSageMakerMetricsTimeSeries

‎tests/integ/sagemaker/experiments/test_run.py

+18
Original file line numberDiff line numberDiff line change
@@ -167,6 +167,12 @@ def validate_tc_artifact_association(is_output, expected_artifact_name):
167167
_PYTHON_SCRIPT_PATH = os.path.join(DATA_DIR, "experiment/scripts/launcher.sh")
168168

169169

170+
@pytest.mark.skip(
171+
reason=(
172+
"Waiting for the CR https://code.amazon.com/reviews/CR-75915367/revisions/1#/details "
173+
"to deploy to us-west-2"
174+
)
175+
)
170176
def test_run_from_local_and_train_job_and_all_exp_cfg_match(sagemaker_session, job_resource_dir):
171177
# Notes:
172178
# 1. The 1st Run TC created locally
@@ -208,6 +214,12 @@ def test_run_from_local_and_train_job_and_all_exp_cfg_match(sagemaker_session, j
208214
)
209215

210216

217+
@pytest.mark.skip(
218+
reason=(
219+
"Waiting for the CR https://code.amazon.com/reviews/CR-75915367/revisions/1#/details "
220+
"to deploy to us-west-2"
221+
)
222+
)
211223
def test_run_from_local_and_train_job_and_exp_cfg_not_match(sagemaker_session, job_resource_dir):
212224
# Notes:
213225
# 1. The 1st Run TC created locally
@@ -249,6 +261,12 @@ def test_run_from_local_and_train_job_and_exp_cfg_not_match(sagemaker_session, j
249261
)
250262

251263

264+
@pytest.mark.skip(
265+
reason=(
266+
"Waiting for the CR https://code.amazon.com/reviews/CR-75915367/revisions/1#/details "
267+
"to deploy to us-west-2"
268+
)
269+
)
252270
def test_run_from_train_job(sagemaker_session, job_resource_dir):
253271
# Notes:
254272
# 1. No Run TC created locally or specified in experiment config

‎tests/unit/sagemaker/huggingface/test_estimator.py

+1
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
"ExperimentName": "exp",
4949
"TrialName": "trial",
5050
"TrialComponentDisplayName": "tc",
51+
"RunName": "rn",
5152
}
5253

5354

‎tests/unit/sagemaker/tensorflow/test_estimator.py

+1
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@
5656
"ExperimentName": "exp",
5757
"TrialName": "trial",
5858
"TrialComponentDisplayName": "tc",
59+
"RunName": "rn",
5960
}
6061

6162

‎tests/unit/sagemaker/training_compiler/test_huggingface_pytorch_compiler.py

+1
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@
5252
"ExperimentName": "exp",
5353
"TrialName": "trial",
5454
"TrialComponentDisplayName": "tc",
55+
"RunName": "rn",
5556
}
5657

5758

‎tests/unit/sagemaker/training_compiler/test_huggingface_tensorflow_compiler.py

+1
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@
5050
"ExperimentName": "exp",
5151
"TrialName": "trial",
5252
"TrialComponentDisplayName": "tc",
53+
"RunName": "rn",
5354
}
5455

5556

‎tests/unit/sagemaker/training_compiler/test_tensorflow_compiler.py

+1
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@
5050
"ExperimentName": "exp",
5151
"TrialName": "trial",
5252
"TrialComponentDisplayName": "tc",
53+
"RunName": "rn",
5354
}
5455

5556

‎tests/unit/test_amazon_estimator.py

+8-2
Original file line numberDiff line numberDiff line change
@@ -239,12 +239,18 @@ def test_fit_pass_experiment_config(sagemaker_session):
239239
labels = [99, 85, 87, 2]
240240
pca.fit(
241241
pca.record_set(np.array(train), np.array(labels)),
242-
experiment_config={"ExperimentName": "exp"},
242+
experiment_config={
243+
"ExperimentName": "exp",
244+
"RunName": "rn",
245+
},
243246
)
244247

245248
called_args = sagemaker_session.train.call_args
246249

247-
assert called_args[1]["experiment_config"] == {"ExperimentName": "exp"}
250+
assert called_args[1]["experiment_config"] == {
251+
"ExperimentName": "exp",
252+
"RunName": "rn",
253+
}
248254

249255

250256
def test_build_shards():

‎tests/unit/test_estimator.py

+8-1
Original file line numberDiff line numberDiff line change
@@ -2489,7 +2489,12 @@ def test_start_new(sagemaker_session):
24892489
hyperparameters=hyperparameters,
24902490
)
24912491

2492-
exp_config = {"ExperimentName": "exp", "TrialName": "t", "TrialComponentDisplayName": "tc"}
2492+
exp_config = {
2493+
"ExperimentName": "exp",
2494+
"TrialName": "t",
2495+
"TrialComponentDisplayName": "tc",
2496+
"RunName": "rn",
2497+
}
24932498

24942499
started_training_job = training_job.start_new(estimator, inputs, experiment_config=exp_config)
24952500
called_args = sagemaker_session.train.call_args
@@ -2680,6 +2685,7 @@ def test_unsupported_type_in_dict():
26802685
"ExperimentName": "exp",
26812686
"TrialName": "trial",
26822687
"TrialComponentDisplayName": "tc",
2688+
"RunName": "rn",
26832689
}
26842690
}
26852691
)
@@ -2884,6 +2890,7 @@ def test_generic_to_fit_with_experiment_config(time, sagemaker_session):
28842890
"ExperimentName": "exp",
28852891
"TrialName": "trial",
28862892
"TrialComponentDisplayName": "tc",
2893+
"RunName": "rn",
28872894
},
28882895
)
28892896

‎tests/unit/test_mxnet.py

+1
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@
6262
"ExperimentName": "exp",
6363
"TrialName": "trial",
6464
"TrialComponentDisplayName": "tc",
65+
"RunName": "rn",
6566
}
6667

6768
MODEL_PKG_RESPONSE = {"ModelPackageArn": "arn:model-pkg-arn"}

‎tests/unit/test_pytorch.py

+1
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@
5454
"ExperimentName": "exp",
5555
"TrialName": "trial",
5656
"TrialComponentDisplayName": "tc",
57+
"RunName": "rn",
5758
}
5859

5960
DISTRIBUTION_PYTORCH_DDP_ENABLED = {"pytorchddp": {"enabled": True}}

‎tests/unit/test_rl.py

+1
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@
4949
"ExperimentName": "exp",
5050
"TrialName": "trial",
5151
"TrialComponentDisplayName": "tc",
52+
"RunName": "rn",
5253
}
5354

5455

‎tests/unit/test_session.py

+1
Original file line numberDiff line numberDiff line change
@@ -700,6 +700,7 @@ def test_training_input_all_arguments():
700700
"ExperimentName": "dummyExp",
701701
"TrialName": "dummyT",
702702
"TrialComponentDisplayName": "dummyTC",
703+
"RunName": "dummyRN",
703704
}
704705
MODEL_CLIENT_CONFIG = {"InvocationsMaxRetries": 2, "InvocationsTimeoutInSeconds": 60}
705706

‎tests/unit/test_sklearn.py

+1
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@
5151
"ExperimentName": "exp",
5252
"TrialName": "trial",
5353
"TrialComponentDisplayName": "tc",
54+
"RunName": "rn",
5455
}
5556

5657

‎tests/unit/test_xgboost.py

+1
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@
5454
"ExperimentName": "exp",
5555
"TrialName": "trial",
5656
"TrialComponentDisplayName": "tc",
57+
"RunName": "rn",
5758
}
5859

5960

0 commit comments

Comments
 (0)