diff --git a/python/pyfunc-server/examples/echo_upi/upi_server.py b/python/pyfunc-server/examples/echo_upi/upi_server.py index 3dd1fccf0..57b320fd5 100644 --- a/python/pyfunc-server/examples/echo_upi/upi_server.py +++ b/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, ) diff --git a/python/sdk/merlin/fluent.py b/python/sdk/merlin/fluent.py index c5a96e1fd..004ae88d8 100644 --- a/python/sdk/merlin/fluent.py +++ b/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: diff --git a/python/sdk/merlin/model.py b/python/sdk/merlin/model.py index f535b42ca..357440d41 100644 --- a/python/sdk/merlin/model.py +++ b/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 diff --git a/python/sdk/merlin/pyfunc.py b/python/sdk/merlin/pyfunc.py index 268677c86..2fdfc97ef 100644 --- a/python/sdk/merlin/pyfunc.py +++ b/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: