commit the code for triton

Author: Ziqun Ye

ads/model/deployment/common/utils.py

@@ -119,19 +119,15 @@ 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)
         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)
         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 +136,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

@@ -878,6 +878,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(
@@ -1390,6 +1391,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) == "TRITON":
+            runtime.set_spec(
+                runtime.CONST_TRITON, True
+            )
 
         self.set_spec(self.CONST_INFRASTRUCTURE, infrastructure)
         self.set_spec(self.CONST_RUNTIME, runtime)
@@ -1566,6 +1571,9 @@ def _build_model_deployment_configuration_details(self) -> Dict:
                 infrastructure.web_concurrency
             )
             runtime.set_spec(runtime.CONST_ENV, environment_variables)
+        if runtime.triton:
+            environment_variables["CONTAINER_TYPE"] = "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_TRITON = "triton"
 
     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_TRITON: "triton"
     }
 
     payload_attribute_map = {
@@ -532,3 +534,29 @@ def with_health_check_port(
             The ModelDeploymentContainerRuntime instance (self).
         """
         return self.set_spec(self.CONST_HEALTH_CHECK_PORT, health_check_port)
+    
+    @property
+    def triton(self) -> str:
+        """Whether container is triton or not.
+
+        Returns
+        -------
+        bool
+            Whether container is triton or not.
+        """
+        return self.get_spec(self.CONST_TRITON, False)
+
+    def with_triton(self, triton: bool = True) -> "ModelDeploymentRuntime":
+        """Sets the flag for triton.
+
+        Parameters
+        ----------
+        triton: bool
+            Whether it is a triton container.
+
+        Returns
+        -------
+        ModelDeploymentRuntime
+            The ModelDeploymentRuntime instance (self).
+        """
+        return self.set_spec(self.CONST_TRITON, triton)

ads/model/generic_model.py

@@ -2054,6 +2054,7 @@ def deploy(
             raise ValueError("`compartment_id` has to be provided.")
         if not (self.properties.project_id or existing_infrastructure.project_id):
             raise ValueError("`project_id` has to be provided.")
+        
         infrastructure = (
             ModelDeploymentInfrastructure()
             .with_compartment_id(