Multi model deployment

mii/__init__.py

@@ -15,6 +15,7 @@
 
 __version__ = "0.0.0"
 non_persistent_models = {}
+multi_model_deployments = {}
 try:
     from .version import __version__
 except ImportError:

mii/config.py

@@ -107,6 +107,7 @@ class Config:
 
 
 class ReplicaConfig(BaseModel):
+    deployment_name: str = ""
     hostname: str = ""
     tensor_parallel_ports: List[int] = []
     torch_dist_port: int = None
@@ -123,4 +124,16 @@ 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 = True
+    GPU_index_map: dict = None
+    mii_config: dict = None
+    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'

mii/deployment.py

@@ -16,16 +16,17 @@
 from .config import ReplicaConfig, LoadBalancerConfig
 
 
-def deploy(task,
-           model,
-           deployment_name,
-           deployment_type=DeploymentType.LOCAL,
-           model_path=None,
+def deploy(task=None,
+           model=None,
+           deployment_name=None,
            enable_deepspeed=True,
            enable_zero=False,
            ds_config=None,
            mii_config={},
-           version=1):
+           deployment_tag=None,
+           deployments=[],
+           deployment_type=DeploymentType.LOCAL,
+           model_path=None):
     """Deploy a task using specified model. For usage examples see:
 
         mii/examples/local/text-generation-example.py
@@ -66,37 +67,49 @@ def deploy(task,
         If deployment_type is `LOCAL`, returns just the name of the deployment that can be used to create a query handle using `mii.mii_query_handle(deployment_name)`
 
     """
+    if len(deployments == 0):
+        assert model is not None and task is not None and deployment_name is not None, "model, task, and deployment name must be set to deploy sigular model"
+        deployments = [Deployment(deployment_name, task, model, enable_deepspeed, enable_zero, None, mii_config, ds_config, version)]
+        deployment_tag = deployment_name + "_tag"
+    else:
+        assert deployment_tag is not None, "deployment_tag must be set to deploy multiple models"
 
+    mii.multi_model_deployments[deployment_tag] = deployments
+    ports = set()
     # parse and validate mii config
-    mii_config = mii.config.MIIConfig(**mii_config)
-    if enable_zero:
-        if ds_config.get("fp16", {}).get("enabled", False):
-            assert (mii_config.dtype == torch.half), "MII Config Error: MII dtype and ZeRO dtype must match"
-        else:
-            assert (mii_config.dtype == torch.float), "MII Config Error: MII dtype and ZeRO dtype must match"
-    assert not (enable_deepspeed and enable_zero), "MII Config Error: DeepSpeed and ZeRO cannot both be enabled, select only one"
+    for deployment in deployments:
+        mii_config = mii.config.MIIConfig(**deployment.mii_config)
+        assert mii_config.port_number not in ports, f"duplicate port numbers not allowed - {mii.config.port_number}"
+        ports.add(mii_config.port_number)
+        if deployment.enable_zero:
+            if deployment.ds_config.get("fp16", {}).get("enabled", False):
+                assert (mii_config.dtype == torch.half), "MII Config Error: MII dtype and ZeRO dtype must match"
+            else:
+                assert (mii_config.dtype == torch.float), "MII Config Error: MII dtype and ZeRO dtype must match"
+        assert not (enable_deepspeed and enable_zero), "MII Config Error: DeepSpeed and ZeRO cannot both be enabled, select only one"
 
     # aml only allows certain characters for deployment names
     if deployment_type == DeploymentType.AML:
         allowed_chars = set(string.ascii_lowercase + string.ascii_uppercase +
                             string.digits + '-')
         assert set(deployment_name) <= allowed_chars, "AML deployment names can only contain a-z, A-Z, 0-9, and '-'"
 
-    task = mii.utils.get_task(task)
+    for deployment in deployments:
+        deployment.task = mii.utils.get_task(deployment.task)
 
-    if not mii_config.skip_model_check:
-        mii.utils.check_if_task_and_model_is_valid(task, model)
-        if enable_deepspeed:
-            mii.utils.check_if_task_and_model_is_supported(task, model)
+        if not mii_config.skip_model_check:
+            mii.utils.check_if_task_and_model_is_valid(deployment.task, deployment.model)
+            if enable_deepspeed:
+                mii.utils.check_if_task_and_model_is_supported(deployment.task, deployment.model)
 
-    if enable_deepspeed:
-        logger.info(
-            f"************* MII is using DeepSpeed Optimizations to accelerate your model *************"
-        )
-    else:
-        logger.info(
-            f"************* DeepSpeed Optimizations not enabled. Please use enable_deepspeed to get better performance *************"
-        )
+        if enable_deepspeed:
+            logger.info(
+                    f"************* MII is using DeepSpeed Optimizations to accelerate your model: {deployment.model} *************"
+            )
+        else:
+            logger.info(
+                    f"************* DeepSpeed Optimizations not enabled. Please use enable_deepspeed to get better performance for: {deployment.model} *************"
+            )
 
     # In local deployments use default path if no model path set
     if model_path is None and deployment_type == DeploymentType.LOCAL:
@@ -126,21 +139,16 @@ def deploy(task,
                                    replica_configs=replica_configs)
 
     if deployment_type != DeploymentType.NON_PERSISTENT:
-        create_score_file(deployment_name=deployment_name,
+        create_score_file(deployment_tag=deployment_tag,
+                          deployments=deployments,
                           deployment_type=deployment_type,
-                          task=task,
-                          model_name=model,
-                          ds_optimize=enable_deepspeed,
-                          ds_zero=enable_zero,
-                          ds_config=ds_config,
-                          mii_config=mii_config,
                           model_path=model_path,
                           lb_config=lb_config)
 
     if deployment_type == DeploymentType.AML:
-        _deploy_aml(deployment_name=deployment_name, model_name=model, version=version)
+        _deploy_aml(deployment_tag=deployment_tag, model_name=model, version=version)
     elif deployment_type == DeploymentType.LOCAL:
-        return _deploy_local(deployment_name, model_path=model_path)
+        return _deploy_local(deployment_tag, model_path=model_path)
     elif deployment_type == DeploymentType.NON_PERSISTENT:
         assert int(os.getenv('WORLD_SIZE', '1')) == mii_config.tensor_parallel, "World Size does not equal number of tensors. When using non-persistent deployment type, please launch with `deepspeed --num_gpus <tensor_parallel>`"
         provider = MODEL_PROVIDER_MAP[get_provider_name(model, task)]
@@ -157,14 +165,14 @@ def deploy(task,
         raise Exception(f"Unknown deployment type: {deployment_type}")
 
 
-def _deploy_local(deployment_name, model_path):
-    mii.utils.import_score_file(deployment_name).init()
+def _deploy_local(deployment_tag, model_path):
+    mii.utils.import_score_file(deployment_tag).init()
 
 
-def _deploy_aml(deployment_name, model_name, version):
+def _deploy_aml(deployment_tag, model_name, version):
     acr_name = mii.aml_related.utils.get_acr_name()
     mii.aml_related.utils.generate_aml_scripts(acr_name=acr_name,
-                                               deployment_name=deployment_name,
+                                               deployment_name=deployment_tag,
                                                model_name=model_name,
                                                version=version)
     print(

mii/models/score/generate.py

@@ -9,53 +9,52 @@
 from mii.constants import DeploymentType
 
 
-def create_score_file(deployment_name,
+def create_score_file(deployment_tag,
                       deployment_type,
-                      task,
-                      model_name,
-                      ds_optimize,
-                      ds_zero,
-                      ds_config,
-                      mii_config,
+                      deployments,
                       model_path,
                       lb_config):
+
     config_dict = {}
-    config_dict[mii.constants.DEPLOYMENT_NAME_KEY] = deployment_name
-    config_dict[mii.constants.TASK_NAME_KEY] = mii.utils.get_task_name(task)
-    config_dict[mii.constants.MODEL_NAME_KEY] = model_name
-    config_dict[mii.constants.ENABLE_DEEPSPEED_KEY] = ds_optimize
-    config_dict[mii.constants.MII_CONFIGS_KEY] = mii_config.dict()
-    config_dict[mii.constants.ENABLE_DEEPSPEED_ZERO_KEY] = ds_zero
-    config_dict[mii.constants.DEEPSPEED_CONFIG_KEY] = ds_config
     config_dict[mii.constants.MODEL_PATH_KEY] = model_path
-
-    if lb_config is not None:
-        config_dict[mii.constants.LOAD_BALANCER_CONFIG_KEY] = lb_config
-
-    if len(mii.__path__) > 1:
-        logger.warning(
-            f"Detected mii path as multiple sources: {mii.__path__}, might cause unknown behavior"
-        )
+    config_dict[mii.constants.DEPLOYMENT_TAG_KEY] = deployment_tag
+    for deployment in deployments:
+        config_dict[deployment.deployment_name] = {}
+        config_dict[deployment.deployment_name][mii.constants.DEPLOYMENT_NAME_KEY] = deployment_name
+        config_dict[deployment.deployment_name][mii.constants.TASK_NAME_KEY] = mii.utils.get_task_name(task)
+        config_dict[deployment.deployment_name][mii.constants.MODEL_NAME_KEY] = model_name
+        config_dict[deployment.deployment_name][mii.constants.ENABLE_DEEPSPEED_KEY] = ds_optimize
+        config_dict[deployment.deployment_name][mii.constants.MII_CONFIGS_KEY] = mii_config.dict()
+        config_dict[deployment.deployment_name][mii.constants.ENABLE_DEEPSPEED_ZERO_KEY] = ds_zero
+        config_dict[deployment.deployment_name][mii.constants.DEEPSPEED_CONFIG_KEY] = ds_config
+
+        if lb_config is not None:
+            config_dict[deployment.deployment_name][mii.constants.LOAD_BALANCER_CONFIG_KEY] = lb_config
+
+        if len(mii.__path__) > 1:
+            logger.warning(
+                f"Detected mii path as multiple sources: {mii.__path__}, might cause unknown behavior"
+            )
 
     with open(os.path.join(mii.__path__[0],
-                           "models/score/score_template.py"),
-              "r") as fd:
+                        "models/score/score_template.py"),
+            "r") as fd:
         score_src = fd.read()
 
     # update score file w. global config dict
     source_with_config = f"{score_src}\n"
     source_with_config += f"configs = {pprint.pformat(config_dict, indent=4)}"
 
-    with open(generated_score_path(deployment_name, deployment_type), "w") as fd:
+    with open(generated_score_path(deployment_tag, deployment_type), "w") as fd:
         fd.write(source_with_config)
         fd.write("\n")
 
 
-def generated_score_path(deployment_name, deployment_type):
+def generated_score_path(deployment_tag, deployment_type):
     if deployment_type == DeploymentType.LOCAL:
-        score_path = os.path.join(mii.utils.mii_cache_path(), deployment_name)
+        score_path = os.path.join(mii.utils.mii_cache_path(), deployment_tag)
     elif deployment_type == DeploymentType.AML:
-        score_path = os.path.join(mii.aml_related.utils.aml_output_path(deployment_name),
+        score_path = os.path.join(mii.aml_related.utils.aml_output_path(deployment_tag),
                                   "code")
     if not os.path.isdir(score_path):
         os.makedirs(score_path)

mii/models/score/score_template.py

@@ -16,6 +16,8 @@
 
 def init():
     model_path = mii.utils.full_model_path(configs[mii.constants.MODEL_PATH_KEY])
+    deployment_tag = configs[mii.constants.DEPLOYMENT_TAG_KEY]
+    deployments = mii.multi_model_deployments[deployment_tag]
 
     deployment_name = configs[mii.constants.DEPLOYMENT_NAME_KEY]
     model_name = configs[mii.constants.MODEL_NAME_KEY]

mii/server.py

@@ -29,14 +29,9 @@ def config_to_b64_str(config):
 class MIIServer():
     '''Initialize the model, setup the server for the model under model_path'''
     def __init__(self,
-                 deployment_name,
-                 task_name,
-                 model_name,
+                 deployment_tag,
+                 deployments,
                  model_path,
-                 ds_optimize=True,
-                 ds_zero=False,
-                 ds_config=None,
-                 mii_configs={},
                  lb_config=None):
 
         mii_configs = mii.config.MIIConfig(**mii_configs)
@@ -55,13 +50,9 @@ def __init__(self,
                 f.write(f"localhost slots={num_gpu}")
             mii.configs.hostfile = hostfile
 
-        processes = self._initialize_service(deployment_name,
-                                             model_name,
+        processes = self._initialize_service(deployment_tag,
+                                             deployments,
                                              model_path,
-                                             ds_optimize,
-                                             ds_zero,
-                                             ds_config,
-                                             mii_configs,
                                              lb_config)
         self._wait_until_server_is_live(processes, lb_config.replica_configs)
 
@@ -278,13 +269,9 @@ def _launch_deepspeed(self,
                                            ds_launch_str=ds_launch_str)
 
     def _initialize_service(self,
-                            deployment_name,
-                            model_name,
+                            deployment_tag,
+                            deployments,
                             model_path,
-                            ds_optimize,
-                            ds_zero,
-                            ds_config,
-                            mii_configs,
                             lb_config):
 
         processes = []
@@ -295,19 +282,20 @@ def _initialize_service(self,
 
         # Start replica instances
         for i, repl_config in enumerate(lb_config.replica_configs):
+            name = repl_config.deployment_name
             hostfile = tempfile.NamedTemporaryFile(delete=False)
             hostfile.write(
                 f'{repl_config.hostname} slots={max(host_gpus[repl_config.hostname])+1}\n'
                 .encode())
             processes.append(
                 self._launch_deepspeed(
-                    deployment_name,
-                    model_name,
+                    name,
+                    deployments[name].model,
                     model_path,
-                    ds_optimize,
-                    ds_zero,
-                    ds_config,
-                    mii_configs,
+                    deployments[name].enable_deepspeed,
+                    deployments[name].enable_zero,
+                    deployments[name].ds_config,
+                    deployments[name].mii_configs,
                     hostfile.name,
                     repl_config.hostname,
                     repl_config.tensor_parallel_ports[0],