From 8e79b7e320761930990b1ea916482d37d8a00e16 Mon Sep 17 00:00:00 2001 From: Dewen Qi Date: Thu, 19 May 2022 13:34:16 -0700 Subject: [PATCH] Revert "fix: Prevent passing PipelineVariable object into image_uris.retrieve (#3054)" This reverts commit edd246342751df50867eecd45ca0138318c10ca5. --- src/sagemaker/image_uris.py | 9 +---- src/sagemaker/workflow/entities.py | 6 +--- .../test_model_create_and_registration.py | 33 ++++++++----------- .../sagemaker/workflow/test_model_steps.py | 30 ++++++++--------- tests/integ/sagemaker/workflow/test_retry.py | 11 ++++--- .../sagemaker/workflow/test_training_steps.py | 6 ++-- .../sagemaker/workflow/test_tuning_steps.py | 14 +++----- .../integ/sagemaker/workflow/test_workflow.py | 33 +++++++++---------- .../sagemaker/image_uris/test_retrieve.py | 17 ---------- .../workflow/test_execution_variables.py | 2 +- .../unit/sagemaker/workflow/test_functions.py | 4 +-- .../sagemaker/workflow/test_parameters.py | 2 +- .../sagemaker/workflow/test_properties.py | 2 +- 13 files changed, 62 insertions(+), 107 deletions(-) diff --git a/src/sagemaker/image_uris.py b/src/sagemaker/image_uris.py index a161557216..2428b1ca85 100644 --- a/src/sagemaker/image_uris.py +++ b/src/sagemaker/image_uris.py @@ -23,7 +23,6 @@ from sagemaker.jumpstart.utils import is_jumpstart_model_input from sagemaker.spark import defaults from sagemaker.jumpstart import artifacts -from sagemaker.workflow import is_pipeline_variable logger = logging.getLogger(__name__) @@ -105,17 +104,11 @@ def retrieve( Raises: NotImplementedError: If the scope is not supported. - ValueError: If the combination of arguments specified is not supported or - any PipelineVariable object is passed in. + ValueError: If the combination of arguments specified is not supported. VulnerableJumpStartModelError: If any of the dependencies required by the script have known security vulnerabilities. DeprecatedJumpStartModelError: If the version of the model is deprecated. """ - args = dict(locals()) - for name, val in args.items(): - if is_pipeline_variable(val): - raise ValueError("%s should not be a pipeline variable (%s)" % (name, type(val))) - if is_jumpstart_model_input(model_id, model_version): return artifacts._retrieve_image_uri( model_id, diff --git a/src/sagemaker/workflow/entities.py b/src/sagemaker/workflow/entities.py index 2880e9e6c9..f530c9b669 100644 --- a/src/sagemaker/workflow/entities.py +++ b/src/sagemaker/workflow/entities.py @@ -78,11 +78,7 @@ def __add__(self, other: Union[Expression, PrimitiveType]): def __str__(self): """Override built-in String function for PipelineVariable""" - raise TypeError( - "Pipeline variables do not support __str__ operation. " - "Please use `.to_string()` to convert it to string type in execution time" - "or use `.expr` to translate it to Json for display purpose in Python SDK." - ) + raise TypeError("Pipeline variables do not support __str__ operation.") def __int__(self): """Override built-in Integer function for PipelineVariable""" diff --git a/tests/integ/sagemaker/workflow/test_model_create_and_registration.py b/tests/integ/sagemaker/workflow/test_model_create_and_registration.py index 0b5a0320f9..eb7bc80abe 100644 --- a/tests/integ/sagemaker/workflow/test_model_create_and_registration.py +++ b/tests/integ/sagemaker/workflow/test_model_create_and_registration.py @@ -90,12 +90,10 @@ def test_conditional_pytorch_training_model_registration( inputs = TrainingInput(s3_data=input_path) instance_count = ParameterInteger(name="InstanceCount", default_value=1) - instance_type = "ml.m5.xlarge" + instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge") good_enough_input = ParameterInteger(name="GoodEnoughInput", default_value=1) in_condition_input = ParameterString(name="Foo", default_value="Foo") - # If image_uri is not provided, the instance_type should not be a pipeline variable - # since instance_type is used to retrieve image_uri in compile time (PySDK) pytorch_estimator = PyTorch( entry_point=entry_point, role=role, @@ -155,6 +153,7 @@ def test_conditional_pytorch_training_model_registration( in_condition_input, good_enough_input, instance_count, + instance_type, ], steps=[step_train, step_cond], sagemaker_session=sagemaker_session, @@ -260,10 +259,8 @@ def test_sklearn_xgboost_sip_model_registration( prefix = "sip" bucket_name = sagemaker_session.default_bucket() instance_count = ParameterInteger(name="InstanceCount", default_value=1) - instance_type = "ml.m5.xlarge" + instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge") - # The instance_type should not be a pipeline variable - # since it is used to retrieve image_uri in compile time (PySDK) sklearn_processor = SKLearnProcessor( role=role, instance_type=instance_type, @@ -334,8 +331,6 @@ def test_sklearn_xgboost_sip_model_registration( source_dir = base_dir code_location = "s3://{0}/{1}/code".format(bucket_name, prefix) - # If image_uri is not provided, the instance_type should not be a pipeline variable - # since instance_type is used to retrieve image_uri in compile time (PySDK) estimator = XGBoost( entry_point=entry_point, source_dir=source_dir, @@ -421,6 +416,7 @@ def test_sklearn_xgboost_sip_model_registration( train_data_path_param, val_data_path_param, model_path_param, + instance_type, instance_count, output_path_param, ], @@ -466,7 +462,7 @@ def test_model_registration_with_drift_check_baselines( pipeline_name, ): instance_count = ParameterInteger(name="InstanceCount", default_value=1) - instance_type = "ml.m5.xlarge" + instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge") # upload model data to s3 model_local_path = os.path.join(DATA_DIR, "mxnet_mnist/model.tar.gz") @@ -554,9 +550,6 @@ def test_model_registration_with_drift_check_baselines( ), ) customer_metadata_properties = {"key1": "value1"} - - # If image_uri is not provided, the instance_type should not be a pipeline variable - # since instance_type is used to retrieve image_uri in compile time (PySDK) estimator = XGBoost( entry_point="training.py", source_dir=os.path.join(DATA_DIR, "sip"), @@ -586,6 +579,7 @@ def test_model_registration_with_drift_check_baselines( parameters=[ model_uri_param, metrics_uri_param, + instance_type, instance_count, ], steps=[step_register], @@ -673,11 +667,9 @@ def test_model_registration_with_model_repack( inputs = TrainingInput(s3_data=input_path) instance_count = ParameterInteger(name="InstanceCount", default_value=1) - instance_type = "ml.m5.xlarge" + instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge") good_enough_input = ParameterInteger(name="GoodEnoughInput", default_value=1) - # If image_uri is not provided, the instance_type should not be a pipeline variable - # since instance_type is used to retrieve image_uri in compile time (PySDK) pytorch_estimator = PyTorch( entry_point=entry_point, role=role, @@ -732,7 +724,7 @@ def test_model_registration_with_model_repack( pipeline = Pipeline( name=pipeline_name, - parameters=[good_enough_input, instance_count], + parameters=[good_enough_input, instance_count, instance_type], steps=[step_cond], sagemaker_session=sagemaker_session, ) @@ -775,10 +767,8 @@ def test_model_registration_with_tensorflow_model_with_pipeline_model( inputs = TrainingInput(s3_data=input_path) instance_count = ParameterInteger(name="InstanceCount", default_value=1) - instance_type = "ml.m5.xlarge" + instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge") - # If image_uri is not provided, the instance_type should not be a pipeline variable - # since instance_type is used to retrieve image_uri in compile time (PySDK) tensorflow_estimator = TensorFlow( entry_point=entry_point, role=role, @@ -819,7 +809,10 @@ def test_model_registration_with_tensorflow_model_with_pipeline_model( pipeline = Pipeline( name=pipeline_name, - parameters=[instance_count], + parameters=[ + instance_count, + instance_type, + ], steps=[step_train, step_register_model], sagemaker_session=sagemaker_session, ) diff --git a/tests/integ/sagemaker/workflow/test_model_steps.py b/tests/integ/sagemaker/workflow/test_model_steps.py index 2c60e1d7d2..71b07e5404 100644 --- a/tests/integ/sagemaker/workflow/test_model_steps.py +++ b/tests/integ/sagemaker/workflow/test_model_steps.py @@ -83,16 +83,15 @@ def test_pytorch_training_model_registration_and_creation_without_custom_inferen inputs = TrainingInput(s3_data=input_path) instance_count = ParameterInteger(name="InstanceCount", default_value=1) + instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge") - # If image_uri is not provided, the instance_type should not be a pipeline variable - # since instance_type is used to retrieve image_uri in compile time (PySDK) pytorch_estimator = PyTorch( entry_point=entry_point, role=role, framework_version="1.5.0", py_version="py3", instance_count=instance_count, - instance_type="ml.m5.xlarge", + instance_type=instance_type, sagemaker_session=pipeline_session, ) train_step_args = pytorch_estimator.fit(inputs=inputs) @@ -141,7 +140,7 @@ def test_pytorch_training_model_registration_and_creation_without_custom_inferen ) pipeline = Pipeline( name=pipeline_name, - parameters=[instance_count], + parameters=[instance_count, instance_type], steps=[step_train, step_model_regis, step_model_create, step_fail], sagemaker_session=pipeline_session, ) @@ -204,16 +203,15 @@ def test_pytorch_training_model_registration_and_creation_with_custom_inference( inputs = TrainingInput(s3_data=input_path) instance_count = ParameterInteger(name="InstanceCount", default_value=1) + instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge") - # If image_uri is not provided, the instance_type should not be a pipeline variable - # since instance_type is used to retrieve image_uri in compile time (PySDK) pytorch_estimator = PyTorch( entry_point=entry_point, role=role, framework_version="1.5.0", py_version="py3", instance_count=instance_count, - instance_type="ml.m5.xlarge", + instance_type=instance_type, sagemaker_session=pipeline_session, output_kms_key=kms_key, ) @@ -269,7 +267,7 @@ def test_pytorch_training_model_registration_and_creation_with_custom_inference( ) pipeline = Pipeline( name=pipeline_name, - parameters=[instance_count], + parameters=[instance_count, instance_type], steps=[step_train, step_model_regis, step_model_create, step_fail], sagemaker_session=pipeline_session, ) @@ -402,6 +400,7 @@ def test_model_registration_with_drift_check_baselines_and_model_metrics( pipeline_name, ): instance_count = ParameterInteger(name="InstanceCount", default_value=1) + instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge") # upload model data to s3 model_local_path = os.path.join(DATA_DIR, "mxnet_mnist/model.tar.gz") @@ -489,12 +488,10 @@ def test_model_registration_with_drift_check_baselines_and_model_metrics( ), ) customer_metadata_properties = {"key1": "value1"} - # If image_uri is not provided, the instance_type should not be a pipeline variable - # since instance_type is used to retrieve image_uri in compile time (PySDK) estimator = XGBoost( entry_point="training.py", source_dir=os.path.join(DATA_DIR, "sip"), - instance_type="ml.m5.xlarge", + instance_type=instance_type, instance_count=instance_count, framework_version="0.90-2", sagemaker_session=pipeline_session, @@ -527,6 +524,7 @@ def test_model_registration_with_drift_check_baselines_and_model_metrics( parameters=[ model_uri_param, metrics_uri_param, + instance_type, instance_count, ], steps=[step_model_register], @@ -608,14 +606,13 @@ def test_model_registration_with_tensorflow_model_with_pipeline_model( ) inputs = TrainingInput(s3_data=input_path) instance_count = ParameterInteger(name="InstanceCount", default_value=1) + instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge") - # If image_uri is not provided, the instance_type should not be a pipeline variable - # since instance_type is used to retrieve image_uri in compile time (PySDK) tensorflow_estimator = TensorFlow( entry_point=entry_point, role=role, instance_count=instance_count, - instance_type="ml.m5.xlarge", + instance_type=instance_type, framework_version=tf_full_version, py_version=tf_full_py_version, sagemaker_session=pipeline_session, @@ -648,7 +645,10 @@ def test_model_registration_with_tensorflow_model_with_pipeline_model( ) pipeline = Pipeline( name=pipeline_name, - parameters=[instance_count], + parameters=[ + instance_count, + instance_type, + ], steps=[step_train, step_register_model], sagemaker_session=pipeline_session, ) diff --git a/tests/integ/sagemaker/workflow/test_retry.py b/tests/integ/sagemaker/workflow/test_retry.py index f2e6de6ba1..17d3e9ff77 100644 --- a/tests/integ/sagemaker/workflow/test_retry.py +++ b/tests/integ/sagemaker/workflow/test_retry.py @@ -31,7 +31,10 @@ _REGISTER_MODEL_NAME_BASE, _CREATE_MODEL_NAME_BASE, ) -from sagemaker.workflow.parameters import ParameterInteger +from sagemaker.workflow.parameters import ( + ParameterInteger, + ParameterString, +) from sagemaker.pytorch.estimator import PyTorch from sagemaker.workflow.pipeline import Pipeline from sagemaker.workflow.retry import ( @@ -182,11 +185,9 @@ def test_model_registration_with_model_repack( inputs = TrainingInput(s3_data=input_path) instance_count = ParameterInteger(name="InstanceCount", default_value=1) - instance_type = "ml.m5.xlarge" + instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge") good_enough_input = ParameterInteger(name="GoodEnoughInput", default_value=1) - # If image_uri is not provided, the instance_type should not be a pipeline variable - # since instance_type is used to retrieve image_uri in compile time (PySDK) pytorch_estimator = PyTorch( entry_point=entry_point, role=role, @@ -251,7 +252,7 @@ def test_model_registration_with_model_repack( ) pipeline = Pipeline( name=pipeline_name, - parameters=[good_enough_input, instance_count], + parameters=[good_enough_input, instance_count, instance_type], steps=[step_train, step_cond], sagemaker_session=pipeline_session, ) diff --git a/tests/integ/sagemaker/workflow/test_training_steps.py b/tests/integ/sagemaker/workflow/test_training_steps.py index 100d234d95..248dacac44 100644 --- a/tests/integ/sagemaker/workflow/test_training_steps.py +++ b/tests/integ/sagemaker/workflow/test_training_steps.py @@ -59,7 +59,7 @@ def test_training_job_with_debugger_and_profiler( pytorch_training_latest_py_version, ): instance_count = ParameterInteger(name="InstanceCount", default_value=1) - instance_type = "ml.m5.xlarge" + instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge") rules = [ Rule.sagemaker(rule_configs.vanishing_gradient()), @@ -78,8 +78,6 @@ def test_training_job_with_debugger_and_profiler( ) inputs = TrainingInput(s3_data=input_path) - # If image_uri is not provided, the instance_type should not be a pipeline variable - # since instance_type is used to retrieve image_uri in compile time (PySDK) pytorch_estimator = PyTorch( entry_point=script_path, role="SageMakerRole", @@ -100,7 +98,7 @@ def test_training_job_with_debugger_and_profiler( pipeline = Pipeline( name=pipeline_name, - parameters=[instance_count], + parameters=[instance_count, instance_type], steps=[step_train], sagemaker_session=sagemaker_session, ) diff --git a/tests/integ/sagemaker/workflow/test_tuning_steps.py b/tests/integ/sagemaker/workflow/test_tuning_steps.py index 0abbe4d0e7..6b994114eb 100644 --- a/tests/integ/sagemaker/workflow/test_tuning_steps.py +++ b/tests/integ/sagemaker/workflow/test_tuning_steps.py @@ -91,10 +91,8 @@ def test_tuning_single_algo( inputs = TrainingInput(s3_data=input_path) instance_count = ParameterInteger(name="InstanceCount", default_value=1) - instance_type = "ml.m5.xlarge" + instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge") - # If image_uri is not provided, the instance_type should not be a pipeline variable - # since instance_type is used to retrieve image_uri in compile time (PySDK) pytorch_estimator = PyTorch( entry_point=entry_point, role=role, @@ -169,7 +167,7 @@ def test_tuning_single_algo( pipeline = Pipeline( name=pipeline_name, - parameters=[instance_count, min_batch_size, max_batch_size], + parameters=[instance_count, instance_type, min_batch_size, max_batch_size], steps=[step_tune, step_best_model, step_second_best_model], sagemaker_session=pipeline_session, ) @@ -223,12 +221,10 @@ def test_tuning_multi_algos( ) instance_count = ParameterInteger(name="InstanceCount", default_value=1) - instance_type = "ml.m5.xlarge" + instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge") input_data = f"s3://sagemaker-sample-data-{region_name}/processing/census/census-income.csv" - # The instance_type should not be a pipeline variable - # since it is used to retrieve image_uri in compile time (PySDK) sklearn_processor = SKLearnProcessor( framework_version="0.20.0", instance_type=instance_type, @@ -263,8 +259,6 @@ def test_tuning_multi_algos( json_get_hp = JsonGet( step_name=step_process.name, property_file=property_file, json_path="train_size" ) - # If image_uri is not provided, the instance_type should not be a pipeline variable - # since instance_type is used to retrieve image_uri in compile time (PySDK) pytorch_estimator = PyTorch( entry_point=entry_point, role=role, @@ -313,7 +307,7 @@ def test_tuning_multi_algos( pipeline = Pipeline( name=pipeline_name, - parameters=[instance_count, min_batch_size, max_batch_size, static_hp_1], + parameters=[instance_count, instance_type, min_batch_size, max_batch_size], steps=[step_process, step_tune], sagemaker_session=sagemaker_session, ) diff --git a/tests/integ/sagemaker/workflow/test_workflow.py b/tests/integ/sagemaker/workflow/test_workflow.py index 6d6a498761..46eeeb4d06 100644 --- a/tests/integ/sagemaker/workflow/test_workflow.py +++ b/tests/integ/sagemaker/workflow/test_workflow.py @@ -157,14 +157,12 @@ def test_three_step_definition( athena_dataset_definition, ): framework_version = "0.20.0" - instance_type = "ml.m5.xlarge" + instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge") instance_count = ParameterInteger(name="InstanceCount", default_value=1) output_prefix = ParameterString(name="OutputPrefix", default_value="output") input_data = f"s3://sagemaker-sample-data-{region_name}/processing/census/census-income.csv" - # The instance_type should not be a pipeline variable - # since it is used to retrieve image_uri in compile time (PySDK) sklearn_processor = SKLearnProcessor( framework_version=framework_version, instance_type=instance_type, @@ -202,8 +200,6 @@ def test_three_step_definition( code=os.path.join(script_dir, "preprocessing.py"), ) - # If image_uri is not provided, the instance_type should not be a pipeline variable - # since instance_type is used to retrieve image_uri in compile time (PySDK) sklearn_train = SKLearn( framework_version=framework_version, entry_point=os.path.join(script_dir, "train.py"), @@ -242,7 +238,7 @@ def test_three_step_definition( pipeline = Pipeline( name=pipeline_name, - parameters=[instance_count, output_prefix], + parameters=[instance_type, instance_count, output_prefix], steps=[step_process, step_train, step_model], sagemaker_session=pipeline_session, ) @@ -252,6 +248,13 @@ def test_three_step_definition( assert set(tuple(param.items()) for param in definition["Parameters"]) == set( [ + tuple( + { + "Name": "InstanceType", + "Type": "String", + "DefaultValue": "ml.m5.xlarge", + }.items() + ), tuple({"Name": "InstanceCount", "Type": "Integer", "DefaultValue": 1}.items()), tuple( { @@ -296,14 +299,14 @@ def test_three_step_definition( ] ) assert processing_args["ProcessingResources"]["ClusterConfig"] == { - "InstanceType": "ml.m5.xlarge", + "InstanceType": {"Get": "Parameters.InstanceType"}, "InstanceCount": {"Get": "Parameters.InstanceCount"}, "VolumeSizeInGB": 30, } assert training_args["ResourceConfig"] == { "InstanceCount": 1, - "InstanceType": "ml.m5.xlarge", + "InstanceType": {"Get": "Parameters.InstanceType"}, "VolumeSizeInGB": 30, } assert training_args["InputDataConfig"][0]["DataSource"]["S3DataSource"]["S3Uri"] == { @@ -336,12 +339,10 @@ def test_steps_with_map_params_pipeline( ): instance_count = ParameterInteger(name="InstanceCount", default_value=2) framework_version = "0.20.0" - instance_type = "ml.m5.xlarge" + instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge") output_prefix = ParameterString(name="OutputPrefix", default_value="output") input_data = f"s3://sagemaker-sample-data-{region_name}/processing/census/census-income.csv" - # The instance_type should not be a pipeline variable - # since it is used to retrieve image_uri in compile time (PySDK) sklearn_processor = SKLearnProcessor( framework_version=framework_version, instance_type=instance_type, @@ -379,8 +380,6 @@ def test_steps_with_map_params_pipeline( code=os.path.join(script_dir, "preprocessing.py"), ) - # If image_uri is not provided, the instance_type should not be a pipeline variable - # since instance_type is used to retrieve image_uri in compile time (PySDK) sklearn_train = SKLearn( framework_version=framework_version, entry_point=os.path.join(script_dir, "train.py"), @@ -436,7 +435,7 @@ def test_steps_with_map_params_pipeline( pipeline = Pipeline( name=pipeline_name, - parameters=[instance_count, output_prefix], + parameters=[instance_type, instance_count, output_prefix], steps=[step_process, step_train, step_cond], sagemaker_session=pipeline_session, ) @@ -1030,10 +1029,8 @@ def test_model_registration_with_tuning_model( inputs = TrainingInput(s3_data=input_path) instance_count = ParameterInteger(name="InstanceCount", default_value=1) - instance_type = "ml.m5.xlarge" + instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge") - # If image_uri is not provided, the instance_type should not be a pipeline variable - # since instance_type is used to retrieve image_uri in compile time (PySDK) pytorch_estimator = PyTorch( entry_point=entry_point, role=role, @@ -1090,7 +1087,7 @@ def test_model_registration_with_tuning_model( pipeline = Pipeline( name=pipeline_name, - parameters=[instance_count, min_batch_size, max_batch_size], + parameters=[instance_count, instance_type, min_batch_size, max_batch_size], steps=[step_tune, step_register_best], sagemaker_session=pipeline_session, ) diff --git a/tests/unit/sagemaker/image_uris/test_retrieve.py b/tests/unit/sagemaker/image_uris/test_retrieve.py index ddf4049448..6dba4aee0e 100644 --- a/tests/unit/sagemaker/image_uris/test_retrieve.py +++ b/tests/unit/sagemaker/image_uris/test_retrieve.py @@ -19,7 +19,6 @@ from mock import patch from sagemaker import image_uris -from sagemaker.workflow.parameters import ParameterString BASE_CONFIG = { "processors": ["cpu", "gpu"], @@ -718,19 +717,3 @@ def test_retrieve_huggingface(config_for_framework): "564829616587.dkr.ecr.us-east-1.amazonaws.com/huggingface-pytorch-training:" "1.6.0-transformers4.3.1-gpu-py37-cu110-ubuntu18.04" == pt_new_version ) - - -def test_retrieve_with_pipeline_variable(): - with pytest.raises(Exception) as error: - image_uris.retrieve( - framework="tensorflow", - version="1.15", - py_version="py3", - instance_type=ParameterString( - name="TrainingInstanceType", - default_value="ml.m5.xlarge", - ), - region="us-east-1", - image_scope="training", - ) - assert "instance_type should not be a pipeline variable" in str(error.value) diff --git a/tests/unit/sagemaker/workflow/test_execution_variables.py b/tests/unit/sagemaker/workflow/test_execution_variables.py index 33ab191ad5..23a0484d01 100644 --- a/tests/unit/sagemaker/workflow/test_execution_variables.py +++ b/tests/unit/sagemaker/workflow/test_execution_variables.py @@ -33,7 +33,7 @@ def test_implicit_value(): with pytest.raises(TypeError) as error: str(var) - assert "Pipeline variables do not support __str__ operation." in str(error.value) + assert str(error.value) == "Pipeline variables do not support __str__ operation." with pytest.raises(TypeError) as error: int(var) diff --git a/tests/unit/sagemaker/workflow/test_functions.py b/tests/unit/sagemaker/workflow/test_functions.py index c0450f216c..a9f03d7c6d 100644 --- a/tests/unit/sagemaker/workflow/test_functions.py +++ b/tests/unit/sagemaker/workflow/test_functions.py @@ -81,7 +81,7 @@ def test_implicit_value_on_join(): with pytest.raises(TypeError) as error: str(func) - assert "Pipeline variables do not support __str__ operation." in str(error.value) + assert str(error.value) == "Pipeline variables do not support __str__ operation." with pytest.raises(TypeError) as error: int(func) @@ -189,7 +189,7 @@ def test_implicit_value_on_json_get(): with pytest.raises(TypeError) as error: str(func) - assert "Pipeline variables do not support __str__ operation." in str(error.value) + assert str(error.value) == "Pipeline variables do not support __str__ operation." with pytest.raises(TypeError) as error: int(func) diff --git a/tests/unit/sagemaker/workflow/test_parameters.py b/tests/unit/sagemaker/workflow/test_parameters.py index b9f371d883..c9aa6cb983 100644 --- a/tests/unit/sagemaker/workflow/test_parameters.py +++ b/tests/unit/sagemaker/workflow/test_parameters.py @@ -76,7 +76,7 @@ def test_parameter_to_string_and_string_implicit_value(): with pytest.raises(TypeError) as error: str(param) - assert "Pipeline variables do not support __str__ operation." in str(error.value) + assert str(error.value) == "Pipeline variables do not support __str__ operation." def test_parameter_integer_implicit_value(): diff --git a/tests/unit/sagemaker/workflow/test_properties.py b/tests/unit/sagemaker/workflow/test_properties.py index 10df1de97c..eac6168646 100644 --- a/tests/unit/sagemaker/workflow/test_properties.py +++ b/tests/unit/sagemaker/workflow/test_properties.py @@ -111,7 +111,7 @@ def test_implicit_value(): with pytest.raises(TypeError) as error: str(prop.CreationTime) - assert "Pipeline variables do not support __str__ operation." in str(error.value) + assert str(error.value) == "Pipeline variables do not support __str__ operation." with pytest.raises(TypeError) as error: int(prop.CreationTime)