deploy.py

# =========== Copyright 2023 @ CAMEL-AI.org. All Rights Reserved. ===========
# Licensed under the Apache License, Version 2.0 (the “License”);
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an “AS IS” BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =========== Copyright 2023 @ CAMEL-AI.org. All Rights Reserved. ===========
import subprocess
import threading
import time
import argparse
import requests

model_paths = {
    'llama-3': 'meta-llama/Meta-Llama-3-8B-Instruct',
    'qwen-2.5': 'qwen-2_5-7b-instruct',
}

allocation_size = {
    'llama-3':  3,
    'qwen-2.5':  3
}


def check_port_open(host, port):
    while True:
        url = f'http://{host}:{port}/health'
        try:
            response = requests.get(url, timeout=5)
            if response.status_code == 200:
                break
            else:
                time.sleep(0.3)
        except Exception:
            time.sleep(0.3)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Deploy models with specified options.')
    parser.add_argument('--model_path', type=str, default=None, help='Path to the model')
    parser.add_argument('--model_name', type=str, default='qwen-2.5', help='Name of the served model')
    args = parser.parse_args()

    if args.model_path is None:
        if args.model_name in model_paths:
            args.model_path = model_paths[args.model_name]
        else:
            raise ValueError(f"Model name '{args.model_name}' is not recognized. Please provide a valid model_path.")
        
    host = "127.0.0.1"
    ports = [
        [8020, 8018, 8019],
        [8024, 8025, 8026],
        [8006, 8007, 8008],
        [8002, 8003, 8005],
        [8011, 8009, 8010],
        [8014, 8012, 8013],
        [8017, 8015, 8016],
        [8021, 8022, 8023]
    ]
    gpus = [0]

    all_ports = [port for i in gpus for port in ports[i]]
    print("All ports: ", all_ports, '\n\n')
    t = None
    gpu_utilization = 0.9 / allocation_size[args.model_name]
    for i in range(allocation_size[args.model_name]):
        for j, gpu in enumerate(gpus):
            cmd = (
                f"CUDA_VISIBLE_DEVICES={gpu} python -m "
                f"vllm.entrypoints.openai.api_server --model"
                f" '{args.model_path}' "
                f"--served-model-name '{args.model_name}' "
                f"--host {host} --port {ports[j][i]} --gpu-memory-utilization "
                f"{gpu_utilization} --disable-log-stats --enforce-eager")
            t = threading.Thread(target=subprocess.run,
                                 args=(cmd, ),
                                 kwargs={"shell": True},
                                 daemon=True)
            t.start()
        check_port_open(host, ports[0][i])
    t.join()