chore: avoid exceeding resource limits for training timeout

Author: evakravi

tests/integ/sagemaker/jumpstart/estimator/test_jumpstart_estimator.py

@@ -40,6 +40,7 @@ def test_jumpstart_estimator(setup):
         role=get_sm_session().get_caller_identity_arn(),
         sagemaker_session=get_sm_session(),
         tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
+        max_run=259200,  # avoid exceeding resource limits
     )
 
     estimator.fit(

tests/integ/sagemaker/jumpstart/script_mode_class/test_transfer_learning.py

@@ -98,6 +98,10 @@ def test_jumpstart_transfer_learning_estimator_class(setup):
         instance_type=training_instance_type,
     )
 
+    # Avoid exceeding resource limits
+    if "max_run" in estimator_kwargs:
+        del estimator_kwargs["max_run"]
+
     estimator = Estimator(
         image_uri=image_uri,
         source_dir=script_uri,