test: stop slow CI Health V2 suite from timing out (#6391)

- jumpstart: run gated training jobs through
  fit_estimator_with_capacity_xfail so a job stuck in Pending waiting
  for ml.g5 capacity becomes an xfail after 30 minutes instead of
  blocking the build; let the helper forward fit() kwargs (accept_eula)
- clarify monitor: upload the model under a per-run key, and move the
  abalone model uploads off the shared integ-test-data/xgboost/model
  key, so concurrently running tests cannot swap in a model the
  xgboost:1 image cannot load (Clarify shadow endpoint ping failures)
- clarify monitor: cut reruns of the scheduled-execution tests from 5
  to 1, since each attempt waits for the next hourly execution
- feature processor: serialize local Spark runs across xdist workers
  with a file lock, so concurrent spark-submit --packages downloads do
  not corrupt the shared ivy cache

Author: lucasjia-aws

tests/integ/sagemaker/experiments/test_run.py

@@ -502,7 +502,7 @@ def test_run_from_transform_job(
     exp_name = unique_name_from_base(_EXP_NAME_BASE_IN_SCRIPT)
     xgb_model_data_s3 = sagemaker_session.upload_data(
         path=os.path.join(_TRANSFORM_MATERIALS, "xgb_model.tar.gz"),
-        key_prefix="integ-test-data/xgboost/model",
+        key_prefix="integ-test-data/xgboost_abalone/model",
     )
     env = _update_env_with_client_config(
         env={

tests/integ/sagemaker/feature_store/feature_processor/test_feature_processor_integ.py

@@ -16,9 +16,11 @@
 import os
 import subprocess
 import sys
+import tempfile
 import time
 from typing import Dict
 from datetime import datetime
+from filelock import FileLock
 from pyspark.sql import DataFrame
 import pytz
 
@@ -109,10 +111,28 @@
 
 SCHEDULE_EXPRESSION_TIMESTAMP_FORMAT = "%Y-%m-%dT%H:%M:%S"  # 2023-01-01T07:00:00
 
+# Local (non-@remote) feature_processor runs start a Spark JVM whose spark-submit resolves the
+# hadoop-aws/hadoop-common Maven packages into the shared ~/.ivy2 cache. Ivy is not safe for
+# concurrent processes: when two pytest-xdist workers on the same host start Spark at the same
+# time against a cold cache, their downloads collide (0-byte jars) and the JVM exits with
+# "Java gateway process exited before sending its port number". Run those tests one at a time.
+_LOCAL_SPARK_LOCK_PATH = os.path.join(
+    tempfile.gettempdir(), "sagemaker-integ-feature-processor-local-spark.lock"
+)
+_LOCAL_SPARK_LOCK_TIMEOUT_SECONDS = 30 * 60
+
+
+@pytest.fixture
+def local_spark_lock():
+    """Hold a host-wide lock so only one test runs a local Spark session at a time."""
+    with FileLock(_LOCAL_SPARK_LOCK_PATH, timeout=_LOCAL_SPARK_LOCK_TIMEOUT_SECONDS):
+        yield
+
 
 @pytest.mark.slow_test
 def test_feature_processor_transform_online_only_store_ingestion(
     sagemaker_session,
+    local_spark_lock,
 ):
     car_data_feature_group_name = get_car_data_feature_group_name()
     car_data_aggregated_feature_group_name = get_car_data_aggregated_feature_group_name()
@@ -234,6 +254,7 @@ def transform(raw_s3_data_as_df):
 @pytest.mark.slow_test
 def test_feature_processor_transform_with_customized_data_source(
     sagemaker_session,
+    local_spark_lock,
 ):
     car_data_feature_group_name = get_car_data_feature_group_name()
     car_data_aggregated_feature_group_name = get_car_data_aggregated_feature_group_name()
@@ -377,6 +398,7 @@ def transform(raw_s3_data_as_df):
 @pytest.mark.flaky(reruns=5, reruns_delay=2)
 def test_feature_processor_transform_offline_only_store_ingestion(
     sagemaker_session,
+    local_spark_lock,
 ):
     car_data_feature_group_name = get_car_data_feature_group_name()
     car_data_aggregated_feature_group_name = get_car_data_aggregated_feature_group_name()

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

@@ -121,12 +121,14 @@ def test_gated_model_training_v1(setup):
         hyperparameters={"max_steps": "1"},
     )
 
-    # uses ml.g5.12xlarge instance
-    estimator.fit(
+    # uses ml.g5.12xlarge instance, which can sit in Pending for hours waiting for
+    # capacity; fail fast as a CapacityError instead of blocking the whole build.
+    fit_estimator_with_capacity_xfail(
+        estimator,
         {
             "training": f"s3://{get_jumpstart_content_bucket(JUMPSTART_DEFAULT_REGION_NAME)}/"
             f"{get_training_dataset_for_model_and_version(model_id, model_version)}",
-        }
+        },
     )
 
     # uses ml.g5.2xlarge instance
@@ -167,12 +169,14 @@ def test_gated_model_training_v2(setup):
         hyperparameters={"max_steps": "1"},
     )
 
-    # uses ml.g5.12xlarge instance
-    estimator.fit(
+    # uses ml.g5.12xlarge instance, which can sit in Pending for hours waiting for
+    # capacity; fail fast as a CapacityError instead of blocking the whole build.
+    fit_estimator_with_capacity_xfail(
+        estimator,
         {
             "training": f"s3://{get_jumpstart_content_bucket(JUMPSTART_DEFAULT_REGION_NAME)}/"
             f"{get_training_dataset_for_model_and_version(model_id, model_version)}",
-        }
+        },
     )
 
     # test that we can create a JumpStartEstimator from existing job with `attach`

tests/integ/sagemaker/jumpstart/private_hub/estimator/test_jumpstart_private_hub_estimator.py

@@ -30,6 +30,7 @@
     get_sm_session,
     get_training_dataset_for_model_and_version,
     add_model_references_to_hub,
+    fit_estimator_with_capacity_xfail,
     x_fail_if_ice,
 )
 
@@ -151,12 +152,15 @@ def test_jumpstart_hub_gated_estimator_with_eula(setup, add_model_references):
         hyperparameters={"max_steps": "1"},
     )
 
-    estimator.fit(
-        accept_eula=True,
-        inputs={
+    # Fail fast as a CapacityError if ml.g5.2xlarge capacity is not granted, instead of
+    # letting the training job sit in Pending until the build times out.
+    fit_estimator_with_capacity_xfail(
+        estimator,
+        {
             "training": f"s3://{get_jumpstart_content_bucket(JUMPSTART_DEFAULT_REGION_NAME)}/"
             f"{get_training_dataset_for_model_and_version(model_id, model_version)}",
         },
+        accept_eula=True,
     )
 
     predictor = estimator.deploy(

tests/integ/sagemaker/jumpstart/utils.py

@@ -104,6 +104,7 @@ def fit_estimator_with_capacity_xfail(
     estimator,
     inputs,
     capacity_wait_timeout_seconds: int = DEFAULT_CAPACITY_WAIT_TIMEOUT_SECONDS,
+    **fit_kwargs,
 ):
     """Submit a training job and wait for it, but fail fast on capacity queueing.
 
@@ -122,8 +123,10 @@ def fit_estimator_with_capacity_xfail(
         inputs: The training inputs passed through to ``estimator.fit()``.
         capacity_wait_timeout_seconds: Max time to allow the job to remain in
             "waiting for capacity" before declaring a capacity shortage.
+        **fit_kwargs: Extra keyword arguments passed through to ``estimator.fit()``
+            (e.g. ``accept_eula=True``).
     """
-    estimator.fit(inputs, wait=False)
+    estimator.fit(inputs, wait=False, **fit_kwargs)
 
     training_job_name = estimator.latest_training_job.name
     sagemaker_client = estimator.sagemaker_session.sagemaker_client

tests/integ/sagemaker/workflow/test_model_steps.py

@@ -704,7 +704,7 @@ def test_xgboost_model_register_and_deploy_with_runtime_repack(
     model_package_group_name = f"{pipeline_name}TestModelPackageGroup"
     xgb_model_data_s3 = pipeline_session.upload_data(
         path=os.path.join(_XGBOOST_PATH, "xgb_model.tar.gz"),
-        key_prefix="integ-test-data/xgboost/model",
+        key_prefix="integ-test-data/xgboost_abalone/model",
     )
     xgb_model_data_param = ParameterString(name="ModelData", default_value=xgb_model_data_s3)
     xgb_model = XGBoostModel(

tests/integ/sagemaker/workflow/test_monitor_batch_transform_step.py

@@ -205,7 +205,7 @@ def step_model_create(
 
     xgb_model_data_s3 = pipeline_session.upload_data(
         path=os.path.join(os.path.join(DATA_DIR, "xgboost_abalone"), "xgb_model.tar.gz"),
-        key_prefix="integ-test-data/xgboost/model",
+        key_prefix="integ-test-data/xgboost_abalone/model",
     )
 
     xgb_model = XGBoostModel(

tests/integ/test_clarify_model_monitor.py

@@ -85,9 +85,12 @@
 @pytest.yield_fixture(scope="module")
 def endpoint_name(sagemaker_session):
     endpoint_name = unique_name_from_base("clarify-xgb-integ")
+    # Upload under a key unique to this run. Clarify re-downloads the model artifact when it
+    # creates its shadow endpoint for every monitoring execution, long after this upload, so
+    # a shared key can by then hold a different model written by a concurrently running test.
     xgb_model_data = sagemaker_session.upload_data(
         path=os.path.join(XGBOOST_DATA_PATH, "xgb_model.tar.gz"),
-        key_prefix="integ-test-data/xgboost/model",
+        key_prefix=f"integ-test-data/xgboost/model/{endpoint_name}",
     )
 
     xgb_image = image_uris.retrieve(
@@ -360,7 +363,9 @@ def test_bias_monitor(sagemaker_session, scheduled_bias_monitor, endpoint_name,
     tests.integ.test_region() in tests.integ.NO_MODEL_MONITORING_REGIONS,
     reason="ModelMonitoring is not yet supported in this region.",
 )
-@pytest.mark.flaky(reruns=5, reruns_delay=2)
+# Each attempt waits for the next hourly scheduled execution, so every rerun costs up to an
+# hour of build time; more than one rerun can push the slow suite past its build timeout.
+@pytest.mark.flaky(reruns=1, reruns_delay=2)
 def test_run_bias_monitor(
     scheduled_bias_monitor, sagemaker_session, endpoint_name, ground_truth_input, upload_actual_data
 ):
@@ -470,7 +475,8 @@ def test_explainability_monitor(sagemaker_session, scheduled_explainability_moni
     tests.integ.test_region() in tests.integ.NO_MODEL_MONITORING_REGIONS,
     reason="ModelMonitoring is not yet supported in this region.",
 )
-@pytest.mark.flaky(reruns=5, reruns_delay=2)
+# See test_run_bias_monitor: each rerun waits for another hourly execution.
+@pytest.mark.flaky(reruns=1, reruns_delay=2)
 def test_run_explainability_monitor(
     scheduled_explainability_monitor,
     sagemaker_session,

tests/integ/test_model_package.py

@@ -44,7 +44,7 @@ def test_update_approval_model_package(sagemaker_session):
 
     xgb_model_data_s3 = sagemaker_session.upload_data(
         path=os.path.join(_XGBOOST_PATH, "xgb_model.tar.gz"),
-        key_prefix="integ-test-data/xgboost/model",
+        key_prefix="integ-test-data/xgboost_abalone/model",
     )
     model = XGBoostModel(
         model_data=xgb_model_data_s3, framework_version="1.3-1", sagemaker_session=sagemaker_session
@@ -86,7 +86,7 @@ def test_update_model_life_cycle_model_package(sagemaker_session):
 
     xgb_model_data_s3 = sagemaker_session.upload_data(
         path=os.path.join(_XGBOOST_PATH, "xgb_model.tar.gz"),
-        key_prefix="integ-test-data/xgboost/model",
+        key_prefix="integ-test-data/xgboost_abalone/model",
     )
     model = XGBoostModel(
         model_data=xgb_model_data_s3, framework_version="1.3-1", sagemaker_session=sagemaker_session
@@ -142,7 +142,7 @@ def test_inference_specification_addition(sagemaker_session):
 
     xgb_model_data_s3 = sagemaker_session.upload_data(
         path=os.path.join(_XGBOOST_PATH, "xgb_model.tar.gz"),
-        key_prefix="integ-test-data/xgboost/model",
+        key_prefix="integ-test-data/xgboost_abalone/model",
     )
     model = XGBoostModel(
         model_data=xgb_model_data_s3, framework_version="1.3-1", sagemaker_session=sagemaker_session
@@ -223,7 +223,7 @@ def test_update_source_uri(sagemaker_session):
 
     xgb_model_data_s3 = sagemaker_session.upload_data(
         path=os.path.join(_XGBOOST_PATH, "xgb_model.tar.gz"),
-        key_prefix="integ-test-data/xgboost/model",
+        key_prefix="integ-test-data/xgboost_abalone/model",
     )
     model = XGBoostModel(
         model_data=xgb_model_data_s3, framework_version="1.3-1", sagemaker_session=sagemaker_session
@@ -287,7 +287,7 @@ def test_update_model_card_with_model_card_object(sagemaker_session):
 
     xgb_model_data_s3 = sagemaker_session.upload_data(
         path=os.path.join(_XGBOOST_PATH, "xgb_model.tar.gz"),
-        key_prefix="integ-test-data/xgboost/model",
+        key_prefix="integ-test-data/xgboost_abalone/model",
     )
     model = XGBoostModel(
         model_data=xgb_model_data_s3, framework_version="1.3-1", sagemaker_session=sagemaker_session
@@ -384,7 +384,7 @@ def test_update_model_card_with_model_card_json(sagemaker_session):
 
     xgb_model_data_s3 = sagemaker_session.upload_data(
         path=os.path.join(_XGBOOST_PATH, "xgb_model.tar.gz"),
-        key_prefix="integ-test-data/xgboost/model",
+        key_prefix="integ-test-data/xgboost_abalone/model",
     )
     model = XGBoostModel(
         model_data=xgb_model_data_s3, framework_version="1.3-1", sagemaker_session=sagemaker_session
@@ -468,7 +468,7 @@ def test_clone_model_package_using_source_uri(sagemaker_session):
 
     xgb_model_data_s3 = sagemaker_session.upload_data(
         path=os.path.join(_XGBOOST_PATH, "xgb_model.tar.gz"),
-        key_prefix="integ-test-data/xgboost/model",
+        key_prefix="integ-test-data/xgboost_abalone/model",
     )
     model = XGBoostModel(
         model_data=xgb_model_data_s3, framework_version="1.3-1", sagemaker_session=sagemaker_session
@@ -545,7 +545,7 @@ def test_register_model_using_source_uri(sagemaker_session):
 
     xgb_model_data_s3 = sagemaker_session.upload_data(
         path=os.path.join(_XGBOOST_PATH, "xgb_model.tar.gz"),
-        key_prefix="integ-test-data/xgboost/model",
+        key_prefix="integ-test-data/xgboost_abalone/model",
     )
 
     model = XGBoostModel(

tests/integ/test_transformer.py

@@ -666,7 +666,7 @@ def test_transformer_and_monitoring_job(
 ):
     xgb_model_data_s3 = pipeline_session.upload_data(
         path=os.path.join(os.path.join(DATA_DIR, "xgboost_abalone"), "xgb_model.tar.gz"),
-        key_prefix="integ-test-data/xgboost/model",
+        key_prefix="integ-test-data/xgboost_abalone/model",
     )
     data_bias_supplied_baseline_constraints = Constraints.from_file_path(
         constraints_file_path=os.path.join(
@@ -729,7 +729,7 @@ def test_transformer_and_monitoring_job_to_pass_with_no_failure_in_violation(
 ):
     xgb_model_data_s3 = pipeline_session.upload_data(
         path=os.path.join(os.path.join(DATA_DIR, "xgboost_abalone"), "xgb_model.tar.gz"),
-        key_prefix="integ-test-data/xgboost/model",
+        key_prefix="integ-test-data/xgboost_abalone/model",
     )
     data_bias_supplied_baseline_constraints = Constraints.from_file_path(
         constraints_file_path=os.path.join(