test(instance-preferences): enforce the public contract in integ tests (#2158)

Instance Preferences is publicly released, so the integ tests no longer
need the pre-release escape hatches: the XFAIL branches for gated
Describe responses, model @required rejections, and InternalServerError
before type resolution are removed, along with the rollout notes that
explained them. Both tests now assert the full contract unconditionally:
create accepted, preferences echoed, top-level instance type absent,
winner among the submitted preferences with the expected count.

The processing test used a private loader for the bundled service model;
the public botocore model now carries the shapes, so it uses a plain
boto3 client. Candidate instance types are overridable via
INSTANCE_PREFERENCES_TEST_INSTANCE_TYPES /
PROCESSING_INSTANCE_PREFERENCES_TEST_INSTANCE_TYPES so the tests can run
in accounts where the default types are quota-saturated.

Verified against prod us-west-2: processing passes (winner resolved,
3 min); training passes (winner resolved, 8 min).

Author: deeppcs

sagemaker-core/tests/integ/processing/test_instance_preferences_processing.py

@@ -28,16 +28,12 @@
 
 - ``PROCESSING_INSTANCE_PREFERENCES_TEST_ROLE_ARN``  - execution role ARN
 - ``PROCESSING_INSTANCE_PREFERENCES_TEST_IMAGE_URI`` - processing image
+- ``PROCESSING_INSTANCE_PREFERENCES_TEST_INSTANCE_TYPES`` - optional comma-separated
+  candidate types (default: ml.m5.xlarge, ml.m5.large)
 - ``SAGEMAKER_ENDPOINT``                             - optional endpoint
-  override (e.g. a pre-GA stage endpoint).
-
-Rollout note: like the training analog
-(``sagemaker-train/tests/integ/train/test_instance_preferences.py``), while
-the server-side scheduling/write-back stages are not yet deployed on the
-target endpoint an accepted job fails with ``InternalServerError`` before an
-instance type is resolved; the test reports that as XFAIL (rollout
-incomplete) and enforces the full contract automatically once deployed.
+  override.
 """
+
 from __future__ import absolute_import
 
 import os
@@ -50,7 +46,13 @@
 IMAGE_URI = os.environ.get("PROCESSING_INSTANCE_PREFERENCES_TEST_IMAGE_URI")
 ENDPOINT = os.environ.get("SAGEMAKER_ENDPOINT")
 
-PREFERENCE_TYPES = ["ml.m5.xlarge", "ml.m5.large"]
+PREFERENCE_TYPES = [
+    t.strip()
+    for t in os.environ.get(
+        "PROCESSING_INSTANCE_PREFERENCES_TEST_INSTANCE_TYPES", "ml.m5.xlarge,ml.m5.large"
+    ).split(",")
+    if t.strip()
+]
 WAIT_TIMEOUT_SECONDS = 30 * 60
 POLL_SECONDS = 30
 
@@ -65,27 +67,11 @@
 
 
 def _sagemaker_client():
-    """Build a botocore client on the bundled service model (which carries the
-    pre-GA InstancePreferences shapes), honoring SAGEMAKER_ENDPOINT."""
-    import pathlib
-
-    import botocore.loaders
-    import botocore.session as bc_session_mod
-    from boto3.session import Session as Boto3Session
-
-    # .../sagemaker-core/tests/integ/processing/<this file> -> .../sagemaker-core/sample
-    sample_dir = pathlib.Path(__file__).resolve().parents[3] / "sample"
-    assert (sample_dir / "sagemaker").exists(), f"bundled model dir not found: {sample_dir}"
-    sample_dir = str(sample_dir)
-    bc_session = bc_session_mod.get_session()
-    loader = botocore.loaders.Loader(
-        extra_search_paths=[sample_dir], include_default_search_paths=True
-    )
-    bc_session.register_component("data_loader", loader)
+    """Build a boto3 SageMaker client, honoring SAGEMAKER_ENDPOINT."""
+    import boto3
+
     region = os.environ.get("AWS_REGION", "us-west-2")
-    return Boto3Session(botocore_session=bc_session, region_name=region).client(
-        "sagemaker", endpoint_url=ENDPOINT
-    )
+    return boto3.client("sagemaker", region_name=region, endpoint_url=ENDPOINT)
 
 
 def test_processor_instance_preferences_e2e():
@@ -107,26 +93,7 @@ def test_processor_instance_preferences_e2e():
     )
 
     job_name = f"instance-prefs-proc-integ-{uuid.uuid4().hex[:8]}"
-    try:
-        processor.run(wait=False, logs=False, job_name=job_name)
-    except Exception as e:  # botocore ClientError
-        message = str(e)
-        if "ValidationException" in message and (
-            "instanceType' failed to satisfy constraint: Member must not be null" in message
-            or "InstanceType and InstanceCount greater than 0 must be specified" in message
-        ):
-            # The public processing model still carries @required on
-            # InstanceType/InstanceCount (the relaxation is bundled with the
-            # GA/Trebuchet ungating), so a preferences-only create is rejected
-            # by the frontend model validation on this endpoint. The SDK-side
-            # request construction and serialization succeeded.
-            pytest.xfail(
-                "Server rollout incomplete on this endpoint: CreateProcessingJob "
-                "rejected a preferences-only ClusterConfig (public model @required "
-                "relaxation on InstanceType/InstanceCount not yet deployed): "
-                f"{message[:200]}"
-            )
-        raise
+    processor.run(wait=False, logs=False, job_name=job_name)
 
     # --- Create accepted; Describe echoes the request contract -------------
     described = client.describe_processing_job(ProcessingJobName=job_name)
@@ -152,12 +119,6 @@ def test_processor_instance_preferences_e2e():
         time.sleep(POLL_SECONDS)
 
     failure_reason = described.get("FailureReason", "")
-    if status == "Failed" and "InternalServerError" in failure_reason:
-        pytest.xfail(
-            "Server rollout incomplete on this endpoint: processing job accepted "
-            "with InstancePreferences but failed with InternalServerError before "
-            f"type resolution (job={job_name})"
-        )
 
     # --- Full contract: resolved winner is surfaced and is a submitted pref -
     cluster_config = described["ProcessingResources"]["ClusterConfig"]

sagemaker-train/tests/integ/train/test_instance_preferences.py

@@ -30,16 +30,10 @@
 - ``INSTANCE_PREFERENCES_TEST_IMAGE_URI``  - training image the account can pull
 - ``INSTANCE_PREFERENCES_TEST_S3_OUTPUT``  - s3:// output path
 - ``INSTANCE_PREFERENCES_TEST_S3_INPUT``   - optional s3:// input channel
-- ``SAGEMAKER_ENDPOINT``                   - optional endpoint override (e.g.
-  a pre-GA stage endpoint); honored by the sagemaker-core client loader.
-
-Rollout note: the server-side feature ships in stages (validation ->
-scheduling sweep -> resolved-type write-back). While the scheduling stage is
-not yet deployed to the target endpoint, an accepted job fails with
-``InternalServerError`` before an instance type is resolved; the test detects
-that specific signature and reports it as XFAIL (server rollout incomplete)
-instead of a test failure, so it starts enforcing the full contract
-automatically once the stage deploys.
+- ``INSTANCE_PREFERENCES_TEST_INSTANCE_TYPES`` - optional comma-separated
+  candidate types (default: ml.m5.xlarge, ml.m4.xlarge)
+- ``SAGEMAKER_ENDPOINT``                   - optional endpoint override;
+  honored by the sagemaker-core client loader.
 """
 
 from __future__ import absolute_import
@@ -57,7 +51,13 @@
 S3_OUTPUT = os.environ.get("INSTANCE_PREFERENCES_TEST_S3_OUTPUT")
 S3_INPUT = os.environ.get("INSTANCE_PREFERENCES_TEST_S3_INPUT")
 
-PREFERENCE_TYPES = ["ml.m5.xlarge", "ml.m4.xlarge"]
+PREFERENCE_TYPES = [
+    t.strip()
+    for t in os.environ.get(
+        "INSTANCE_PREFERENCES_TEST_INSTANCE_TYPES", "ml.m5.xlarge,ml.m4.xlarge"
+    ).split(",")
+    if t.strip()
+]
 WAIT_TIMEOUT_SECONDS = 30 * 60
 POLL_SECONDS = 30
 
@@ -118,14 +118,9 @@ def test_model_trainer_instance_preferences_e2e():
     echoed = described.resource_config
     # Top-level instance type was never set; must not come back populated.
     assert _get_value(echoed.instance_type) is None
-    if _get_value(echoed.instance_preferences) is None:
-        # Gated (pre-GA) endpoints strip instance_preferences/selected_* from
-        # public Describe responses; the full contract is assertable post-ungate.
-        pytest.xfail(
-            "Dark-launch gating active on this endpoint: instance_preferences/"
-            "selected_* are stripped from the public Describe response "
-            f"(job={job_name}). Re-run after the GA ungate."
-        )
+    assert [p.instance_type for p in _get_value(echoed.instance_preferences) or []] == (
+        PREFERENCE_TYPES
+    ), f"Describe did not echo instance_preferences (job={job_name})"
 
     # Wait until terminal or winner visible (Selected* propagation lags the
     # secondary-status transitions; gate on the winner, not on secondary).
@@ -141,24 +136,15 @@ def test_model_trainer_instance_preferences_e2e():
         time.sleep(POLL_SECONDS)
 
     failure_reason = _get_value(described.failure_reason) or ""
-    if status == "Failed" and "InternalServerError" in failure_reason:
-        # Accepted by validation but failed before type resolution: the
-        # scheduling stage of the server rollout is not deployed on this
-        # endpoint yet. Everything client-side (model, serialization,
-        # acceptance, describe echo) has been asserted above.
-        pytest.xfail(
-            "Server rollout incomplete on this endpoint: job accepted with "
-            "instance_preferences but failed with InternalServerError before "
-            f"type resolution (job={job_name})"
-        )
 
     # --- Full contract: resolved winner is surfaced and is a submitted pref -
     final_rc = described.resource_config
     selected_type = _get_value(final_rc.selected_instance_type)
     selected_count = _get_value(final_rc.selected_instance_count)
     assert selected_type in PREFERENCE_TYPES, (
         f"selected_instance_type={selected_type!r} not among submitted "
-        f"preferences {PREFERENCE_TYPES} (job={job_name}, status={status})"
+        f"preferences {PREFERENCE_TYPES} (job={job_name}, status={status}, "
+        f"failure={failure_reason!r})"
     )
     assert selected_count == 1
     assert _get_value(final_rc.instance_type) is None