ODSC 39392/triton (#128)

Author: z7ye

ads/model/deployment/common/utils.py

@@ -119,19 +119,14 @@ def send_request(
     Returns:
         A JSON representive of a requests.Response object.
     """
-    headers = dict()
     if is_json_payload:
-        headers["Content-Type"] = (
-            header.get("content_type") or DEFAULT_CONTENT_TYPE_JSON
-        )
+        header["Content-Type"] =  header.pop("content_type", DEFAULT_CONTENT_TYPE_JSON) or DEFAULT_CONTENT_TYPE_JSON
         request_kwargs = {"json": data}
     else:
-        headers["Content-Type"] = (
-            header.get("content_type") or DEFAULT_CONTENT_TYPE_BYTES
-        )
+        header["Content-Type"] = header.pop("content_type", DEFAULT_CONTENT_TYPE_BYTES) or DEFAULT_CONTENT_TYPE_BYTES
         request_kwargs = {"data": data}  # should pass bytes when using data
-
-    request_kwargs["headers"] = headers
+    
+    request_kwargs["headers"] = header
 
     if dry_run:
         request_kwargs["headers"]["Accept"] = "*/*"
@@ -140,7 +135,7 @@ def send_request(
             return json.loads(req.body)
         return req.body
     else:
-        request_kwargs["auth"] = header.get("signer")
+        request_kwargs["auth"] = header.pop("signer")
         return requests.post(endpoint, **request_kwargs).json()
 
 

ads/model/deployment/model_deployment.py

@@ -17,7 +17,6 @@
 from ads.common import auth as authutil
 import pandas as pd
 from ads.model.serde.model_input import JsonModelInputSERDE
-from ads.common import auth, oci_client
 from ads.common.oci_logging import (
     LOG_INTERVAL,
     LOG_RECORDS_LIMIT,
@@ -63,6 +62,7 @@
 
 MODEL_DEPLOYMENT_KIND = "deployment"
 MODEL_DEPLOYMENT_TYPE = "modelDeployment"
+MODEL_DEPLOYMENT_INFERENCE_SERVER_TRITON = "TRITON"
 
 MODEL_DEPLOYMENT_INSTANCE_SHAPE = "VM.Standard.E4.Flex"
 MODEL_DEPLOYMENT_INSTANCE_OCPUS = 1
@@ -828,6 +828,8 @@ def predict(
         data: Any = None,
         serializer: "ads.model.ModelInputSerializer" = model_input_serializer,
         auto_serialize_data: bool = False,
+        model_name: str = None,
+        model_version: str = None,
         **kwargs,
     ) -> dict:
         """Returns prediction of input data run against the model deployment endpoint.
@@ -860,6 +862,10 @@ def predict(
             If `auto_serialize_data=False`, `data` required to be bytes or json serializable
             and `json_input` required to be json serializable. If `auto_serialize_data` set
             to True, data will be serialized before sending to model deployment endpoint.
+        model_name: str
+            Defaults to None. When the `Inference_server="triton"`, the name of the model to invoke.
+        model_version: str
+            Defaults to None. When the `Inference_server="triton"`, the version of the model to invoke.
         kwargs:
             content_type: str
                 Used to indicate the media type of the resource.
@@ -878,6 +884,7 @@ def predict(
             "signer": signer,
             "content_type": kwargs.get("content_type", None),
         }
+        header.update(kwargs.pop("headers", {}))
 
         if data is None and json_input is None:
             raise AttributeError(
@@ -916,9 +923,13 @@ def predict(
             raise TypeError(
                 "`data` is not bytes or json serializable. Set `auto_serialize_data` to `True` to serialize the input data."
             )
-
+        if model_name and model_version:
+            header['model-name'] = model_name
+            header['model-version'] = model_version
+        elif bool(model_version) ^ bool(model_name):
+            raise ValueError("`model_name` and `model_version` have to be provided together.")
         prediction = send_request(
-            data=data, endpoint=endpoint, is_json_payload=is_json_payload, header=header
+            data=data, endpoint=endpoint, is_json_payload=is_json_payload, header=header,
         )
         return prediction
 
@@ -1390,6 +1401,10 @@ def _update_from_oci_model(self, oci_model_instance) -> "ModelDeployment":
             infrastructure.CONST_WEB_CONCURRENCY,
             runtime.env.get("WEB_CONCURRENCY", None),
         )
+        if runtime.env.get("CONTAINER_TYPE", None) == MODEL_DEPLOYMENT_INFERENCE_SERVER_TRITON:
+            runtime.set_spec(
+                runtime.CONST_INFERENCE_SERVER, MODEL_DEPLOYMENT_INFERENCE_SERVER_TRITON.lower()
+            )
 
         self.set_spec(self.CONST_INFRASTRUCTURE, infrastructure)
         self.set_spec(self.CONST_RUNTIME, runtime)
@@ -1566,6 +1581,9 @@ def _build_model_deployment_configuration_details(self) -> Dict:
                 infrastructure.web_concurrency
             )
             runtime.set_spec(runtime.CONST_ENV, environment_variables)
+        if hasattr(runtime, "inference_server") and runtime.inference_server and runtime.inference_server.upper() == MODEL_DEPLOYMENT_INFERENCE_SERVER_TRITON:
+            environment_variables["CONTAINER_TYPE"] = MODEL_DEPLOYMENT_INFERENCE_SERVER_TRITON
+            runtime.set_spec(runtime.CONST_ENV, environment_variables)
         environment_configuration_details = {
             runtime.CONST_ENVIRONMENT_CONFIG_TYPE: runtime.environment_config_type,
             runtime.CONST_ENVIRONMENT_VARIABLES: runtime.env,

ads/model/deployment/model_deployment_runtime.py

@@ -330,6 +330,7 @@ class ModelDeploymentContainerRuntime(ModelDeploymentRuntime):
     CONST_ENTRYPOINT = "entrypoint"
     CONST_SERVER_PORT = "serverPort"
     CONST_HEALTH_CHECK_PORT = "healthCheckPort"
+    CONST_INFERENCE_SERVER = "inferenceServer"
 
     attribute_map = {
         **ModelDeploymentRuntime.attribute_map,
@@ -339,6 +340,7 @@ class ModelDeploymentContainerRuntime(ModelDeploymentRuntime):
         CONST_ENTRYPOINT: "entrypoint",
         CONST_SERVER_PORT: "server_port",
         CONST_HEALTH_CHECK_PORT: "health_check_port",
+        CONST_INFERENCE_SERVER: "inference_server"
     }
 
     payload_attribute_map = {
@@ -532,3 +534,57 @@ def with_health_check_port(
             The ModelDeploymentContainerRuntime instance (self).
         """
         return self.set_spec(self.CONST_HEALTH_CHECK_PORT, health_check_port)
+    
+    @property
+    def inference_server(self) -> str:
+        """Returns the inference server.
+
+        Returns
+        -------
+        str
+            The inference server.
+        """
+        return self.get_spec(self.CONST_INFERENCE_SERVER, None)
+
+    def with_inference_server(self, inference_server: str = "triton") -> "ModelDeploymentRuntime":
+        """Sets the inference server. Current supported inference server is "triton".
+        Note if you are using byoc, you do not need to set the inference server.
+
+        Parameters
+        ----------
+        inference_server: str
+            Set the inference server.
+
+        Returns
+        -------
+        ModelDeploymentRuntime
+            The ModelDeploymentRuntime instance (self).
+            
+        Example
+        -------
+        >>> from ads.model.deployment import ModelDeployment, ModelDeploymentContainerRuntime, ModelDeploymentInfrastructure
+        >>> import ads
+        >>> ads.set_auth("resource_principal")
+        >>> infrastructure = ModelDeploymentInfrastructure()\
+        ...                 .with_project_id(<project_id>)\
+        ...                 .with_compartment_id(<comparment_id>)\
+        ...                 .with_shape_name("VM.Standard.E4.Flex")\
+        ...                 .with_replica(2)\
+        ...                 .with_bandwidth_mbps(10)\
+        ...                 .with_access_log(log_group_id=<deployment_log_group_id>, log_id=<deployment_access_log_id>)\
+        ...                 .with_predict_log(log_group_id=<deployment_log_group_id>, log_id=<deployment_predict_log_id>)
+
+        >>> runtime = ModelDeploymentContainerRuntime()\
+        ...                 .with_image(<container_image>)\
+        ...                 .with_server_port(<server_port>)\
+        ...                 .with_health_check_port(<health_check_port>)\
+        ...                 .with_model_uri(<model_id>)\
+        ...                 .with_env({"key":"value", "key2":"value2"})\
+        ...                 .with_inference_server("triton")
+        >>> deployment = ModelDeployment()\
+        ...                 .with_display_name("Triton Example")\
+        ...                 .with_infrastructure(infrastructure)\
+        ...                 .with_runtime(runtime)
+        >>> deployment.deploy()
+        """
+        return self.set_spec(self.CONST_INFERENCE_SERVER, inference_server.lower())

tests/unitary/default_setup/model_deployment/test_model_deployment_v2.py

@@ -308,6 +308,62 @@ def initialize_model_deployment_from_spec(self):
                 "runtime": runtime,
             }
         )
+        
+    def initialize_model_deployment_triton_builder(self):
+        infrastructure = ModelDeploymentInfrastructure()\
+            .with_compartment_id("fakeid.compartment.oc1..xxx")\
+            .with_project_id("fakeid.datascienceproject.oc1.iad.xxx")\
+            .with_shape_name("VM.Standard.E4.Flex")\
+            .with_replica(2)\
+            .with_bandwidth_mbps(10)\
+
+        runtime = ModelDeploymentContainerRuntime()\
+            .with_image("fake_image")\
+            .with_server_port(5000)\
+            .with_health_check_port(5000)\
+            .with_model_uri("fake_model_id")\
+            .with_env({"key":"value", "key2":"value2"})\
+            .with_inference_server("triton")
+        
+        deployment = ModelDeployment()\
+            .with_display_name("triton case")\
+            .with_infrastructure(infrastructure)\
+            .with_runtime(runtime)
+        return deployment
+
+    def initialize_model_deployment_triton_yaml(self):
+        yaml_string = """
+kind: deployment
+spec:
+  displayName: triton
+  infrastructure:
+    kind: infrastructure
+    spec:
+      bandwidthMbps: 10
+      compartmentId: fake_compartment_id
+      deploymentType: SINGLE_MODEL
+      policyType: FIXED_SIZE
+      replica: 2
+      shapeConfigDetails:
+        memoryInGBs: 16.0
+        ocpus: 1.0
+      shapeName: VM.Standard.E4.Flex
+    type: datascienceModelDeployment
+  runtime:
+    kind: runtime
+    spec:
+      env:
+        key: value
+        key2: value2
+      inference_server: triton
+      healthCheckPort: 8000
+      image: fake_image
+      modelUri: fake_model_id
+      serverPort: 8000
+    type: container
+"""
+        deployment_from_yaml = ModelDeployment.from_yaml(yaml_string)
+        return deployment_from_yaml
 
     def initialize_model_deployment_from_kwargs(self):
         infrastructure = (
@@ -435,11 +491,34 @@ def test_initialize_model_deployment_with_error(self):
                 },
             )
 
+
     def test_initialize_model_deployment_with_spec_kwargs(self):
         model_deployment_kwargs = self.initialize_model_deployment_from_kwargs()
         model_deployment_builder = self.initialize_model_deployment()
 
         assert model_deployment_kwargs.to_dict() == model_deployment_builder.to_dict()
+        
+        
+    def test_initialize_model_deployment_triton_builder(self):
+        temp_model_deployment = self.initialize_model_deployment_triton_builder()
+        assert isinstance(
+            temp_model_deployment.runtime, ModelDeploymentContainerRuntime
+        )
+        assert isinstance(
+            temp_model_deployment.infrastructure, ModelDeploymentInfrastructure
+        )
+        assert temp_model_deployment.runtime.inference_server == "triton"
+    
+    def test_initialize_model_deployment_triton_yaml(self):
+        temp_model_deployment = self.initialize_model_deployment_triton_yaml()
+        assert isinstance(
+            temp_model_deployment.runtime, ModelDeploymentContainerRuntime
+        )
+        assert isinstance(
+            temp_model_deployment.infrastructure, ModelDeploymentInfrastructure
+        )
+        assert temp_model_deployment.runtime.inference_server == "triton"
+        
 
     def test_model_deployment_to_dict(self):
         model_deployment = self.initialize_model_deployment()