Multi model deployment

examples/multi_model/query.py

@@ -0,0 +1,42 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import mii
+import time
+
+generator = mii.mii_query_handle("first_test")
+result = generator.query(
+    {"query": ["DeepSpeed is",
+               "Seattle is"]},
+    "bloom560m_deployment",
+    do_sample=True,
+    max_new_tokens=30,
+)
+print(result)
+
+time.sleep(5)
+result = generator.query({'query': "DeepSpeed is the greatest"},
+                         "microsoft/DialogRPT-human-vs-rand_deployment")
+print(result)
+
+time.sleep(5)
+
+result = generator.query(
+    {
+        'text': "DeepSpeed is the greatest",
+        'conversation_id': 3,
+        'past_user_inputs': [],
+        'generated_responses': []
+    },
+    "microsoft/DialoGPT-large_deployment")
+print(result)
+
+results = generator.query(
+    {
+        'question': "What is the greatest?",
+        'context': "DeepSpeed is the greatest"
+    },
+    "deepset/roberta-large-squad2" + "-qa-deployment")
+print(results)

examples/multi_model/shutdown.py

@@ -0,0 +1,7 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+import mii
+
+mii.terminate("first_test")

examples/multi_model/text-generation-bloom560m-example.py

@@ -0,0 +1,46 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+import mii
+
+gpu_index_map1 = {'master': [0]}
+gpu_index_map2 = {'master': [1]}
+gpu_index_map3 = {'master': [0, 1]}
+
+deployments = []
+mii_configs1 = {"tensor_parallel": 2, "dtype": "fp16"}
+deployments.append(
+    mii.Deployment(task='text-generation',
+                   model="bigscience/bloom-560m",
+                   deployment_name="bloom560m_deployment",
+                   GPU_index_map=gpu_index_map3,
+                   mii_config=mii.config.MIIConfig(**mii_configs1)))
+
+# gpt2
+name = "microsoft/DialogRPT-human-vs-rand"
+deployments.append(
+    mii.Deployment(task='text-classification',
+                   model=name,
+                   deployment_name=name + "_deployment",
+                   GPU_index_map=gpu_index_map2))
+
+mii_configs2 = {"tensor_parallel": 1}
+
+name = "microsoft/DialoGPT-large"
+
+deployments.append(
+    mii.Deployment(task='conversational',
+                   model=name,
+                   deployment_name=name + "_deployment",
+                   GPU_index_map=gpu_index_map1,
+                   mii_config=mii.config.MIIConfig(**mii_configs2)))
+
+name = "deepset/roberta-large-squad2"
+deployments.append(
+    mii.Deployment(task="question-answering",
+                   model=name,
+                   deployment_name=name + "-qa-deployment",
+                   GPU_index_map=gpu_index_map2))
+
+mii.deploy(deployment_tag="first_test", deployments=deployments)

mii/__init__.py

@@ -10,7 +10,7 @@
 from .constants import DeploymentType, Tasks
 from .aml_related.utils import aml_output_path
 
-from .config import MIIConfig, LoadBalancerConfig
+from .config import MIIConfig, LoadBalancerConfig, Deployment
 from .grpc_related.proto import modelresponse_pb2_grpc
 
 __version__ = "0.0.0"

mii/client.py

@@ -12,17 +12,27 @@
 from mii.method_table import GRPC_METHOD_TABLE
 
 
-def _get_deployment_info(deployment_name):
-    configs = mii.utils.import_score_file(deployment_name).configs
-    task = configs[mii.constants.TASK_NAME_KEY]
-    mii_configs_dict = configs[mii.constants.MII_CONFIGS_KEY]
+def _get_deployment_info(deployment_tag):
+    deployments = []
+    configs = mii.utils.import_score_file(deployment_tag).configs
+    for deployment in configs:
+        if not isinstance(configs[deployment], dict):
+            continue
+        deployments.append(configs[deployment])
+        mii_configs_dict = configs[deployment][mii.constants.MII_CONFIGS_KEY]
+        mii_configs = mii.config.MIIConfig(**mii_configs_dict)
+    return deployments
+    """
+    task = configs[deployment_name][mii.constants.TASK_NAME_KEY]
+    mii_configs_dict = configs[deployment_name][mii.constants.MII_CONFIGS_KEY]
     mii_configs = mii.config.MIIConfig(**mii_configs_dict)
 
     assert task is not None, "The task name should be set before calling init"
     return task, mii_configs
+    """
 
 
-def mii_query_handle(deployment_name):
+def mii_query_handle(deployment_tag):
     """Get a query handle for a local deployment:
 
         mii/examples/local/gpt2-query-example.py
@@ -35,12 +45,14 @@ def mii_query_handle(deployment_name):
         query_handle: A query handle with a single method `.query(request_dictionary)` using which queries can be sent to the model.
     """
 
-    if deployment_name in mii.non_persistent_models:
-        inference_pipeline, task = mii.non_persistent_models[deployment_name]
-        return MIINonPersistentClient(task, deployment_name)
+    if deployment_tag in mii.non_persistent_models:
+        inference_pipeline, task = mii.non_persistent_models[deployment_tag]
+        return MIINonPersistentClient(task, deployment_tag)
 
-    task_name, mii_configs = _get_deployment_info(deployment_name)
-    return MIIClient(task_name, "localhost", mii_configs.port_number)
+    deployments = _get_deployment_info(deployment_tag)
+    mii_configs_dict = deployments[0][mii.constants.MII_CONFIGS_KEY]
+    mii_configs = mii.config.MIIConfig(**mii_configs_dict)
+    return MIIClient(deployments, "localhost", mii_configs.port_number)
 
 
 def create_channel(host, port):
@@ -55,24 +67,37 @@ class MIIClient():
     """
     Client to send queries to a single endpoint.
     """
-    def __init__(self, task_name, host, port):
+    def __init__(self, deployments, host, port):
         self.asyncio_loop = asyncio.get_event_loop()
         channel = create_channel(host, port)
         self.stub = modelresponse_pb2_grpc.ModelResponseStub(channel)
-        self.task = get_task(task_name)
+        #self.task = get_task(task_name)
+        self.deployments = deployments
 
-    async def _request_async_response(self, request_dict, **query_kwargs):
-        if self.task not in GRPC_METHOD_TABLE:
-            raise ValueError(f"unknown task: {self.task}")
+    async def _request_async_response(self, request_dict, task, **query_kwargs):
+        if task not in GRPC_METHOD_TABLE:
+            raise ValueError(f"unknown task: {task}")
 
-        task_methods = GRPC_METHOD_TABLE[self.task]
+        task_methods = GRPC_METHOD_TABLE[task]
         proto_request = task_methods.pack_request_to_proto(request_dict, **query_kwargs)
         proto_response = await getattr(self.stub, task_methods.method)(proto_request)
         return task_methods.unpack_response_from_proto(proto_response)
 
-    def query(self, request_dict, **query_kwargs):
+    def query(self, request_dict, deployment_name=None, **query_kwargs):
+        task = None
+        if deployment_name is None:  #mii.terminate() or single model
+            #assert len(self.deployments) == 1, "Must pass deployment_name to query when using multiple deployments"
+            deployment_name = self.deployments[0]['deployment_name']
+            task = get_task(self.deployments[0]['task_name'])
+        else:
+            for deployment in self.deployments:
+                if deployment[mii.constants.DEPLOYMENT_NAME_KEY] == deployment_name:
+                    task = get_task(deployment[mii.constants.TASK_NAME_KEY])
+                    break
+        query_kwargs['deployment_name'] = deployment_name
         return self.asyncio_loop.run_until_complete(
             self._request_async_response(request_dict,
+                                         task,
                                          **query_kwargs))
 
     async def terminate_async(self):
@@ -86,17 +111,35 @@ async def create_session_async(self, session_id):
         return await self.stub.CreateSession(
             modelresponse_pb2.SessionID(session_id=session_id))
 
-    def create_session(self, session_id):
-        assert self.task == Tasks.TEXT_GENERATION, f"Session creation only available for task '{Tasks.TEXT_GENERATION}'."
+    def create_session(self, session_id, deployment_name=None):
+        task = None
+        if deployment_name is None:  #mii.terminate() or single model
+            deployment_name = self.deployments[0][mii.constants.DEPLOYMENT_NAME_KEY]
+            task = get_task(self.deployments[0][mii.constants.TASK_NAME_KEY])
+        else:
+            for deployment in self.deployments:
+                if deployment[mii.constants.DEPLOYMENT_NAME_KEY] == deployment_name:
+                    task = get_task(deployment[mii.constants.TASK_NAME_KEY])
+                    break
+        assert task == Tasks.TEXT_GENERATION, f"Session creation only available for task '{Tasks.TEXT_GENERATION}'."
         return self.asyncio_loop.run_until_complete(
             self.create_session_async(session_id))
 
     async def destroy_session_async(self, session_id):
         await self.stub.DestroySession(modelresponse_pb2.SessionID(session_id=session_id)
                                        )
 
-    def destroy_session(self, session_id):
-        assert self.task == Tasks.TEXT_GENERATION, f"Session deletion only available for task '{Tasks.TEXT_GENERATION}'."
+    def destroy_session(self, session_id, deployment_name=None):
+        task = None
+        if deployment_name is None:  #mii.terminate() or single model
+            deployment_name = self.deployments[0][mii.constants.DEPLOYMENT_NAME_KEY]
+            task = get_task(self.deployments[0][mii.constants.TASK_NAME_KEY])
+        else:
+            for deployment in self.deployments:
+                if deployment[mii.constants.DEPLOYMENT_NAME_KEY] == deployment_name:
+                    task = get_task(deployment[mii.constants.TASK_NAME_KEY])
+                    break
+        assert task == Tasks.TEXT_GENERATION, f"Session deletion only available for task '{Tasks.TEXT_GENERATION}'."
         self.asyncio_loop.run_until_complete(self.destroy_session_async(session_id))
 
 
@@ -157,7 +200,7 @@ def destroy_session(self, session_id):
 
 class MIINonPersistentClient():
     def __init__(self, task, deployment_name):
-        self.task = task
+        self.task = get_task(task)
         self.deployment_name = deployment_name
 
     def query(self, request_dict, **query_kwargs):
@@ -188,7 +231,10 @@ def terminate(self):
         del mii.non_persistent_models[self.deployment_name]
 
 
-def terminate_restful_gateway(deployment_name):
-    _, mii_configs = _get_deployment_info(deployment_name)
-    if mii_configs.enable_restful_api:
-        requests.get(f"http://localhost:{mii_configs.restful_api_port}/terminate")
+def terminate_restful_gateway(deployment_tag):
+    deployments = _get_deployment_info(deployment_tag)
+    for deployment in deployments:
+        mii_configs_dict = deployment[mii.constants.MII_CONFIGS_KEY]
+        mii_configs = mii.config.MIIConfig(**mii_configs_dict)
+        if mii_configs.enable_restful_api:
+            requests.get(f"http://localhost:{mii_configs.restful_api_port}/terminate")

mii/config.py

@@ -6,7 +6,6 @@
 from typing import Union, List
 from enum import Enum
 from pydantic import BaseModel, validator, root_validator
-
 from deepspeed.launcher.runner import DLTS_HOSTFILE
 
 
@@ -107,6 +106,8 @@ class Config:
 
 
 class ReplicaConfig(BaseModel):
+    task: str = ""
+    deployment_name: str = ""
     hostname: str = ""
     tensor_parallel_ports: List[int] = []
     torch_dist_port: int = None
@@ -123,4 +124,18 @@ class LoadBalancerConfig(BaseModel):
 
     class Config:
         validate_all = True
-        validate_assignment = True
+
+
+validate_assignment = True
+
+
+class Deployment(BaseModel):
+    deployment_name: str
+    task: str
+    model: str
+    enable_deepspeed: bool = True
+    enable_zero: bool = False
+    GPU_index_map: dict = None
+    mii_config: MIIConfig = MIIConfig.parse_obj({})
+    ds_config: dict = None
+    version: int = 1

mii/constants.py

@@ -94,7 +94,7 @@ class ModelProvider(enum.Enum):
 DEPLOYMENT_NAME_KEY = 'deployment_name'
 MODEL_PATH_KEY = 'model_path'
 LOAD_BALANCER_CONFIG_KEY = 'load_balancer_config'
-
+DEPLOYMENT_TAG_KEY = 'deployment_tag'
 ENABLE_DEEPSPEED_KEY = 'ds_optimize'
 ENABLE_DEEPSPEED_ZERO_KEY = 'ds_zero'
 DEEPSPEED_CONFIG_KEY = 'ds_config'