fix: migrate inference integration tests to SageMaker SDK v3 API

The PyPI sagemaker package v3.x removed the top-level sagemaker.Session
class along with sagemaker.tensorflow.serving.TensorFlowModel and
sagemaker.multidatamodel.MultiDataModel. Our conftest used the legacy v2
API. Rewrite to use boto3 directly (sagemaker, sagemaker-runtime, s3
clients) so tests work with the >=3.0.0 pin and are SDK-version
independent. boto3 maps 1:1 to the create_model / create_endpoint_config
/ create_endpoint / invoke_endpoint flow these tests already needed,
which is simpler than adopting v3's heavier ModelBuilder abstraction
for integration tests. The sagemaker dep is dropped from
requirements.txt since it is no longer imported.

Author: Bhanu Teja Goshikonda

test/tensorflow/integration/inference/conftest.py

@@ -1,5 +1,12 @@
 """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.
+
 Fixtures intentionally defer all AWS calls until test-execution time so that
 ``pytest --collect-only`` works in environments without AWS credentials.
 """
@@ -46,11 +53,51 @@ def boto_session(aws_region: str):
 
 
 @pytest.fixture(scope="session")
-def sagemaker_session(boto_session):
-    """A SageMaker SDK session for high-level deploy/predict calls."""
-    import sagemaker
+def sagemaker_client(boto_session):
+    """Low-level SageMaker control-plane client (create/delete model, endpoint, ...)."""
+    return boto_session.client("sagemaker")
+
 
-    return sagemaker.Session(boto_session=boto_session)
+@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.
+
+    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.
+    """
+    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
 
 
 @pytest.fixture
@@ -68,13 +115,61 @@ def _make(prefix: str) -> str:
 
 
 @pytest.fixture
-def cleanup_endpoint(sagemaker_session):
+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):
     """Yield-style fixture that tears down endpoint, endpoint config, and model.
 
     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] = []
 
@@ -83,25 +178,22 @@ def _register(endpoint_name: str, model_name: str | None = None) -> None:
 
     yield _register
 
-    sm_client = sagemaker_session.boto_session.client("sagemaker")
     for item in registered:
         endpoint_name = item["endpoint_name"]
         model_name = item["model_name"]
 
         for delete_call, kwargs in (
-            (sm_client.delete_endpoint, {"EndpointName": endpoint_name}),
-            (sm_client.delete_endpoint_config, {"EndpointConfigName": endpoint_name}),
+            (sagemaker_client.delete_endpoint, {"EndpointName": endpoint_name}),
+            (sagemaker_client.delete_endpoint_config, {"EndpointConfigName": endpoint_name}),
         ):
             try:
                 delete_call(**kwargs)
-            except sm_client.exceptions.ClientError:
-                # Swallow NotFound / already-deleted; teardown should be best-effort.
-                pass
             except Exception:
+                # Swallow NotFound / already-deleted; teardown is best-effort.
                 pass
 
         if model_name:
             try:
-                sm_client.delete_model(ModelName=model_name)
+                sagemaker_client.delete_model(ModelName=model_name)
             except Exception:
                 pass

test/tensorflow/integration/inference/requirements.txt

@@ -1,3 +1,2 @@
 boto3
 pytest
-sagemaker>=3.0.0

test/tensorflow/integration/inference/test_multi_model_endpoint.py

@@ -3,6 +3,12 @@
 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.
 """
 
 from __future__ import annotations
@@ -18,9 +24,9 @@
 INSTANCE_TYPE = "ml.c5.xlarge"
 
 
-def _decode(response) -> dict:
-    """Normalize an MME runtime invoke response Body into a Python dict."""
-    body = response["Body"].read() if hasattr(response, "get") else response
+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)
@@ -36,15 +42,16 @@ def _values_from_predictions(predictions) -> list:
 
 
 def test_mme_two_models(
-    sagemaker_session,
+    sagemaker_client,
+    sagemaker_runtime_client,
     sagemaker_role_arn,
     inference_image_uri,
+    default_bucket,
+    upload_to_s3,
     unique_name,
     cleanup_endpoint,
+    wait_for_endpoint,
 ):
-    from sagemaker.multidatamodel import MultiDataModel
-    from sagemaker.tensorflow.serving import TensorFlowModel
-
     with tempfile.TemporaryDirectory(prefix="tf220-mme-") as workdir:
         workdir_path = Path(workdir)
 
@@ -59,48 +66,55 @@ 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.
-        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}/"
+        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}/"
 
         endpoint_name = unique_name("tf220-mme")
         base_model_name = unique_name("tf220-mme-model")
         cleanup_endpoint(endpoint_name, model_name=base_model_name)
 
-        # The TensorFlowModel acts as the container template; MultiDataModel
-        # reuses its image_uri + execution role for the endpoint config.
-        base_model = TensorFlowModel(
-            model_data=model1_tar,  # placeholder; MME ignores model_data at runtime
-            role=sagemaker_role_arn,
-            image_uri=inference_image_uri,
-            sagemaker_session=sagemaker_session,
-            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",
+            },
         )
 
-        mme = MultiDataModel(
-            name=base_model_name,
-            model_data_prefix=s3_model_prefix,
-            model=base_model,
-            sagemaker_session=sagemaker_session,
+        sagemaker_client.create_endpoint_config(
+            EndpointConfigName=endpoint_name,
+            ProductionVariants=[
+                {
+                    "VariantName": "AllTraffic",
+                    "ModelName": base_model_name,
+                    "InitialInstanceCount": 1,
+                    "InstanceType": INSTANCE_TYPE,
+                    "InitialVariantWeight": 1.0,
+                }
+            ],
         )
 
-        mme.deploy(
-            initial_instance_count=1,
-            instance_type=INSTANCE_TYPE,
-            endpoint_name=endpoint_name,
+        sagemaker_client.create_endpoint(
+            EndpointName=endpoint_name,
+            EndpointConfigName=endpoint_name,
         )
+        wait_for_endpoint(endpoint_name)
 
-        runtime = sagemaker_session.boto_session.client("sagemaker-runtime")
         payload = json.dumps({"instances": [[1.0, 2.0, 3.0]]})
 
         # Invoke model1 (x * 2.0)
-        resp1 = runtime.invoke_endpoint(
+        resp1 = sagemaker_runtime_client.invoke_endpoint(
             EndpointName=endpoint_name,
             ContentType="application/json",
             TargetModel="model1.tar.gz",
@@ -112,7 +126,7 @@ def test_mme_two_models(
         assert values1 == pytest.approx([2.0, 4.0, 6.0]), f"model1 got {values1!r}"
 
         # Invoke model2 (x * 3.0)
-        resp2 = runtime.invoke_endpoint(
+        resp2 = sagemaker_runtime_client.invoke_endpoint(
             EndpointName=endpoint_name,
             ContentType="application/json",
             TargetModel="model2.tar.gz",