| import atexit |
| import os |
| import signal |
| import time |
|
|
| import torch.multiprocessing as mp |
|
|
| from dataclasses import fields |
| from time import perf_counter |
|
|
| from tqdm.auto import tqdm |
| from transformers import AutoTokenizer |
|
|
| from diffulex.config import Config |
| from diffulex.distributed.parallel_state import get_world_size |
| from diffulex.engine.model_runner import AutoModelRunner |
| from diffulex.engine.request import AutoReq |
| from diffulex.engine.scheduler import AutoScheduler, DataParallelScheduler, SchedulerBase |
| from diffulex.logger import get_logger |
| from diffulex.mixin.async_serving.engine import DiffulexAsyncEngineMixin |
| from diffulex.profiling import TorchProfileSession, record_function |
| from diffulex.sampling_params import SamplingParams |
| from diffulex.utils.output import GenerationOutputs |
|
|
| logger = get_logger(__name__) |
|
|
|
|
| def _set_parent_death_signal(sig: int = signal.SIGTERM) -> None: |
| if os.name != "posix": |
| return |
| try: |
| import ctypes |
|
|
| libc = ctypes.CDLL("libc.so.6", use_errno=True) |
| pr_set_pdeathsig = 1 |
| libc.prctl(pr_set_pdeathsig, sig) |
| except Exception: |
| logger.debug("Failed to set parent-death signal for worker process.", exc_info=True) |
|
|
|
|
| def _run_model_runner_worker(config: Config, rank: int, event) -> None: |
| _set_parent_death_signal() |
| if os.getppid() == 1: |
| raise SystemExit("Diffulex worker parent exited before worker initialization.") |
| AutoModelRunner.from_config(config, rank, event) |
|
|
|
|
| class DiffulexEngine(DiffulexAsyncEngineMixin): |
| def __init__(self, model, **kwargs): |
| config_fields = {field.name for field in fields(Config)} |
| config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields} |
| self.config = config = Config(model, **config_kwargs) |
| self.model_parallel_world_size = get_world_size( |
| config.tensor_parallel_size, |
| config.expert_parallel_size, |
| dp_size=config.data_parallel_size, |
| ) |
| if len(config.device_ids) < self.model_parallel_world_size: |
| raise ValueError( |
| "Not enough CUDA devices for the requested topology, " |
| f"need {self.model_parallel_world_size}, got device_ids={config.device_ids}." |
| ) |
| self.ps = [] |
| self.events = [] |
| ctx = mp.get_context("spawn") |
| for i in range(1, self.model_parallel_world_size): |
| event = ctx.Event() |
| process = ctx.Process(target=_run_model_runner_worker, args=(config, i, event)) |
| process.start() |
| self.ps.append(process) |
| self.events.append(event) |
| self._exited = False |
| self.profile_session = TorchProfileSession("engine") |
| atexit.register(self.exit) |
| self._install_signal_handlers() |
|
|
| try: |
| self.tokenizer = AutoTokenizer.from_pretrained(config.model, use_fast=True, trust_remote_code=True) |
| config.tokenizer_vocab_size = len(self.tokenizer) |
| config.eos = self.tokenizer.eos_token_id |
|
|
| if ( |
| getattr(self.tokenizer, "mask_token_id", None) is not None |
| and config.mask_token_id != self.tokenizer.mask_token_id |
| ): |
| logger.warning( |
| "Overriding mask_token_id from %s to tokenizer mask_token_id %s.", |
| config.mask_token_id, |
| self.tokenizer.mask_token_id, |
| ) |
| config.mask_token_id = self.tokenizer.mask_token_id |
|
|
| self.model_runner = AutoModelRunner.from_config(config, 0, self.events) |
| self.scheduler: SchedulerBase | DataParallelScheduler = AutoScheduler.from_config(config) |
| except BaseException: |
| self.exit() |
| raise |
|
|
| def _install_signal_handlers(self) -> None: |
| if getattr(self, "_signal_handlers_installed", False): |
| return |
| self._signal_handlers_installed = True |
| self._previous_signal_handlers = {} |
| for sig in (signal.SIGINT, signal.SIGTERM): |
| try: |
| previous = signal.getsignal(sig) |
| self._previous_signal_handlers[sig] = previous |
|
|
| def handler(signum, frame, *, _previous=previous): |
| self.exit() |
| if callable(_previous): |
| _previous(signum, frame) |
| else: |
| raise SystemExit(128 + signum) |
|
|
| signal.signal(sig, handler) |
| except Exception: |
| logger.debug("Failed to install signal handler for %s.", sig, exc_info=True) |
|
|
| @staticmethod |
| def _join_or_stop_process(process, *, timeout: float = 5.0) -> None: |
| try: |
| process.join(timeout=timeout) |
| except Exception: |
| logger.debug("Failed to join worker process %s.", getattr(process, "pid", None), exc_info=True) |
| if not process.is_alive(): |
| return |
|
|
| logger.warning("Terminating stale worker process pid=%s.", process.pid) |
| try: |
| process.terminate() |
| except Exception: |
| logger.debug("Failed to terminate worker process %s.", process.pid, exc_info=True) |
| try: |
| process.join(timeout=timeout) |
| except Exception: |
| logger.debug("Failed to join terminated worker process %s.", process.pid, exc_info=True) |
| if not process.is_alive(): |
| return |
|
|
| logger.warning("Killing stale worker process pid=%s.", process.pid) |
| try: |
| process.kill() |
| except Exception: |
| logger.debug("Failed to kill worker process %s.", process.pid, exc_info=True) |
| try: |
| process.join(timeout=timeout) |
| except Exception: |
| logger.debug("Failed to join killed worker process %s.", process.pid, exc_info=True) |
|
|
| def exit(self): |
| if getattr(self, "_exited", False): |
| return |
| self._exited = True |
| if hasattr(self, "profile_session"): |
| self.profile_session.stop() |
| if hasattr(self, "model_runner") and self.model_runner is not None: |
| try: |
| self.model_runner.call("exit") |
| except Exception: |
| pass |
| try: |
| del self.model_runner |
| except Exception: |
| pass |
| for p in getattr(self, "ps", []): |
| self._join_or_stop_process(p) |
| time.sleep(0) |
|
|
| def add_request(self, prompt: str | list[int], sampling_params: SamplingParams): |
| if isinstance(prompt, str): |
| with record_function("diffulex.engine.tokenizer_encode"): |
| prompt = self.tokenizer.encode(prompt) |
|
|
| with record_function("diffulex.engine.add_request"): |
| req = AutoReq.create(self.config, prompt, sampling_params) |
| req.page_size = self.config.kv_cache_page_size |
| with record_function("diffulex.engine.scheduler_add"): |
| self.scheduler.add(req) |
| return req.req_id |
|
|
| def step(self): |
| self.profile_session.start() |
| with record_function("diffulex.engine.scheduler_schedule"): |
| reqs, is_prefill = self.scheduler.schedule() |
| with record_function("diffulex.engine.prepare_reqs_for_execution"): |
| self._prepare_reqs_for_execution(reqs) |
| try: |
| with record_function("diffulex.engine.model_runner_run"): |
| sample_output = self.model_runner.call("run", reqs) |
| finally: |
| with record_function("diffulex.engine.clear_execution_prepared"): |
| self._clear_execution_prepared(reqs) |
| with record_function("diffulex.engine.scheduler_postprocess"): |
| self.scheduler.postprocess(reqs, sample_output) |
| finished_req_ids = [req.req_id for req in reqs if (req.is_completed or req.is_finished)] |
| if finished_req_ids: |
| with record_function("diffulex.engine.evict_sampler_state"): |
| self.model_runner.call("evict_sampler_state", finished_req_ids) |
| self.profile_session.step() |
| return reqs, is_prefill |
|
|
| @staticmethod |
| def _prepare_reqs_for_execution(reqs): |
| for req in reqs: |
| step_fn = getattr(req, "step", None) |
| if callable(step_fn): |
| step_fn() |
| mark_fn = getattr(req, "mark_execution_prepared", None) |
| if callable(mark_fn): |
| mark_fn() |
|
|
| @staticmethod |
| def _clear_execution_prepared(reqs): |
| for req in reqs: |
| clear_fn = getattr(req, "clear_execution_prepared", None) |
| if callable(clear_fn): |
| clear_fn() |
|
|
| def is_finished(self): |
| return self.scheduler.is_finished() |
|
|
| def abort_request(self, req_id: int) -> bool: |
| return self.scheduler.abort_request(req_id) |
|
|
| def generate( |
| self, |
| prompts: list[str] | list[list[int]], |
| sampling_params: SamplingParams | list[SamplingParams], |
| use_tqdm: bool = True, |
| ) -> list[str]: |
| with record_function("diffulex.engine.generate"): |
| if use_tqdm: |
| pbar = tqdm(total=len(prompts), desc="Diffulex Generating", dynamic_ncols=True) |
|
|
| if not isinstance(sampling_params, list): |
| sampling_params = [sampling_params] * len(prompts) |
|
|
| req_id_to_prompt_id = {} |
| for prompt_id, (prompt, sp) in tqdm( |
| enumerate(zip(prompts, sampling_params)), |
| total=len(prompts), |
| desc="Adding Requests to Scheduler", |
| dynamic_ncols=True, |
| ): |
| req_id = self.add_request(prompt, sp) |
| req_id_to_prompt_id[req_id] = prompt_id |
|
|
| step = 0 |
| outputs = GenerationOutputs(len(prompts)) |
| while not self.is_finished(): |
| step += 1 |
|
|
| start = perf_counter() |
| reqs, is_prefill = self.step() |
| step_time = perf_counter() - start |
|
|
| with record_function("diffulex.engine.record_outputs"): |
| outputs.record_step(reqs, step_time, req_id_to_prompt_id) |
|
|
| if use_tqdm: |
| pbar.set_postfix(outputs.fast_postfix()) |
|
|
| for req in reqs: |
| if (req.is_completed or req.is_finished) and use_tqdm: |
| pbar.update(1) |
|
|
| if use_tqdm: |
| pbar.close() |
|
|
| outputs.log_summary() |
| with record_function("diffulex.engine.convert_outputs_to_text"): |
| outputs.convert_to_text(self.tokenizer) |
|
|
| return outputs |
|
|
|
|
| __all__ = ["DiffulexEngine"] |
|
|
