Add Non-persistent deployment type

mii/__init__.py

@@ -14,6 +14,7 @@
 from .grpc_related.proto import modelresponse_pb2_grpc
 
 __version__ = "0.0.0"
+non_persistent_model = {}
 try:
     from .version import __version__
 except ImportError:

mii/client.py

@@ -34,14 +34,21 @@ def mii_query_handle(deployment_name):
     Returns:
         query_handle: A query handle with a single method `.query(request_dictionary)` using which queries can be sent to the model.
     """
-    task_name, mii_configs = _get_deployment_info(deployment_name)
-    if mii_configs.enable_load_balancing:
-        return MIIClient(task_name, "localhost", mii_configs.port_number)
+    if len(mii.non_persistent_model) != 0:
+        assert deployment_name in mii.non_persistent_model, f"Could not find '{deployment_name}'"
+        inference_pipeline, task = mii.non_persistent_model[deployment_name]
+        assert task is not None, "The task name should be set before calling init"
+        return MIINonPersistentClient(task, inference_pipeline, deployment_name)
+
     else:
-        return MIITensorParallelClient(
-            task_name,
-            "localhost",
-            [mii_configs.port_number + i for i in range(mii_configs.tensor_parallel)])
+        task_name, mii_configs = _get_deployment_info(deployment_name)
+        if mii_configs.enable_load_balancing:
+            return MIIClient(task_name, "localhost", mii_configs.port_number)
+        else:
+            return MIITensorParallelClient(
+                task_name,
+                "localhost",
+                [mii_configs.port_number + i for i in range(mii_configs.tensor_parallel)])
 
 
 def create_channel(host, port):
@@ -155,6 +162,25 @@ def destroy_session(self, session_id):
         for client in self.clients:
             client.destroy_session(session_id)
 
+class MIINonPersistentClient():
+    def __init__(self, task, inference_pipeline, deployment_name):
+        self.inference_pipeline = inference_pipeline
+        self.task = task
+        self.deployment_name = deployment_name
+
+    def query(self, request_dict, **query_kwargs):
+        task_methods = GRPC_METHOD_TABLE[self.task]
+        if self.task == Tasks.QUESTION_ANSWERING:
+            return task_methods.run_inference(self.inference_pipeline, request_dict, **query_kwargs)
+
+        query = request_dict['query']
+        return task_methods.run_inference(self.inference_pipeline, query, **query_kwargs)
+
+    def terminate(self):
+        print("Terminating ...")
+        del mii.persistent_model[self.deployment_name]
+
+
 
 def terminate_restful_gateway(deployment_name):
     _, mii_configs = _get_deployment_info(deployment_name)

mii/constants.py

@@ -9,7 +9,7 @@
 class DeploymentType(enum.Enum):
     LOCAL = 1
     AML = 2
-
+    NON_PERSISTENT = 3
 
 MII_CONFIGS_KEY = 'mii_configs'
 

mii/deployment.py

@@ -9,8 +9,8 @@
 
 from deepspeed.launcher.runner import fetch_hostfile
 
-from .constants import DeploymentType, MII_MODEL_PATH_DEFAULT
-from .utils import logger
+from .constants import DeploymentType, MII_MODEL_PATH_DEFAULT, MODEL_PROVIDER_MAP
+from .utils import logger, get_task, get_provider_name
 from .models.score import create_score_file
 from .config import ReplicaConfig, LoadBalancerConfig
 
@@ -140,6 +140,10 @@ def deploy(task,
         _deploy_aml(deployment_name=deployment_name, model_name=model, version=version)
     elif deployment_type == DeploymentType.LOCAL:
         return _deploy_local(deployment_name, model_path=model_path)
+    elif deployment_type == DeploymentType.NON_PERSISTENT:
+        assert not mii_config.enable_load_balancing, "Cannot use Load Balancing with Non persistent deployment"
+        provider = MODEL_PROVIDER_MAP[get_provider_name(model, task)]
+        mii.non_persistent_model[deployment_name] = (load_models(task, model, model_path, enable_deepspeed, enable_zero, provider, mii_config), get_task(task))
     else:
         raise Exception(f"Unknown deployment type: {deployment_type}")
 

mii/server.py

@@ -13,7 +13,7 @@
 from collections import defaultdict
 
 import mii
-from mii.utils import get_num_gpus, logger
+from mii.utils import get_num_gpus, logger, get_provider_name
 from mii.config import ReplicaConfig
 
 
@@ -120,14 +120,16 @@ def _build_server_args(self,
         server_args_str += " --ds-optimize" if ds_optimize else ""
 
         # XXX: fetch model provider based on model name in a more general way
-        if model_name == "gpt-neox":
+        """if model_name == "gpt-neox":
             provider = mii.constants.MODEL_PROVIDER_NAME_EA
         elif ("bigscience/bloom" == model_name) or ("microsoft/bloom" in model_name):
             provider = mii.constants.MODEL_PROVIDER_NAME_HF_LLM
         elif self.task == mii.Tasks.TEXT2IMG:
             provider = mii.constants.MODEL_PROVIDER_NAME_DIFFUSERS
         else:
             provider = mii.constants.MODEL_PROVIDER_NAME_HF
+        """
+        provider = get_provider_name(model_name, self.task)
         server_args_str += f" --provider {provider}"
 
         server_args_str += f" --config {b64_config_str}"

mii/utils.py

@@ -208,6 +208,16 @@ def get_num_gpus(mii_configs):
     ) >= num_gpus, f"Available GPU count: {torch.cuda.device_count()} does not meet the required gpu count: {num_gpus}"
     return num_gpus
 
+def get_provider_name(model_name, task):
+    if model_name == "gpt-neox":
+        provider = mii.constants.MODEL_PROVIDER_NAME_EA
+    elif ("bigscience/bloom" == model_name) or ("microsoft/bloom" in model_name):
+        provider = mii.constants.MODEL_PROVIDER_NAME_HF_LLM
+    elif task == mii.Tasks.TEXT2IMG:
+        provider = mii.constants.MODEL_PROVIDER_NAME_DIFFUSERS
+    else:
+        provider = mii.constants.MODEL_PROVIDER_NAME_HF
+    return provider
 
 log_levels = {
     "debug": logging.DEBUG,