web-demo.py

# Copyright 2025 ByteDance Ltd. and/or its affiliates
#
# 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.

import os
import subprocess
import sys
import signal
import tinychat.utils.exp_configs as exp_configs
import argparse
import time

parser = argparse.ArgumentParser()
parser.add_argument("--model", type=str, default="llama3.3")
parser.add_argument("--tp", type=int, default=8)
parser.add_argument("--target_tp", type=int, default=4)
parser.add_argument("--draft_tp", type=int, default=4)
args = parser.parse_args()

# Global variable to store the subprocess
process = None


def signal_handler(sig, frame):
    """Handle termination signals and cleanup all processes"""
    print(f"\nReceived signal {sig}. Cleaning up processes...")

    if process:
        # Kill the entire process group to ensure all spawned processes are terminated
        try:
            # Send SIGTERM to the process group
            os.killpg(os.getpgid(process.pid), signal.SIGTERM)
            print("Sent SIGTERM to process group")

            # Wait a bit for graceful shutdown
            time.sleep(2)

            # If still running, force kill
            if process.poll() is None:
                os.killpg(os.getpgid(process.pid), signal.SIGKILL)
                print("Force killed process group with SIGKILL")

        except (ProcessLookupError, OSError) as e:
            print(f"Process cleanup completed or already terminated: {e}")

    # Additional cleanup: find and kill any remaining Python processes on the target GPUs
    gpu_devices = os.environ.get("CUDA_VISIBLE_DEVICES", "").split(",")
    if gpu_devices and gpu_devices[0]:  # If CUDA_VISIBLE_DEVICES is set
        try:
            # Use nvidia-smi to find processes on our GPUs and kill them
            cmd = [
                "nvidia-smi",
                "--query-compute-apps=pid,gpu_name",
                "--format=csv,noheader,nounits",
            ]
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=5)

            for line in result.stdout.strip().split("\n"):
                if line.strip():
                    parts = line.split(",")
                    if len(parts) >= 2:
                        pid = parts[0].strip()
                        try:
                            # Check if this is a python process
                            check_cmd = ["ps", "-p", pid, "-o", "comm="]
                            check_result = subprocess.run(
                                check_cmd, capture_output=True, text=True, timeout=2
                            )
                            if "python" in check_result.stdout.lower():
                                os.kill(int(pid), signal.SIGKILL)
                                print(f"Killed remaining Python process {pid}")
                        except (
                            ValueError,
                            ProcessLookupError,
                            subprocess.TimeoutExpired,
                        ):
                            pass
        except (subprocess.TimeoutExpired, FileNotFoundError):
            pass

    print("Cleanup completed. Exiting...")
    sys.exit(0)


# prepare AWQ
# exp_configs.init()

if args.model == "llama3.3":
    exp_config = exp_configs.llama_3_3_exp_config
elif args.model == "llama3":
    exp_config = exp_configs.llama_exp_config
elif args.model == "qwen":
    exp_config = exp_configs.qwen_exp_config
elif args.model == "r1llama":
    exp_config = exp_configs.r1_llama_exp_config
elif args.model == "r1qwen":
    exp_config = exp_configs.r1_qwen_exp_config
elif args.model == "deepseek":
    exp_config = exp_configs.deepseek_coder_exp_config
elif args.model == "llama-small-test":
    exp_config = exp_configs.llama_small_test_exp_config
else:
    raise ValueError(f"Model {args.model} not supported")


def main():
    global process

    # Register signal handlers for proper cleanup
    signal.signal(signal.SIGINT, signal_handler)  # Ctrl+C
    signal.signal(signal.SIGTERM, signal_handler)  # Termination signal
    if hasattr(signal, "SIGHUP"):  # Not available on Windows
        signal.signal(signal.SIGHUP, signal_handler)

    # Model configurations
    # MODEL = exp_config.target_model_config.model_path
    # TARGET_FILE_NAME_PREFIX = exp_config.target_model_config.ckpt_prefix + "/" + exp_config.target_model_config.model_id
    EAMODEL = exp_config.draft_model_config.model_path
    DRAFT_FILE_NAME_PREFIX = (
        exp_config.draft_model_config.ckpt_prefix
        + "/"
        + exp_config.draft_model_config.model_id
    )
    MODEL = EAMODEL
    TARGET_FILE_NAME_PREFIX = DRAFT_FILE_NAME_PREFIX
    # CUDA and device configurations
    os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(str(i) for i in range(args.tp))
    os.environ["NUM_DEVICES"] = str(args.tp)
    os.environ["SPEC_TYPE"] = "pipe"
    os.environ["DTYPE"] = "bfloat16"
    os.environ["GRADIO_SERVER_NAME"] = "0.0.0.0"

    MODEL_TYPE = "llama-3-instruct"
    NUM_DEVICES = args.tp

    # Construct the torchrun command
    cmd = [
        "torchrun",
        f"--nproc-per-node={NUM_DEVICES}",
        "../tinychat/webui.py",
        "--ea-model-path",
        EAMODEL,
        "--base-model-path",
        MODEL,
        "--model-type",
        MODEL_TYPE,
        "--load-ckpt-file-name",
        TARGET_FILE_NAME_PREFIX,
        "--draft-load-ckpt-file-name",
        DRAFT_FILE_NAME_PREFIX,
        "--spec-type",
        "pipe",
        "--dtype",
        "bfloat16",
        "--tp",
        str(args.tp),
        "--draft-tp",
        str(args.draft_tp),
        "--target-tp",
        str(args.target_tp),
    ]

    try:
        # Run the command with process group creation for proper cleanup
        print(f"Starting torchrun with command: {' '.join(cmd)}")
        print(f"Using GPUs: {os.environ['CUDA_VISIBLE_DEVICES']}")

        process = subprocess.Popen(
            cmd,
            preexec_fn=os.setsid,  # Create new process group
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            universal_newlines=True,
            bufsize=1,
        )

        # Stream output in real-time
        try:
            for line in iter(process.stdout.readline, ""):
                print(line.rstrip())

            # Wait for process to complete
            returncode = process.wait()
            print(f"Command completed with return code: {returncode}")

        except KeyboardInterrupt:
            # This should be handled by signal_handler, but just in case
            signal_handler(signal.SIGINT, None)

    except subprocess.CalledProcessError as e:
        print(f"Command failed with return code: {e.returncode}")
        sys.exit(e.returncode)
    except Exception as e:
        print(f"Error running command: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()