integration test for jumpstart with mig profile (#334)

* integration test for jumpstart with mig profile

* template fix for mig with jumpstart

* skipped mig tests until instances setup finished

---------

Co-authored-by: Chad Chiang <[email protected]>

Author: chad119

hyperpod-jumpstart-inference-template/hyperpod_jumpstart_inference_template/v1_1/template.py

@@ -1,8 +1,8 @@
 TEMPLATE_CONTENT = """
-apiVersion: inference.sagemaker.aws.amazon.com/v1alpha1
+apiVersion: inference.sagemaker.aws.amazon.com/v1
 kind: JumpStartModel
 metadata:
-  name:                {{ model_id }}
+  name:                {{ metadata_name or endpoint_name }}
   namespace:           {{ namespace or "default" }}
 spec:
   model:
@@ -18,4 +18,6 @@
     {% if accelerator_partition_validation is not none %}validations: 
     {% if accelerator_partition_validation is not none %}  acceleratorPartitionValidation: {{ accelerator_partition_validation }}{% endif %}
     {% endif %}
+  tlsConfig:
+    tlsCertificateOutputS3Uri: {{ tls_certificate_output_s3_uri or "" }}
 """

test/integration_tests/inference/cli/test_cli_jumpstart_inference_with_mig.py

@@ -0,0 +1,121 @@
+import time
+import pytest
+import boto3
+from click.testing import CliRunner
+from sagemaker.hyperpod.cli.commands.inference import (
+    js_create, custom_invoke, js_list, js_describe, js_delete, js_get_operator_logs, js_list_pods
+)
+from sagemaker.hyperpod.inference.hp_jumpstart_endpoint import HPJumpStartEndpoint
+from test.integration_tests.utils import get_time_str
+
+# --------- Test Configuration ---------
+NAMESPACE = "integration"
+VERSION = "1.1"
+REGION = "us-east-2"
+TIMEOUT_MINUTES = 20
+POLL_INTERVAL_SECONDS = 30
+
+@pytest.fixture(scope="module")
+def runner():
+    return CliRunner()
+
+@pytest.fixture(scope="module")
+def js_endpoint_name():
+    return "js-mig-cli-integration-" + get_time_str()
+
+@pytest.fixture(scope="module")
+def sagemaker_client():
+    return boto3.client("sagemaker", region_name=REGION)
+
+# --------- JumpStart Endpoint Tests ---------
+@pytest.mark.skip(reason="Temporarily disabled")
+@pytest.mark.dependency(name="create")
+def test_js_create(runner, js_endpoint_name):
+    result = runner.invoke(js_create, [
+        "--namespace", NAMESPACE,
+        "--version", VERSION,
+        "--model-id", "deepseek-llm-r1-distill-qwen-1-5b",
+        "--instance-type", "ml.p4d.24xlarge",
+        "--endpoint-name", js_endpoint_name,
+        "--accelerator-partition-type", "mig-7g.40gb",
+        "--accelerator-partition-validation", "true",
+    ])
+    assert result.exit_code == 0, result.output
+
+
+@pytest.mark.dependency(depends=["create"])
+def test_js_list(runner, js_endpoint_name):
+    result = runner.invoke(js_list, ["--namespace", NAMESPACE])
+    assert result.exit_code == 0
+    assert js_endpoint_name in result.output
+
+
+@pytest.mark.dependency(name="describe", depends=["create"])
+def test_js_describe(runner, js_endpoint_name):
+    result = runner.invoke(js_describe, [
+        "--name", js_endpoint_name,
+        "--namespace", NAMESPACE,
+        "--full"
+    ])
+    assert result.exit_code == 0
+    assert js_endpoint_name in result.output
+
+
+@pytest.mark.dependency(depends=["create", "describe"])
+def test_wait_until_inservice(js_endpoint_name):
+    """Poll SDK until specific JumpStart endpoint reaches DeploymentComplete"""
+    print(f"[INFO] Waiting for JumpStart endpoint '{js_endpoint_name}' to be DeploymentComplete...")
+    deadline = time.time() + (TIMEOUT_MINUTES * 60)
+    poll_count = 0
+
+    while time.time() < deadline:
+        poll_count += 1
+        print(f"[DEBUG] Poll #{poll_count}: Checking endpoint status...")
+
+        try:
+            ep = HPJumpStartEndpoint.get(name=js_endpoint_name, namespace=NAMESPACE)
+            state = ep.status.endpoints.sagemaker.state
+            print(f"[DEBUG] Current state: {state}")
+            if state == "CreationCompleted":
+                print("[INFO] Endpoint is in CreationCompleted state.")
+                return
+
+            deployment_state = ep.status.deploymentStatus.deploymentObjectOverallState
+            if deployment_state == "DeploymentFailed":
+                pytest.fail("Endpoint deployment failed.")
+
+        except Exception as e:
+            print(f"[ERROR] Exception during polling: {e}")
+
+        time.sleep(POLL_INTERVAL_SECONDS)
+
+    pytest.fail("[ERROR] Timed out waiting for endpoint to be DeploymentComplete")
+
+
+@pytest.mark.dependency(depends=["create"])
+def test_custom_invoke(runner, js_endpoint_name):
+    result = runner.invoke(custom_invoke, [
+        "--endpoint-name", js_endpoint_name,
+        "--body", '{"inputs": "What is the capital of USA?"}'
+    ])
+    assert result.exit_code == 0
+    assert "error" not in result.output.lower()
+
+@pytest.mark.skip(reason="Temporarily disabled")
+def test_js_get_operator_logs(runner):
+    result = runner.invoke(js_get_operator_logs, ["--since-hours", "1"])
+    assert result.exit_code == 0
+
+@pytest.mark.skip(reason="Temporarily disabled")
+def test_js_list_pods(runner):
+    result = runner.invoke(js_list_pods, ["--namespace", NAMESPACE])
+    assert result.exit_code == 0
+
+
+@pytest.mark.dependency(depends=["create"])
+def test_js_delete(runner, js_endpoint_name):
+    result = runner.invoke(js_delete, [
+        "--name", js_endpoint_name,
+        "--namespace", NAMESPACE
+    ])
+    assert result.exit_code == 0

test/integration_tests/inference/sdk/test_sdk_jumpstart_inference_with_mig.py

@@ -0,0 +1,123 @@
+import time
+import pytest
+import boto3
+from sagemaker.hyperpod.inference.hp_jumpstart_endpoint import HPJumpStartEndpoint
+from sagemaker.hyperpod.inference.config.hp_jumpstart_endpoint_config import (
+    Model, Server, SageMakerEndpoint, Validations
+)
+import sagemaker_core.main.code_injection.codec as codec
+from test.integration_tests.utils import get_time_str
+from sagemaker.hyperpod.common.config.metadata import Metadata
+
+# --------- Config ---------
+NAMESPACE = "integration"
+REGION = "us-east-2"
+ENDPOINT_NAME = "js-mig-sdk-integration-" + get_time_str()
+
+INSTANCE_TYPE = "ml.p4d.24xlarge"
+MODEL_ID = "deepseek-llm-r1-distill-qwen-1-5b"
+
+TIMEOUT_MINUTES = 20
+POLL_INTERVAL_SECONDS = 30
+
+@pytest.fixture(scope="module")
+def sagemaker_client():
+    return boto3.client("sagemaker", region_name=REGION)
+
+@pytest.fixture(scope="module")
+def endpoint_obj():
+    model = Model(model_id=MODEL_ID)
+    validations = Validations(accelerator_partition_validation=True)
+    server = Server(
+        instance_type=INSTANCE_TYPE,
+        accelerator_partition_type="mig-7g.40gb",
+        validations=validations
+    )
+    sm_endpoint = SageMakerEndpoint(name=ENDPOINT_NAME)
+    metadata = Metadata(name=ENDPOINT_NAME, namespace=NAMESPACE)
+
+    return HPJumpStartEndpoint(metadata=metadata, model=model, server=server, sage_maker_endpoint=sm_endpoint)
+
+@pytest.mark.skip(reason="Temporarily disabled")
+@pytest.mark.dependency(name="create")
+def test_create_endpoint(endpoint_obj):
+    endpoint_obj.create()
+    assert endpoint_obj.metadata.name == ENDPOINT_NAME
+
+@pytest.mark.dependency(depends=["create"])
+def test_list_endpoint():
+    endpoints = HPJumpStartEndpoint.list(namespace=NAMESPACE)
+    names = [ep.metadata.name for ep in endpoints]
+    assert ENDPOINT_NAME in names
+
+@pytest.mark.dependency(name="describe", depends=["create"])
+def test_get_endpoint():
+    ep = HPJumpStartEndpoint.get(name=ENDPOINT_NAME, namespace=NAMESPACE)
+    assert ep.metadata.name == ENDPOINT_NAME
+    assert ep.model.modelId == MODEL_ID
+
+@pytest.mark.dependency(depends=["create", "describe"])
+def test_wait_until_inservice():
+    """Poll SDK until specific JumpStart endpoint reaches DeploymentComplete"""
+    print(f"[INFO] Waiting for JumpStart endpoint '{ENDPOINT_NAME}' to be DeploymentComplete...")
+    deadline = time.time() + (TIMEOUT_MINUTES * 60)
+    poll_count = 0
+
+    while time.time() < deadline:
+        poll_count += 1
+        print(f"[DEBUG] Poll #{poll_count}: Checking endpoint status...")
+
+        try:
+            ep = HPJumpStartEndpoint.get(name=ENDPOINT_NAME, namespace=NAMESPACE)
+            state = ep.status.endpoints.sagemaker.state
+            print(f"[DEBUG] Current state: {state}")
+            if state == "CreationCompleted":
+                print("[INFO] Endpoint is in CreationCompleted state.")
+                return
+
+            deployment_state = ep.status.deploymentStatus.deploymentObjectOverallState
+            if deployment_state == "DeploymentFailed":
+                pytest.fail("Endpoint deployment failed.")
+
+        except Exception as e:
+            print(f"[ERROR] Exception during polling: {e}")
+
+        time.sleep(POLL_INTERVAL_SECONDS)
+
+    pytest.fail("[ERROR] Timed out waiting for endpoint to be DeploymentComplete")
+
+
+@pytest.mark.dependency(depends=["create"])
+def test_invoke_endpoint(monkeypatch):
+    original_transform = codec.transform  # Save original
+
+    def mock_transform(data, shape, object_instance=None):
+        if "Body" in data:
+            return {"body": data["Body"].read().decode("utf-8")}
+        return original_transform(data, shape, object_instance)  # Call original
+
+    monkeypatch.setattr("sagemaker_core.main.resources.transform", mock_transform)
+
+    ep = HPJumpStartEndpoint.get(name=ENDPOINT_NAME, namespace=NAMESPACE)
+    data = '{"inputs":"What is the capital of USA?"}'
+    response = ep.invoke(body=data)
+    
+    assert "error" not in response.body.lower()
+
+
+@pytest.mark.skip(reason="Temporarily disabled")
+def test_get_operator_logs():
+    ep = HPJumpStartEndpoint.get(name=ENDPOINT_NAME, namespace=NAMESPACE)
+    logs = ep.get_operator_logs(since_hours=1)
+    assert logs
+
+@pytest.mark.skip(reason="Temporarily disabled")
+def test_list_pods():
+    ep = HPJumpStartEndpoint.get(name=ENDPOINT_NAME, namespace=NAMESPACE)
+    pods = ep.list_pods(NAMESPACE)
+    assert pods
+
+@pytest.mark.dependency(depends=["create"])
+def test_delete_endpoint():
+    ep = HPJumpStartEndpoint.get(name=ENDPOINT_NAME, namespace=NAMESPACE)
+    ep.delete()