Add protocol argument to run_pyfunc_model & start_server

python/pyfunc-server/examples/echo_upi/upi_server.py

@@ -5,6 +5,7 @@
 import merlin
 from caraml.upi.v1 import upi_pb2
 from merlin.model import PyFuncModel
+from merlin.protocol import Protocol
 from prometheus_client import Counter, Gauge
 
 
@@ -46,7 +47,5 @@ def upiv1_infer(
     merlin.run_pyfunc_model(
         model_instance=EchoUPIModel(model_name, model_version),
         conda_env="env.yaml",
-        env_vars={
-            "CARAML_PROTOCOL": "UPI_V1",
-        },
+        protocol=Protocol.UPI_V1,
     )

python/sdk/merlin/fluent.py

@@ -362,18 +362,18 @@ def deploy(
     """
     Deploy a model version.
 
-        :param model_version: If model_version is not given it will deploy active model version
-        :param environment_name: target environment to which the model version will be deployed to. If left empty it will deploy to default environment.
-        :param resource_request: The resource requirement and replicas requests for model version endpoint.
-        :param image_builder_resource_request: The resource requirement and replicas requests for image builder job.
-        :param env_vars: List of environment variables to be passed to the model container.
-        :param transformer: The service to be deployed alongside the model for pre/post-processing steps.
-        :param logger: Response/Request logging configuration for model or transformer.
-        :param deployment_mode: mode of deployment for the endpoint (default: DeploymentMode.SERVERLESS)
-        :param autoscaling_policy: autoscaling policy to be used for the deployment (default: None)
-        :param protocol: protocol to be used by the deployed model (default: HTTP_JSON)
-        :param enable_model_observability: flag to determine whether model observability enabled for the endpoint
-        :return: VersionEndpoint object
+    :param model_version: If model_version is not given it will deploy active model version
+    :param environment_name: target environment to which the model version will be deployed to. If left empty it will deploy to default environment.
+    :param resource_request: The resource requirement and replicas requests for model version endpoint.
+    :param image_builder_resource_request: The resource requirement and replicas requests for image builder job.
+    :param env_vars: List of environment variables to be passed to the model container.
+    :param transformer: The service to be deployed alongside the model for pre/post-processing steps.
+    :param logger: Response/Request logging configuration for model or transformer.
+    :param deployment_mode: mode of deployment for the endpoint (default: DeploymentMode.SERVERLESS)
+    :param autoscaling_policy: autoscaling policy to be used for the deployment (default: None)
+    :param protocol: protocol to be used by the deployed model (default: HTTP_JSON)
+    :param enable_model_observability: flag to determine whether model observability enabled for the endpoint
+    :return: VersionEndpoint object
     """
     _check_active_client()
     if model_version is None:

python/sdk/merlin/model.py

@@ -1412,6 +1412,7 @@ def start_server(
         self,
         env_vars: Dict[str, str] = None,
         port: int = 8080,
+        protocol: Protocol = Protocol.HTTP_JSON,
         pyfunc_base_image: str = None,
         kill_existing_server: bool = False,
         tmp_dir: Optional[str] = os.environ.get("MERLIN_TMP_DIR"),
@@ -1423,6 +1424,7 @@ def start_server(
 
         :param env_vars: dictionary of environment variables to be passed to the server
         :param port: host port that will be used to expose model server
+        :param protocol: protocol to be used by the deployed model (default: HTTP_JSON)
         :param pyfunc_base_image: (optional, default=None) docker image to be used as base image for building pyfunc model
         :param kill_existing_server: (optional, default=False) kill existing server if has been started previously
         :param tmp_dir: (optional, default=None) specify base path for storing model artifact
@@ -1471,6 +1473,7 @@ def start_server(
                 pyfunc_base_image=pyfunc_base_image,
                 port=port,
                 env_vars=env_vars,
+                protocol=protocol,
                 debug=debug,
             )
             return

python/sdk/merlin/pyfunc.py

@@ -403,6 +403,7 @@ def run_pyfunc_model(
     pyfunc_base_image: str = None,
     port: int = 8080,
     env_vars: Dict[str, str] = None,
+    protocol: Protocol = Protocol.HTTP_JSON,
     debug: bool = False,
 ):
     """
@@ -415,6 +416,7 @@ def run_pyfunc_model(
     :param pyfunc_base_image: base image for building pyfunc model
     :param port: port to expose the model
     :param env_vars: dictionary of environment variables to be passed to the server
+    :param protocol: protocol to be used by the deployed model (default: HTTP_JSON)
     :param debug: flag to enable debug mode that will print docker build log
     """
 
@@ -452,6 +454,7 @@ def run_pyfunc_model(
         pyfunc_base_image=pyfunc_base_image,
         port=port,
         env_vars=env_vars,
+        protocol=protocol,
         debug=debug,
     )
 
@@ -465,6 +468,7 @@ def run_pyfunc_local_server(
     pyfunc_base_image: str = None,
     port: int = 8080,
     env_vars: Dict[str, str] = None,
+    protocol: Protocol = None,
     debug: bool = False,
 ):
     if pyfunc_base_image is None:
@@ -506,6 +510,7 @@ def run_pyfunc_local_server(
         model_full_name=f"{model_name}-{model_version}",
         port=port,
         env_vars=env_vars,
+        protocol=protocol,
     )
 
 
@@ -552,6 +557,7 @@ def _run_container(
     model_full_name,
     port,
     env_vars: Dict[str, str] = None,
+    protocol: Protocol = None,
 ):
     docker_client = docker.from_env()
 
@@ -573,7 +579,7 @@ def _run_container(
     env_vars["WORKERS"] = "1"
 
     ports = {"8080/tcp": port}
-    if "CARAML_PROTOCOL" in env_vars and env_vars["CARAML_PROTOCOL"] == "UPI_V1":
+    if protocol == Protocol.UPI_V1:
         ports = {"9000/tcp": port}
 
     try: