fix: rewrite inference integration tests using SageMaker SDK v3

Bhanu Teja Goshikonda · Bhanu Teja Goshikonda · commit 0426f698938d · 2026-06-15T16:23:25.000-07:00
Replaces commit 4b99d0f which used boto3 directly. The SageMaker Python SDK v3 is the supported entrypoint for these tests; v3 removed the v2 sagemaker.Session, sagemaker.tensorflow.serving.TensorFlowModel, and sagemaker.multidatamodel.MultiDataModel classes the old fixtures relied on. For DLC integration tests we already supply image_uri and a pre-built model.tar.gz, so ModelBuilder's auto-detection adds no value. We use the v3 sagemaker-core resource layer directly (the same surface ModelBuilder calls underneath): - conftest sagemaker_session fixture uses sagemaker.core.helper.session_helper.Session (default_bucket / upload_data); cleanup_endpoint uses Endpoint.get(...).delete(), EndpointConfig.get(...).delete(), Model.get(...).delete() - single-model test uses Model.create(primary_container=ContainerDefinition(...)) + EndpointConfig.create([ProductionVariant(...)]) + Endpoint.create() with endpoint.wait_for_status("InService") and endpoint.invoke(...) - MME test expresses the multi-model contract directly: ContainerDefinition(mode="MultiModel", model_data_url=<S3 prefix>) and endpoint.invoke(target_model="modelN.tar.gz") - sagemaker>=3.0.0 retained in requirements.txt
diff --git a/test/tensorflow/integration/inference/conftest.py b/test/tensorflow/integration/inference/conftest.py
@@ -1,11 +1,15 @@
 """Pytest fixtures for TF 2.20 inference integration tests on SageMaker.
 
-These fixtures use boto3 directly (sagemaker, sagemaker-runtime, s3 clients)
-rather than the SageMaker Python SDK. The PyPI ``sagemaker`` package v3.x
-removed the legacy v2 surfaces this suite relied on (``sagemaker.Session``,
-``sagemaker.tensorflow.serving.TensorFlowModel``, ``sagemaker.multidatamodel``),
-and the v3 ``ModelBuilder`` flow is heavier than what these integration tests
-need. Going boto3-only keeps the tests transparent and SDK-version-independent.
+Uses the SageMaker Python SDK v3 (``sagemaker>=3.0.0``) — the v2 Estimator /
+Model / Predictor classes were removed in v3 in favor of the unified
+``ModelBuilder`` and the ``sagemaker-core`` resource layer
+(``Endpoint``, ``EndpointConfig``, ``Model``, ``ContainerDefinition``,
+``ProductionVariant``). For these DLC tests we already have a custom
+``image_uri`` and a pre-built ``model.tar.gz``, so the simplest v3 path is
+the resource layer directly: ``Model.create -> EndpointConfig.create ->
+Endpoint.create -> endpoint.invoke()``. ``ModelBuilder`` is the right choice
+when the SDK should auto-detect the framework / container / packaging — for
+us, those are all fixed by the test fixture inputs.
 
 Fixtures intentionally defer all AWS calls until test-execution time so that
 ``pytest --collect-only`` works in environments without AWS credentials.
@@ -46,58 +50,29 @@ def inference_image_uri() -> str:
 
 @pytest.fixture(scope="session")
 def boto_session(aws_region: str):
-    """A boto3 session bound to the configured region."""
+    """A boto3 session bound to the configured region.
+
+    Used purely as a transport for ``sagemaker.core.helper.session_helper.Session``
+    and for the underlying ``s3`` client when uploading model artifacts; no
+    SageMaker control-plane calls go through it directly.
+    """
     import boto3
 
     return boto3.Session(region_name=aws_region)
 
 
 @pytest.fixture(scope="session")
-def sagemaker_client(boto_session):
-    """Low-level SageMaker control-plane client (create/delete model, endpoint, ...)."""
-    return boto_session.client("sagemaker")
-
-
-@pytest.fixture(scope="session")
-def sagemaker_runtime_client(boto_session):
-    """SageMaker runtime client used to invoke endpoints."""
-    return boto_session.client("sagemaker-runtime")
-
-
-@pytest.fixture(scope="session")
-def s3_client(boto_session):
-    """S3 client used to upload sample model tarballs."""
-    return boto_session.client("s3")
-
-
-@pytest.fixture(scope="session")
-def default_bucket(boto_session, aws_region: str, s3_client) -> str:
-    """Resolve the ``sagemaker-<region>-<account>`` default bucket, creating it if absent.
+def sagemaker_session(boto_session):
+    """A SageMaker SDK v3 session.
 
-    Mirrors the behaviour of the v2 SDK's ``Session.default_bucket()`` so test
-    bodies can keep using a single, predictable bucket without callers having
-    to plumb one in.
+    ``sagemaker.core.helper.session_helper.Session`` is the v3 replacement for
+    the v2 ``sagemaker.Session``. We use it for ``default_bucket()`` and
+    ``upload_data()``; resource-layer ``create()`` calls accept it via the
+    ``session=`` kwarg.
     """
-    sts = boto_session.client("sts")
-    account_id = sts.get_caller_identity()["Account"]
-    bucket = f"sagemaker-{aws_region}-{account_id}"
-
-    try:
-        s3_client.head_bucket(Bucket=bucket)
-    except Exception:
-        # Bucket missing or inaccessible — try to create it. us-east-1 must omit
-        # LocationConstraint; every other region requires it.
-        create_kwargs: dict = {"Bucket": bucket}
-        if aws_region != "us-east-1":
-            create_kwargs["CreateBucketConfiguration"] = {"LocationConstraint": aws_region}
-        try:
-            s3_client.create_bucket(**create_kwargs)
-        except Exception:
-            # Race or pre-existing-but-403 — leave the original error to surface
-            # at upload time rather than masking it here.
-            pass
-
-    return bucket
+    from sagemaker.core.helper.session_helper import Session
+
+    return Session(boto_session=boto_session)
 
 
 @pytest.fixture
@@ -115,61 +90,20 @@ def _make(prefix: str) -> str:
 
 
 @pytest.fixture
-def upload_to_s3(s3_client):
-    """Yield-style helper that uploads a local file to ``s3://<bucket>/<key>``.
-
-    Returns the resulting ``s3://`` URI. Failures propagate; teardown is
-    intentionally not provided since SageMaker model artifacts are typically
-    left in the bucket for forensic inspection.
-    """
-
-    def _upload(local_path: str, bucket: str, key: str) -> str:
-        s3_client.upload_file(local_path, bucket, key)
-        return f"s3://{bucket}/{key}"
-
-    return _upload
-
-
-@pytest.fixture
-def wait_for_endpoint(sagemaker_client):
-    """Poll ``describe_endpoint`` until the endpoint reaches ``InService``.
-
-    Raises ``RuntimeError`` if the endpoint enters ``Failed`` / ``OutOfService``
-    or if the wait exceeds ``timeout_seconds`` (default 1800s = 30 min, which
-    matches typical first-pull cold-start latency for inference DLCs).
-    """
-
-    def _wait(endpoint_name: str, timeout_seconds: int = 1800, poll_seconds: int = 30) -> None:
-        deadline = time.time() + timeout_seconds
-        while time.time() < deadline:
-            resp = sagemaker_client.describe_endpoint(EndpointName=endpoint_name)
-            status = resp["EndpointStatus"]
-            if status == "InService":
-                return
-            if status in {"Failed", "OutOfService"}:
-                reason = resp.get("FailureReason", "<no FailureReason>")
-                raise RuntimeError(
-                    f"endpoint {endpoint_name} entered terminal state {status}: {reason}"
-                )
-            time.sleep(poll_seconds)
-        raise RuntimeError(
-            f"endpoint {endpoint_name} did not reach InService within {timeout_seconds}s"
-        )
-
-    return _wait
-
-
-@pytest.fixture
-def cleanup_endpoint(sagemaker_client):
+def cleanup_endpoint(boto_session):
     """Yield-style fixture that tears down endpoint, endpoint config, and model.
 
+    Uses the v3 ``sagemaker-core`` resource layer (``Endpoint.get(...).delete()``,
+    etc.) rather than raw boto3 SDK calls, so cleanup code matches the deploy
+    code in the tests. The ``session=`` kwarg on resource ``get`` / ``create``
+    methods accepts a raw ``boto3.session.Session`` (see
+    ``sagemaker.core.utils.utils.SageMakerClient``); pass ``boto_session``
+    rather than the helper ``Session``.
+
     Usage:
         def test_x(cleanup_endpoint, ...):
             cleanup_endpoint(endpoint_name, model_name=model_name)
             # ... deploy + predict ...
-
-    Endpoint config is registered under the same name as the endpoint, matching
-    what the test bodies pass to ``create_endpoint_config``.
     """
     registered: list[dict] = []
 
@@ -178,22 +112,26 @@ def _register(endpoint_name: str, model_name: str | None = None) -> None:
 
     yield _register
 
+    # Import lazily so collection works without the SDK installed.
+    from sagemaker.core.resources import Endpoint, EndpointConfig, Model
+
     for item in registered:
         endpoint_name = item["endpoint_name"]
         model_name = item["model_name"]
 
-        for delete_call, kwargs in (
-            (sagemaker_client.delete_endpoint, {"EndpointName": endpoint_name}),
-            (sagemaker_client.delete_endpoint_config, {"EndpointConfigName": endpoint_name}),
+        # Endpoint config name == endpoint name in our deploy flow below.
+        for resource_cls, get_kwargs in (
+            (Endpoint, {"endpoint_name": endpoint_name}),
+            (EndpointConfig, {"endpoint_config_name": endpoint_name}),
         ):
             try:
-                delete_call(**kwargs)
+                resource_cls.get(session=boto_session, **get_kwargs).delete()
             except Exception:
-                # Swallow NotFound / already-deleted; teardown is best-effort.
+                # Best-effort teardown: swallow NotFound / already-deleted.
                 pass
 
         if model_name:
             try:
-                sagemaker_client.delete_model(ModelName=model_name)
+                Model.get(model_name=model_name, session=boto_session).delete()
             except Exception:
                 pass
diff --git a/test/tensorflow/integration/inference/requirements.txt b/test/tensorflow/integration/inference/requirements.txt
@@ -1,2 +1,3 @@
 boto3
 pytest
+sagemaker>=3.0.0
diff --git a/test/tensorflow/integration/inference/test_multi_model_endpoint.py b/test/tensorflow/integration/inference/test_multi_model_endpoint.py
@@ -2,13 +2,16 @@
 
 Builds two tiny SavedModels (``y = 2x`` and ``y = 3x``), uploads both to a
 shared S3 prefix, deploys a SageMaker MME backed by the v2 inference image,
-and asserts that ``TargetModel`` routes invocations correctly.
-
-Uses boto3 directly (sagemaker, sagemaker-runtime, s3 clients) rather than
-the SageMaker Python SDK, because the v3.x ``sagemaker`` package removed
-the v2 ``MultiDataModel`` / ``TensorFlowModel`` classes this test originally
-relied on. MME is a thin server-side construct anyway: ``Mode: MultiModel``
-on the PrimaryContainer plus a model_data S3 prefix.
+and asserts that ``target_model`` routes invocations correctly.
+
+Uses the SageMaker Python SDK v3 ``sagemaker-core`` resource layer — v3
+removed ``sagemaker.multidatamodel.MultiDataModel``. The native MME wire
+contract (``ContainerDefinition.mode = "MultiModel"``,
+``model_data_url = s3://bucket/prefix/``, plus the
+``X-Amzn-SageMaker-Target-Model`` header on invoke) is unchanged, so we
+express it directly: ``Model.create`` with ``mode="MultiModel"`` and an S3
+prefix in ``model_data_url``, then ``endpoint.invoke(target_model=...)``
+which sets the runtime header for us.
 """
 
 from __future__ import annotations
@@ -24,14 +27,6 @@
 INSTANCE_TYPE = "ml.c5.xlarge"
 
 
-def _decode(invoke_response) -> dict:
-    """Decode an ``invoke_endpoint`` response Body into a Python dict."""
-    body = invoke_response["Body"].read()
-    if isinstance(body, (bytes, bytearray)):
-        body = body.decode("utf-8")
-    return json.loads(body)
-
-
 def _values_from_predictions(predictions) -> list:
     """Pull the numeric output list out of either signature-keyed or raw rows."""
     assert predictions and isinstance(predictions, list)
@@ -42,16 +37,21 @@ def _values_from_predictions(predictions) -> list:
 
 
 def test_mme_two_models(
-    sagemaker_client,
-    sagemaker_runtime_client,
+    boto_session,
+    sagemaker_session,
     sagemaker_role_arn,
     inference_image_uri,
-    default_bucket,
-    upload_to_s3,
     unique_name,
     cleanup_endpoint,
-    wait_for_endpoint,
 ):
+    from sagemaker.core.resources import (
+        ContainerDefinition,
+        Endpoint,
+        EndpointConfig,
+        Model,
+        ProductionVariant,
+    )
+
     with tempfile.TemporaryDirectory(prefix="tf220-mme-") as workdir:
         workdir_path = Path(workdir)
 
@@ -66,73 +66,78 @@ def test_mme_two_models(
             output_dir=model2_dir, multiplier=3.0, model_name="model", tar_filename="model2.tar.gz"
         )
 
+        bucket = sagemaker_session.default_bucket()
         run_id = unique_name("mme")
         s3_key_prefix = f"tf220-inference-tests/mme-models/{run_id}"
 
         # Upload each tarball under the shared MME prefix so the runtime can
-        # resolve TargetModel relative to the same S3 location.
-        upload_to_s3(model1_tar, default_bucket, f"{s3_key_prefix}/model1.tar.gz")
-        upload_to_s3(model2_tar, default_bucket, f"{s3_key_prefix}/model2.tar.gz")
-        s3_model_prefix = f"s3://{default_bucket}/{s3_key_prefix}/"
+        # resolve target_model relative to the same S3 location.
+        sagemaker_session.upload_data(path=model1_tar, bucket=bucket, key_prefix=s3_key_prefix)
+        sagemaker_session.upload_data(path=model2_tar, bucket=bucket, key_prefix=s3_key_prefix)
+        s3_model_prefix = f"s3://{bucket}/{s3_key_prefix}/"
 
         endpoint_name = unique_name("tf220-mme")
-        base_model_name = unique_name("tf220-mme-model")
-        cleanup_endpoint(endpoint_name, model_name=base_model_name)
-
-        # ``Mode: MultiModel`` plus an S3 prefix is the entire MME contract on
-        # the control plane; the runtime resolves ``TargetModel`` relative to
-        # ModelDataUrl.
-        sagemaker_client.create_model(
-            ModelName=base_model_name,
-            ExecutionRoleArn=sagemaker_role_arn,
-            PrimaryContainer={
-                "Image": inference_image_uri,
-                "ModelDataUrl": s3_model_prefix,
-                "Mode": "MultiModel",
-            },
+        model_name = unique_name("tf220-mme-model")
+        cleanup_endpoint(endpoint_name, model_name=model_name)
+
+        # 1. Create a multi-model SageMaker Model. The MME contract is
+        #    expressed at the container definition level: mode="MultiModel"
+        #    plus an S3 *prefix* (not a single tar) in model_data_url.
+        Model.create(
+            model_name=model_name,
+            primary_container=ContainerDefinition(
+                image=inference_image_uri,
+                mode="MultiModel",
+                model_data_url=s3_model_prefix,
+            ),
+            execution_role_arn=sagemaker_role_arn,
+            session=boto_session,
         )
 
-        sagemaker_client.create_endpoint_config(
-            EndpointConfigName=endpoint_name,
-            ProductionVariants=[
-                {
-                    "VariantName": "AllTraffic",
-                    "ModelName": base_model_name,
-                    "InitialInstanceCount": 1,
-                    "InstanceType": INSTANCE_TYPE,
-                    "InitialVariantWeight": 1.0,
-                }
+        # 2. Endpoint config + endpoint — same shape as single-model.
+        EndpointConfig.create(
+            endpoint_config_name=endpoint_name,
+            production_variants=[
+                ProductionVariant(
+                    variant_name="AllTraffic",
+                    model_name=model_name,
+                    initial_instance_count=1,
+                    instance_type=INSTANCE_TYPE,
+                ),
             ],
+            session=boto_session,
         )
 
-        sagemaker_client.create_endpoint(
-            EndpointName=endpoint_name,
-            EndpointConfigName=endpoint_name,
+        endpoint = Endpoint.create(
+            endpoint_name=endpoint_name,
+            endpoint_config_name=endpoint_name,
+            session=boto_session,
         )
-        wait_for_endpoint(endpoint_name)
+        endpoint.wait_for_status("InService")
 
         payload = json.dumps({"instances": [[1.0, 2.0, 3.0]]})
 
-        # Invoke model1 (x * 2.0)
-        resp1 = sagemaker_runtime_client.invoke_endpoint(
-            EndpointName=endpoint_name,
-            ContentType="application/json",
-            TargetModel="model1.tar.gz",
-            Body=payload,
+        # 3. Invoke each model by name. ``target_model`` maps to the
+        #    X-Amzn-SageMaker-Target-Model header that selects the tarball
+        #    within the MME's S3 prefix.
+        resp1 = endpoint.invoke(
+            body=payload,
+            content_type="application/json",
+            accept="application/json",
+            target_model="model1.tar.gz",
         )
-        body1 = _decode(resp1)
+        body1 = json.loads(resp1.body.read().decode("utf-8"))
         assert "predictions" in body1, f"model1 response missing predictions: {body1!r}"
         values1 = _values_from_predictions(body1["predictions"])
         assert values1 == pytest.approx([2.0, 4.0, 6.0]), f"model1 got {values1!r}"
 
-        # Invoke model2 (x * 3.0)
-        resp2 = sagemaker_runtime_client.invoke_endpoint(
-            EndpointName=endpoint_name,
-            ContentType="application/json",
-            TargetModel="model2.tar.gz",
-            Body=payload,
+        resp2 = endpoint.invoke(
+            body=payload,
+            content_type="application/json",
+            accept="application/json",
+            target_model="model2.tar.gz",
         )
-        body2 = _decode(resp2)
+        body2 = json.loads(resp2.body.read().decode("utf-8"))
         assert "predictions" in body2, f"model2 response missing predictions: {body2!r}"
         values2 = _values_from_predictions(body2["predictions"])
         assert values2 == pytest.approx([3.0, 6.0, 9.0]), f"model2 got {values2!r}"
diff --git a/test/tensorflow/integration/inference/test_single_model_endpoint.py b/test/tensorflow/integration/inference/test_single_model_endpoint.py

Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,3 @@`
`1`	`1`	`boto3`
`2`	`2`	`pytest`
	`3`	`+sagemaker>=3.0.0`