Merge pull request #40 from zhaoqizqwang/add-inference-classes

Make function classmethod and update unit tests

Author: zhaoqizqwang

sagemaker-hyperpod/setup.py

@@ -16,11 +16,12 @@
         "pyyaml==6.0.2",
         "ratelimit==2.2.1",
         "tabulate==0.9.0",
+        "pydantic==2.11.7",
         "pytest==8.3.2",
         "pytest-cov==5.0.0",
         "pytest-order==1.3.0",
         "tox==4.18.0",
         "ruff==0.6.2",
         "hera-workflows==5.16.3",
     ],
-)
+)

sagemaker-hyperpod/src/sagemaker/hyperpod/hyperpod_manager.py

@@ -125,7 +125,7 @@ def list_clusters(
 
         print(tabulate(table_data, headers=headers))
 
-    def set_context_cluster(
+    def set_context(
         self,
         cluster_name: str,
         region: Optional[str] = None,

sagemaker-hyperpod/src/sagemaker/hyperpod/inference/config/constants.py

@@ -10,6 +10,7 @@
 INFERENCE_ENDPOINT_CONFIG_KIND = "InferenceEndpointConfig"
 INFERENCE_ENDPOINT_CONFIG_PLURAL = "inferenceendpointconfigs"
 MODEL_HUB_NAME = "SageMakerPublicHub"
+DEFAULT_MOUNT_PATH = "/opt/ml/model"
 
 KIND_PLURAL_MAP = {
     JUMPSTART_MODEL_KIND: JUMPSTART_MODEL_PLURAL,

sagemaker-hyperpod/src/sagemaker/hyperpod/inference/hp_endpoint.py

@@ -11,7 +11,8 @@
 )
 from sagemaker.hyperpod.inference.hp_endpoint_base import HPEndpointBase
 from datetime import datetime
-from typing import Union, Dict, Literal
+from typing import Dict, Literal
+import boto3
 
 
 class HPEndpoint(HPEndpointBase):
@@ -93,9 +94,11 @@ def _validate_inputs(
                 raise TypeError(
                     f"fsx_mount_name must be of type str, got {type(fsx_mount_name)}"
                 )
-        
+
         # Validate model_volume_mount_name if provided
-        if model_volume_mount_name is not None and not isinstance(model_volume_mount_name, str):
+        if model_volume_mount_name is not None and not isinstance(
+            model_volume_mount_name, str
+        ):
             raise TypeError(
                 f"model_volume_mount_name must be of type str, got {type(model_volume_mount_name)}"
             )
@@ -114,9 +117,10 @@ def _get_default_endpoint_name(self, model_name):
 
         return model_name + "-" + time_str
 
+    @classmethod
     def create(
-        self,
-        namespace: str,
+        cls,
+        namespace: str = None,
         model_name: str = None,
         model_version: str = None,
         instance_type: str = None,
@@ -132,7 +136,9 @@ def create(
         model_volume_mount_name: str = None,
         model_volume_mount_path: str = None,
     ):
-        self._validate_inputs(
+        instance = cls()
+
+        instance._validate_inputs(
             model_name,
             instance_type,
             image,
@@ -143,16 +149,14 @@ def create(
             fsx_dns_name,
             fsx_file_system_id,
             fsx_mount_name,
-            endpoint_name,
             model_volume_mount_name,
-            model_volume_mount_path,
         )
 
         if not endpoint_name:
-            endpoint_name = self._get_default_endpoint_name(model_name)
+            endpoint_name = instance._get_default_endpoint_name(model_name)
 
         if not model_volume_mount_path:
-            model_volume_mount_path = "/opt/ml/model"
+            model_volume_mount_path = DEFAULT_MOUNT_PATH
 
         if model_source_type == "s3":
             model_source_config = ModelSourceConfig(
@@ -181,7 +185,7 @@ def create(
             model_invocation_port=ModelInvocationPort(container_port=container_port),
             resources=Resources(),
         )
-        
+
         # create spec config
         spec = InferenceEndpointConfigSpec(
             instance_type=instance_type,
@@ -192,56 +196,80 @@ def create(
             endpoint_name=endpoint_name,
         )
 
-        self.call_create_api(
+        instance.call_create_api(
             name=spec.modelName,  # use model name as metadata name
             kind=INFERENCE_ENDPOINT_CONFIG_KIND,
             namespace=namespace,
             spec=spec,
         )
 
+        return super().get_endpoint(endpoint_name)
+
+    @classmethod
     def create_from_spec(
-        self,
+        cls,
         spec: InferenceEndpointConfigSpec,
         namespace: str = None,
     ):
-        self.call_create_api(
+        cls().call_create_api(
             name=spec.modelName,  # use model name as metadata name
             kind=INFERENCE_ENDPOINT_CONFIG_KIND,
             namespace=namespace,
             spec=spec,
         )
 
-    def create_from_dict(self, input: Dict, namespace: str):
+        region = boto3.session.Session().region_name
+
+        return super().get_endpoint(endpoint_name=spec.endpointName, region=region)
+
+    @classmethod
+    def create_from_dict(
+        cls,
+        input: Dict,
+        namespace: str = None,
+    ):
         spec = InferenceEndpointConfigSpec.model_validate(input, by_name=True)
 
-        self.call_create_api(
+        cls().call_create_api(
+            name=spec.modelName,  # use model name as metadata name
+            kind=INFERENCE_ENDPOINT_CONFIG_KIND,
             namespace=namespace,
             spec=spec,
         )
 
-    def list_endpoints(self, namespace: str):
-        return self.call_list_api(
+        region = boto3.session.Session().region_name
+
+        return super().get_endpoint(endpoint_name=spec.endpointName, region=region)
+
+    @classmethod
+    def list_endpoints(
+        cls,
+        namespace: str = None,
+    ):
+        return cls().call_list_api(
             kind=INFERENCE_ENDPOINT_CONFIG_KIND,
             namespace=namespace,
         )
 
+    @classmethod
     def describe_endpoint(
-        self,
+        cls,
         name: str,
-        namespace: str,
+        namespace: str = None,
     ):
-        return self.call_get_api(
+        return cls().call_get_api(
             name=name,
             kind=INFERENCE_ENDPOINT_CONFIG_KIND,
             namespace=namespace,
         )
 
+    @classmethod
     def delete_endpoint(
-        self,
+        cls,
         name: str,
-        namespace: str,
+        namespace: str = None,
     ):
-        return self.call_delete_api(
+        cls().call_delete_api(
             name=name,  # use model id as metadata name
             kind=INFERENCE_ENDPOINT_CONFIG_KIND,
             namespace=namespace,

sagemaker-hyperpod/src/sagemaker/hyperpod/inference/hp_endpoint_base.py

@@ -9,9 +9,35 @@
     InferenceEndpointConfigSpec,
 )
 from types import SimpleNamespace
+import boto3
+from sagemaker_core.main.resources import Endpoint
+
+
+def get_current_region():
+    session = boto3.session.Session()
+    return session.region_name
 
 
 class HPEndpointBase:
+    _endpoint = None
+
+    def get_name(self):
+        if not self._endpoint or not hasattr(self._endpoint, "endpoint_name"):
+            print(f"Endpoint is not set!")
+
+        print(f"Endpoint name is: {self._endpoint.endpoint_name}")
+
+    @classmethod
+    def get_endpoint(
+        cls,
+        endpoint_name: str,
+        region: str = None,
+    ):
+        if not region:
+            region = get_current_region()
+
+        return Endpoint.get(endpoint_name, region=region)
+
     def _validate_connection(self):
         try:
             k8s_config.load_kube_config()
@@ -51,6 +77,8 @@ def call_create_api(
                 plural=KIND_PLURAL_MAP[kind],
                 body=body,
             )
+
+            self.set_endpoint(spec.sageMakerEndpoint.name)
             print("\nSuccessful deployed model and its endpoint!")
         except Exception as e:
             print(f"\nFailed to deploy model and its endpoint: {e}")
@@ -125,3 +153,12 @@ def call_list_api(
             )
         except Exception as e:
             print(f"\nFailed to list endpoint: {e}")
+
+    def invoke(self, body, **kwargs):
+        if self._endpoint is None:
+            raise Exception("Endpoint not initialized. Please set endpoint first.")
+
+        try:
+            self._endpoint.invoke(body, **kwargs)
+        except Exception as e:
+            print(f"\nFailed to invoke endpoint: {e}")

sagemaker-hyperpod/src/sagemaker/hyperpod/inference/hp_jumpstart_endpoint.py

@@ -18,7 +18,6 @@ def _validate_inputs(
         model_id: str,
         instance_type: str,
     ):
-
         # Validate required parameters when spec is None
         if model_id is None or instance_type is None:
             raise ValueError("Must provide both model_id and instance_type.")
@@ -45,19 +44,21 @@ def _get_default_endpoint_name(self, model_id: str):
 
         return model_id + "-" + time_str
 
+    @classmethod
     def create(
-        self,
-        namespace: str,
+        cls,
+        namespace: str = None,
         model_id: str = None,
         model_version: str = None,
         instance_type: str = None,
         sagemaker_endpoint: str = None,
         accept_eula: bool = False,
     ):
-        self._validate_inputs(model_id, instance_type)
+        instance = cls()
+        instance._validate_inputs(model_id, instance_type)
 
         if not sagemaker_endpoint:
-            sagemaker_endpoint = self._get_default_endpoint_name(model_id)
+            sagemaker_endpoint = instance._get_default_endpoint_name(model_id)
 
         spec = JumpStartModelSpec(
             model=Model(
@@ -69,35 +70,53 @@ def create(
             sage_maker_endpoint=SageMakerEndpoint(name=sagemaker_endpoint),
         )
 
-        self.call_create_api(
+        instance.call_create_api(
             name=spec.model.modelId,  # use model id as metadata name
             kind=JUMPSTART_MODEL_KIND,
             namespace=namespace,
             spec=spec,
         )
 
+        return cls.get_endpoint(endpoint_name=sagemaker_endpoint)
+
+    @classmethod
     def create_from_spec(
-        self,
+        cls,
         spec: JumpStartModelSpec,
         namespace: str = None,
     ):
-        self.call_create_api(
+        cls().call_create_api(
             name=spec.model.modelId,  # use model id as metadata name
             kind=JUMPSTART_MODEL_KIND,
             namespace=namespace,
             spec=spec,
         )
 
-    def create_from_dict(self, input: Dict, namespace: str):
+        return super().get_endpoint(endpoint_name=spec.sageMakerEndpoint.name)
+
+    @classmethod
+    def create_from_dict(
+        cls,
+        input: Dict,
+        namespace: str = None,
+    ):
         spec = JumpStartModelSpec.model_validate(input, by_name=True)
 
-        self.call_create_api(
+        cls().call_create_api(
+            name=spec.model.modelId,  # use model id as metadata name
+            kind=JUMPSTART_MODEL_KIND,
             namespace=namespace,
             spec=spec,
         )
 
-    def list_endpoints(self, namespace: str):
-        response = self.call_list_api(
+        return super().get_endpoint(endpoint_name=spec.sageMakerEndpoint.name)
+
+    @classmethod
+    def list_endpoints(
+        cls,
+        namespace: str = None,
+    ):
+        response = cls().call_list_api(
             kind=JUMPSTART_MODEL_KIND,
             namespace=namespace,
         )
@@ -109,27 +128,31 @@ def list_endpoints(self, namespace: str):
         headers = ["METADATA NAME", "CREATE TIME"]
 
         print(tabulate(output_data, headers=headers))
+        return response
 
+    @classmethod
     def describe_endpoint(
-        self,
+        cls,
         name: str,
-        namespace: str,
+        namespace: str = None,
     ):
-        response = self.call_get_api(
+        response = cls().call_get_api(
             name=name,
             kind=JUMPSTART_MODEL_KIND,
             namespace=namespace,
         )
 
         response["metadata"].pop("managedFields")
         print(yaml.dump(response))
+        return response
 
+    @classmethod
     def delete_endpoint(
-        self,
+        cls,
         name: str,
-        namespace: str,
+        namespace: str = None,
     ):
-        return self.call_delete_api(
+        return cls().call_delete_api(
             name=name,  # use model id as metadata name
             kind=JUMPSTART_MODEL_KIND,
             namespace=namespace,