From b1531f430472488cb3e547494375c5c783181b2c Mon Sep 17 00:00:00 2001 From: BioMikeUkr Date: Thu, 14 May 2026 16:49:56 +0300 Subject: [PATCH 1/9] ray serving added --- gliclass/serve/__init__.py | 18 ++ gliclass/serve/__main__.py | 239 +++++++++++++++ gliclass/serve/client.py | 110 +++++++ gliclass/serve/config.py | 113 ++++++++ gliclass/serve/memory.py | 159 ++++++++++ gliclass/serve/server.py | 512 +++++++++++++++++++++++++++++++++ gliclass/serve/server_model.py | 336 ++++++++++++++++++++++ 7 files changed, 1487 insertions(+) create mode 100644 gliclass/serve/__init__.py create mode 100644 gliclass/serve/__main__.py create mode 100644 gliclass/serve/client.py create mode 100644 gliclass/serve/config.py create mode 100644 gliclass/serve/memory.py create mode 100644 gliclass/serve/server.py create mode 100644 gliclass/serve/server_model.py diff --git a/gliclass/serve/__init__.py b/gliclass/serve/__init__.py new file mode 100644 index 0000000..36bcb35 --- /dev/null +++ b/gliclass/serve/__init__.py @@ -0,0 +1,18 @@ +"""GLiClass serving module.""" + +from .client import GLiClassClient +from .config import GLiClassServeConfig +from .memory import GLiClassMemoryEstimator +from .server import GLiClassServer, GLiClassFactory, shutdown, serve_gliclass +from .server_model import GLiClassServerModel + +__all__ = [ + "GLiClassClient", + "GLiClassFactory", + "GLiClassMemoryEstimator", + "GLiClassServeConfig", + "GLiClassServer", + "GLiClassServerModel", + "serve_gliclass", + "shutdown", +] diff --git a/gliclass/serve/__main__.py b/gliclass/serve/__main__.py new file mode 100644 index 0000000..1e64c3b --- /dev/null +++ b/gliclass/serve/__main__.py @@ -0,0 +1,239 @@ +"""CLI entry point for GLiClass serving.""" + +import sys +import signal +import logging +import argparse + +import ray +from ray import serve + +from .config import GLiClassServeConfig +from .server import serve_gliclass + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", +) +logger = logging.getLogger(__name__) + + +def main(): + """Main entry point for GLiClass serving.""" + parser = argparse.ArgumentParser( + description="GLiClass Ray Serve deployment", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + + # Config file + parser.add_argument( + "--config", + type=str, + default=None, + help="Path to YAML config file (CLI args override config values)", + ) + + # Model configuration + parser.add_argument( + "--model", + type=str, + default=None, + help="Model name or path", + ) + parser.add_argument( + "--device", + type=str, + default=None, + help="Device to run on (cuda or cpu)", + ) + parser.add_argument( + "--dtype", + type=str, + default=None, + choices=["float32", "float16", "bfloat16"], + help="Data type for model weights", + ) + parser.add_argument( + "--max-model-len", + type=int, + default=None, + help="Maximum sequence length", + ) + parser.add_argument( + "--max-batch-size", + type=int, + default=None, + help="Maximum batch size", + ) + parser.add_argument( + "--max-labels", + type=int, + default=None, + help="Maximum number of labels (-1 for unlimited)", + ) + parser.add_argument( + "--max-labels-alloc", + type=str, + default=None, + help='Label memory allocation: "dynamic", "fixed", or integer (e.g., "50")', + ) + + # Server configuration + parser.add_argument( + "--host", + type=str, + default="0.0.0.0", + help="Host to bind to", + ) + parser.add_argument( + "--port", + type=int, + default=None, + help="Port to bind to", + ) + parser.add_argument( + "--route-prefix", + type=str, + default=None, + help="HTTP route prefix", + ) + parser.add_argument( + "--num-replicas", + type=int, + default=None, + help="Number of model replicas", + ) + + # Performance configuration + parser.add_argument( + "--calibrate-on-startup", + action="store_true", + default=None, + help="Run memory calibration on startup", + ) + parser.add_argument( + "--precompile-on-startup", + action="store_true", + default=None, + help="Precompile model on startup", + ) + parser.add_argument( + "--use-memory-aware-batching", + action="store_true", + default=None, + help="Use memory-aware dynamic batching", + ) + parser.add_argument( + "--enable-compilation", + action="store_true", + default=None, + help="Enable torch.compile", + ) + parser.add_argument( + "--tokenizer-threads", + type=int, + default=None, + help="Number of tokenizer threads", + ) + + # Calibration configuration + parser.add_argument( + "--calibration-min-batch-size", + type=int, + default=None, + help="Minimum batch size for calibration", + ) + parser.add_argument( + "--calibration-max-batch-size", + type=int, + default=None, + help="Maximum batch size for calibration", + ) + parser.add_argument( + "--calibration-min-seq-len", + type=int, + default=None, + help="Minimum sequence length for calibration", + ) + + args = parser.parse_args() + + if args.config: + logger.info(f"Loading config from: {args.config}") + config = GLiClassServeConfig.from_yaml(args.config) + else: + config = GLiClassServeConfig(model=args.model or "knowledgator/gliclass-edge-v3.0") + + # Convert max_labels_alloc to int if it's a digit string + max_labels_alloc_value = args.max_labels_alloc + if max_labels_alloc_value and max_labels_alloc_value.isdigit(): + max_labels_alloc_value = int(max_labels_alloc_value) + + cli_overrides = { + "model": args.model, + "device": args.device, + "dtype": args.dtype, + "max_model_len": args.max_model_len, + "max_batch_size": args.max_batch_size, + "max_labels": args.max_labels, + "max_labels_alloc": max_labels_alloc_value, + "http_port": args.port, + "route_prefix": args.route_prefix, + "num_replicas": args.num_replicas, + "calibrate_on_startup": args.calibrate_on_startup, + "precompile_on_startup": args.precompile_on_startup, + "use_memory_aware_batching": args.use_memory_aware_batching, + "enable_compilation": args.enable_compilation, + "tokenizer_threads": args.tokenizer_threads, + "calibration_min_batch_size": args.calibration_min_batch_size, + "calibration_max_batch_size": args.calibration_max_batch_size, + "calibration_min_seq_len": args.calibration_min_seq_len, + } + config.update(**cli_overrides) + + logger.info("=" * 60) + logger.info("GLiClass Serve Configuration:") + logger.info(f" Model: {config.model}") + logger.info(f" Device: {config.device}") + logger.info(f" Dtype: {config.dtype}") + logger.info(f" Max model length: {config.max_model_len}") + logger.info(f" Max batch size: {config.max_batch_size}") + logger.info(f" Max labels: {config.max_labels}") + logger.info(f" Max labels alloc: {config.max_labels_alloc}") + logger.info(f" HTTP port: {config.http_port}") + logger.info(f" Route prefix: {config.route_prefix}") + logger.info(f" Num replicas: {config.num_replicas}") + logger.info(f" Calibrate on startup: {config.calibrate_on_startup}") + logger.info(f" Precompile on startup: {config.precompile_on_startup}") + logger.info(f" Memory-aware batching: {config.use_memory_aware_batching}") + logger.info("=" * 60) + + logger.info("Initializing Ray...") + ray.init(ignore_reinit_error=True) + + logger.info("Starting Ray Serve...") + serve.start(http_options={"host": args.host, "port": config.http_port}) + + logger.info(f"Deploying GLiClass with model: {config.model}") + _app = serve_gliclass(config, blocking=False) # Keep reference to prevent GC + + logger.info(f"GLiClass server running at http://{args.host}:{config.http_port}{config.route_prefix}") + logger.info("Press Ctrl+C to stop the server") + + def signal_handler(_sig, _frame): + logger.info("Shutting down server...") + serve.shutdown() + ray.shutdown() + sys.exit(0) + + signal.signal(signal.SIGINT, signal_handler) + signal.signal(signal.SIGTERM, signal_handler) + + import time + + while True: + time.sleep(1) + + +if __name__ == "__main__": + main() diff --git a/gliclass/serve/client.py b/gliclass/serve/client.py new file mode 100644 index 0000000..65d91e4 --- /dev/null +++ b/gliclass/serve/client.py @@ -0,0 +1,110 @@ +"""Client for GLiClass serving endpoint.""" + +import requests + + +class GLiClassClient: + """Client for interacting with GLiClass Ray Serve deployment.""" + + def __init__(self, url: str = "http://localhost:8000/gliclass"): + """Initialize the client. + + Args: + url: Base URL of the GLiClass server + """ + self.url = url.rstrip("/") + + def __call__( + self, + texts: str | list[str], + labels: list[str] | list[list[str]], + threshold: float = 0.5, + multi_label: bool = True, + examples: list[dict] | None = None, + prompt: str | list[str] | None = None, + ) -> list[list[dict]]: + """Classify text(s) - same interface as pipeline. + + Args: + texts: Single text or list of texts to classify + labels: List of labels (same for all) or list of label lists (per text) + threshold: Confidence threshold for predictions + multi_label: Whether to enable multi-label classification + examples: Optional list of example classifications + prompt: Optional task description prompt (string or list) + + Returns: + List of results, one per text. Each result is a list of {"label": ..., "score": ...} + """ + payload = { + "texts": texts, + "labels": labels, + "threshold": threshold, + "multi_label": multi_label, + } + if examples is not None: + payload["examples"] = examples + if prompt is not None: + payload["prompt"] = prompt + + response = requests.post(self.url, json=payload, timeout=30) + response.raise_for_status() + + return response.json() + + def classify( + self, + text: str, + labels: list[str], + threshold: float = 0.5, + multi_label: bool = True, + examples: list[dict] | None = None, + prompt: str | None = None, + ) -> list[dict]: + """Classify a single text (convenience method). + + Args: + text: Input text to classify + labels: List of possible labels + threshold: Confidence threshold for predictions + multi_label: Whether to enable multi-label classification + examples: Optional list of example classifications + prompt: Optional task description prompt + + Returns: + List of predictions: [{"label": ..., "score": ...}, ...] + """ + results = self(text, labels, threshold, multi_label, examples, prompt) + return results[0] + + def health_check(self) -> bool: + """Check if the server is healthy. + + Returns: + True if server is healthy, False otherwise + """ + try: + response = requests.get(f"{self.url}/-/healthz", timeout=5) + return response.status_code == 200 + except requests.RequestException: + return False + + +if __name__ == "__main__": + client = GLiClassClient() + + # Single text + result = client.classify( + text="This is a great product! I love it.", + labels=["positive", "negative", "neutral"], + threshold=0.3, + ) + print("Single prediction:", result) + + # Batch + results = client( + texts=["Great product!", "Terrible experience"], + labels=["positive", "negative", "neutral"], + threshold=0.3, + ) + print("Batch predictions:", results) diff --git a/gliclass/serve/config.py b/gliclass/serve/config.py new file mode 100644 index 0000000..382394b --- /dev/null +++ b/gliclass/serve/config.py @@ -0,0 +1,113 @@ +"""Configuration for GLiClass Ray Serve deployment.""" + +from pathlib import Path +from dataclasses import field, asdict, dataclass + +import yaml + + +@dataclass +class GLiClassServeConfig: + """Configuration for GLiClass Ray Serve deployment. + + This config controls model loading, serving parameters, and dynamic batching behavior. + """ + + model: str + device: str = "cuda" + dtype: str = "bfloat16" + + quantization: str | None = None + + max_model_len: int = 2048 + max_labels: int = -1 + max_labels_alloc: str | int = "dynamic" + + default_threshold: float = 0.5 + + num_replicas: int = 1 + num_gpus_per_replica: float = 1.0 + num_cpus_per_replica: float = 1.0 + + max_batch_size: int = 32 + batch_wait_timeout_ms: float = 20.0 + request_timeout_s: float = 30.0 + max_ongoing_requests: int = 256 + queue_capacity: int = 4096 + + route_prefix: str = "/gliclass" + + tokenizer_threads: int = 4 + + enable_compilation: bool = True + calibrate_on_startup: bool = False + precompile_on_startup: bool = True + use_memory_aware_batching: bool = False + + precompiled_batch_sizes: list[int] = field(default_factory=lambda: [1, 2, 4, 8, 16, 32]) + + target_memory_fraction: float = 0.8 + memory_overhead_factor: float = 1.3 + + calibration_min_seq_len: int = 64 + calibration_min_batch_size: int = 1 + calibration_max_batch_size: int = 64 + calibration_probe_batch_size: int = 2 + + warmup_iterations: int = 3 + + http_port: int = 8000 + + ray_address: str | None = None + + def __post_init__(self): + if self.max_batch_size not in self.precompiled_batch_sizes: + self.precompiled_batch_sizes = sorted(set(self.precompiled_batch_sizes) | {self.max_batch_size}) + self.precompiled_batch_sizes = sorted(self.precompiled_batch_sizes) + + def to_env_vars(self) -> dict: + """Convert config to environment variables for model loading.""" + env = {} + if self.tokenizer_threads > 0: + env["TOKENIZERS_PARALLELISM"] = "true" + return env + + @classmethod + def from_yaml(cls, config_path: str | Path) -> "GLiClassServeConfig": + """Load configuration from YAML file. + + Args: + config_path: Path to YAML config file + + Returns: + GLiClassServeConfig instance + """ + config_path = Path(config_path) + with config_path.open("r") as f: + config_dict = yaml.safe_load(f) + return cls(**config_dict) + + def to_yaml(self, config_path: str | Path) -> None: + """Save configuration to YAML file. + + Args: + config_path: Path to save YAML config + """ + config_path = Path(config_path) + config_path.parent.mkdir(parents=True, exist_ok=True) + with config_path.open("w") as f: + yaml.dump(asdict(self), f, default_flow_style=False, sort_keys=False) + + def update(self, **kwargs) -> "GLiClassServeConfig": + """Update config with provided kwargs (for CLI override). + + Args: + **kwargs: Fields to update (None values are ignored) + + Returns: + Updated config instance + """ + for key, value in kwargs.items(): + if value is not None and hasattr(self, key): + setattr(self, key, value) + return self diff --git a/gliclass/serve/memory.py b/gliclass/serve/memory.py new file mode 100644 index 0000000..85249f1 --- /dev/null +++ b/gliclass/serve/memory.py @@ -0,0 +1,159 @@ +"""Memory estimation for GLiClass via precomputed calibration table. + +Startup calibration runs the model on probe batches at power-of-two sequence +lengths and records peak GPU memory per sample. At request time ``batch_size_fn`` +picks the largest precompiled batch size that satisfies + + per_sample(seq_len) * N <= total_gpu - cuda_context - model_weights + +using a pessimistic (rounded-up) seq_len and a safety factor on per-sample +memory. +""" + +import logging +from typing import Dict, List, Callable + +import torch + +logger = logging.getLogger(__name__) + + +def _power_of_two_seq_lens(max_seq_len: int, min_seq_len: int = 64) -> List[int]: + """Return power-of-two sequence lengths from min_seq_len up to max_seq_len.""" + lens: List[int] = [] + s = max(1, min_seq_len) + while s < max_seq_len: + lens.append(s) + s *= 2 + lens.append(max_seq_len) + return lens + + +class GLiClassMemoryEstimator: + """Precomputed memory table for GLiClass inference.""" + + def __init__( + self, + safety_factor: float = 1.3, + target_memory_fraction: float = 0.9, + calibration_probe_batch_size: int = 2, + ): + self.safety_factor = safety_factor + self.target_memory_fraction = target_memory_fraction + self.calibration_probe_batch_size = max(2, calibration_probe_batch_size) + + self.total_gpu_memory: int = 0 + self.cuda_context_bytes: int = 0 + self.model_memory_bytes: int = 0 + + self.per_sample_table: Dict[int, int] = {} + + def measure_cuda_context(self) -> None: + """Record CUDA context overhead. Must be called before the model loads.""" + if not torch.cuda.is_available(): + return + torch.cuda.synchronize() + free, total = torch.cuda.mem_get_info() + self.total_gpu_memory = total + self.cuda_context_bytes = total - free + logger.info("CUDA context: %.1f MiB", self.cuda_context_bytes / (1024**2)) + + def measure_model_memory(self) -> None: + """Record model weight memory. Must be called after the model loads.""" + if not torch.cuda.is_available(): + return + torch.cuda.synchronize() + torch.cuda.empty_cache() + free, total = torch.cuda.mem_get_info() + self.total_gpu_memory = total + used = total - free + self.model_memory_bytes = max(0, used - self.cuda_context_bytes) + logger.info("Model weights: %.1f MiB", self.model_memory_bytes / (1024**2)) + + def available_memory(self) -> int: + """Budget for a batch: ``total_gpu - cuda_context - model_weights``.""" + if not torch.cuda.is_available(): + return 0 + budget = self.total_gpu_memory - self.cuda_context_bytes - self.model_memory_bytes + return max(0, int(budget * self.target_memory_fraction)) + + def calibrate( + self, + predict_method: Callable, + max_seq_len: int, + min_seq_len: int = 64, + ) -> None: + """Populate ``per_sample_table`` across power-of-two seq lengths. + + Uses a small set of dummy labels to probe classification performance. + """ + if not torch.cuda.is_available(): + return + + seq_lens = _power_of_two_seq_lens(max_seq_len, min_seq_len=min_seq_len) + dummy_labels = ["label1", "label2", "label3"] + probe_b = self.calibration_probe_batch_size + + logger.info("Calibrating memory table: seq_lens=%s, probe_batch=%s", seq_lens, probe_b) + + for seq_len in seq_lens: + dummy_text = "word " * max(1, seq_len // 2) + peak = self._measure_peak(predict_method, [dummy_text] * probe_b, dummy_labels) + per_sample = max(1, peak // probe_b) + self.per_sample_table[seq_len] = per_sample + logger.info(" seq_len=%5d: per_sample=%.1f MiB", seq_len, per_sample / (1024**2)) + + def _measure_peak( + self, + predict_method: Callable, + texts: List[str], + labels: List[str], + ) -> int: + """Run a probe batch and return peak allocated bytes above baseline.""" + torch.cuda.synchronize() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + baseline = torch.cuda.memory_allocated() + + predict_method(texts, labels, threshold=0.5) + + torch.cuda.synchronize() + peak = torch.cuda.max_memory_allocated() + return max(0, peak - baseline) + + def _lookup_seq_len(self, seq_len: int) -> int: + """Round ``seq_len`` up to the nearest calibrated entry (pessimistic).""" + if not self.per_sample_table: + raise RuntimeError("Memory estimator has not been calibrated") + for key in sorted(self.per_sample_table.keys()): + if key >= seq_len: + return key + return max(self.per_sample_table.keys()) + + def per_sample_at(self, seq_len: int) -> int: + """Pessimistic per-sample memory at or above ``seq_len``.""" + probe_seq_len = self._lookup_seq_len(seq_len) + return int(self.per_sample_table[probe_seq_len] * self.safety_factor) + + def batch_size_fn( + self, + seq_len: int, + precompiled_sizes: List[int], + ) -> int: + """Largest precompiled batch size satisfying ``per_sample * N <= budget``. + + Budget = ``total_gpu - cuda_context - model_weights`` (times the + configured ``target_memory_fraction``). + """ + if not precompiled_sizes: + return 1 + + available = self.available_memory() + if available <= 0: + return min(precompiled_sizes) + + per_sample = self.per_sample_at(seq_len) + for size in sorted(precompiled_sizes, reverse=True): + if per_sample * size <= available: + return size + return min(precompiled_sizes) diff --git a/gliclass/serve/server.py b/gliclass/serve/server.py new file mode 100644 index 0000000..fbd8d7c --- /dev/null +++ b/gliclass/serve/server.py @@ -0,0 +1,512 @@ +"""Ray Serve deployment for GLiClass with dynamic batching.""" + +import os +import logging +from typing import Any + +import torch +from ray import serve + +from .config import GLiClassServeConfig +from .memory import GLiClassMemoryEstimator +from .server_model import GLiClassServerModel + +logger = logging.getLogger("ray.serve") + + +class GLiClassServer: + """GLiClass Ray Serve deployment with dynamic batching.""" + + def __init__(self, config: GLiClassServeConfig): + """Initialize GLiClass server deployment. + + Args: + config: Server configuration with model and serving parameters + """ + self.config = config + + env_vars = config.to_env_vars() + for key, value in env_vars.items(): + os.environ[key] = value + + if config.tokenizer_threads > 0: + os.environ["TOKENIZERS_PARALLELISM"] = "true" + torch.set_num_threads(config.tokenizer_threads) + + torch.set_float32_matmul_precision("high") + + dtype_map = { + "float32": torch.float32, + "float16": torch.float16, + "fp16": torch.float16, + "bfloat16": torch.bfloat16, + "bf16": torch.bfloat16, + } + self.torch_dtype = dtype_map.get(config.dtype.lower(), torch.bfloat16) + self.device = torch.device(config.device) + + self.memory_estimator = GLiClassMemoryEstimator( + safety_factor=config.memory_overhead_factor, + target_memory_fraction=config.target_memory_fraction, + calibration_probe_batch_size=config.calibration_probe_batch_size, + ) + + if torch.cuda.is_available(): + self.memory_estimator.measure_cuda_context() + + logger.info("Loading model: %s", config.model) + + self.model = GLiClassServerModel( + model_name=config.model, + device=self.device, + dtype=self.torch_dtype, + max_length=config.max_model_len, + max_classes=config.max_labels if config.max_labels > 0 else 100, + max_labels_alloc=config.max_labels_alloc, + ) + + if torch.cuda.is_available(): + self.memory_estimator.measure_model_memory() + + if config.enable_compilation: + self._precompile() + + if torch.cuda.is_available(): + self._calibrate_memory() + + logger.info("GLiClass server initialized successfully") + + def _precompile(self) -> None: + logger.info("Precompiling model for batch sizes: %s", self.config.precompiled_batch_sizes) + + if hasattr(self.model.model, "compile"): + self.model.model.compile() + + dummy_labels = ["person", "organization", "location"] + + for batch_size in self.config.precompiled_batch_sizes: + dummy_texts = [f"Sample text number {i} for precompilation warmup." for i in range(batch_size)] + + for _ in range(self.config.warmup_iterations): + self._run_batch_internal( + dummy_texts, + dummy_labels, + threshold=0.5, + multi_label=True, + ) + + logger.info(" Batch size %d: compiled", batch_size) + + if torch.cuda.is_available(): + torch.cuda.synchronize() + + logger.info("Precompilation complete.") + + def _calibrate_memory(self) -> None: + logger.info("Calibrating memory table...") + + self.memory_estimator.calibrate( + self._run_batch_internal, + max_seq_len=self.config.max_model_len, + min_seq_len=self.config.calibration_min_seq_len, + ) + + logger.info("Memory calibration complete.") + + def batch_size_fn(self, seq_len: int | None = None) -> int: + """Largest precompiled batch size that fits at seq_len. + + Args: + seq_len: Sequence length (text + label words). If None, uses max_model_len. + + Returns: + Optimal batch size from precompiled sizes + """ + if not torch.cuda.is_available(): + return self.config.precompiled_batch_sizes[-1] + + if seq_len is None: + seq_len = self.config.max_model_len + + return self.memory_estimator.batch_size_fn( + seq_len=seq_len, + precompiled_sizes=self.config.precompiled_batch_sizes, + ) + + def observed_seq_len( + self, + texts: list[str], + labels: list[str] | None = None, + ) -> int: + """Total input word count: longest text + all label words. + + Labels are concatenated into input, so they extend effective seq length + for every sample in the batch. + + Args: + texts: Input texts + labels: Label list + + Returns: + Estimated sequence length + """ + max_text_words = max((len(t.split()) for t in texts if t.strip()), default=0) + prompt_words = 0 + if labels: + prompt_words += sum(len(label.split()) for label in labels) + total = max_text_words + prompt_words + return min(max(total, self.config.calibration_min_seq_len), self.config.max_model_len) + + def _filter_labels(self, labels: list[str]) -> list[str]: + if self.config.max_labels > 0 and len(labels) > self.config.max_labels: + logger.warning("Truncating labels from %d to %d", len(labels), self.config.max_labels) + return labels[: self.config.max_labels] + return labels + + @torch.inference_mode() + def _run_batch_internal( + self, + texts: list[str], + labels: list[str], + threshold: float = 0.5, + multi_label: bool = True, + examples: list[dict[str, Any]] | None = None, + prompt: str | list[str] | None = None, + ) -> list[dict[str, Any]]: + """Run batch inference using low-level methods. + + This is the core inference method that uses prepare_batch, tokenize_batch, + collate_batch, run_batch, and decode_batch directly. + + Args: + texts: List of input texts + labels: Label list (same for all texts) + threshold: Classification threshold + multi_label: Multi-label classification + examples: Few-shot examples + prompt: Task description + + Returns: + List of prediction dicts + """ + prepared = self.model.prepare_batch(texts, labels, examples, prompt) + tokenized = self.model.tokenize_batch(prepared) + batch = self.model.collate_batch(tokenized) + outputs = self.model.run_batch(batch) + results = self.model.decode_batch(outputs, threshold, multi_label) + return results + + def predict( + self, + texts: str | list[str], + labels: list[str], + threshold: float | None = None, + multi_label: bool = True, + examples: list[dict[str, Any]] | None = None, + prompt: str | list[str] | None = None, + ) -> list[dict[str, Any]]: + if isinstance(texts, str): + texts = [texts] + + if threshold is None: + threshold = self.config.default_threshold + + labels = self._filter_labels(labels) + + results = self._run_batch_internal( + texts, + labels, + threshold=threshold, + multi_label=multi_label, + examples=examples, + prompt=prompt, + ) + + return results + + +def _build_deployment(config: GLiClassServeConfig): + batch_wait_s = max(config.batch_wait_timeout_ms, 0.0) / 1000.0 + initial_max_batch_size = config.max_batch_size + + @serve.deployment( + num_replicas=config.num_replicas, + ray_actor_options={ + "num_gpus": config.num_gpus_per_replica, + "num_cpus": config.num_cpus_per_replica, + }, + max_ongoing_requests=config.max_ongoing_requests, + ) + class GLiClassDeployment: + def __init__(self, serve_config: GLiClassServeConfig): + self.server = GLiClassServer(serve_config) + # Initialize dynamic batch sizing + self._infer_batch.set_max_batch_size(self.server.batch_size_fn()) + logger.info( + "Ray Serve batch size initialized to %d (precompiled: %s)", + self.server.batch_size_fn(), + serve_config.precompiled_batch_sizes, + ) + + @serve.batch( + max_batch_size=initial_max_batch_size, + batch_wait_timeout_s=batch_wait_s, + ) + async def _infer_batch( + self, + texts: list[str], + labels_list: list[list[str]], + thresholds: list[float], + multi_label_list: list[bool], + examples_list: list[list[dict[str, Any]] | None], + prompts_list: list[str | None], + ) -> list[list[dict[str, Any]]]: + """Single forward pass over the Ray-accumulated batch. + + Before dispatch, re-sizes Ray's batcher via set_max_batch_size + using batch_size_fn on the observed seq length — so the next + accumulation picks the largest precompiled size that fits. + + Assumes batch requests are homogeneous — labels/thresholds/flags + are taken from the first request. + """ + # Dynamically adjust batch size based on observed sequence length + next_max_batch = self.server.batch_size_fn( + seq_len=self.server.observed_seq_len( + texts, + labels=labels_list[0] if labels_list else None, + ) + ) + self._infer_batch.set_max_batch_size(next_max_batch) + + # Use first request's parameters (assume homogeneous) + labels = labels_list[0] + threshold = thresholds[0] + multi_label = multi_label_list[0] + examples = examples_list[0] + prompt = prompts_list[0] + + # Process entire batch at once + results = self.server._run_batch_internal( + texts, + labels, + threshold=threshold, + multi_label=multi_label, + examples=examples, + prompt=prompt, + ) + + return results + + async def predict( + self, + text: str, + labels: list[str], + threshold: float | None = None, + multi_label: bool = True, + examples: list[dict[str, Any]] | None = None, + prompt: str | None = None, + ) -> list[dict[str, Any]]: + """Single prediction endpoint - one text per request.""" + if threshold is None: + threshold = self.server.config.default_threshold + + # Call batched method - Ray will accumulate these + results = await self._infer_batch( + text, + labels, + threshold, + multi_label, + examples, + prompt, + ) + return results + + async def __call__(self, request) -> list[dict[str, Any]]: + """HTTP endpoint - accepts single text per request.""" + payload = await request.json() + text = payload.get("text") or payload.get("texts") + if isinstance(text, list): + # If list provided, take first element for compatibility + text = text[0] if text else "" + return await self.predict( + text=text, + labels=payload["labels"], + threshold=payload.get("threshold"), + multi_label=payload.get("multi_label", True), + examples=payload.get("examples"), + prompt=payload.get("prompt"), + ) + + return GLiClassDeployment.bind(config) + + +def serve_gliclass( + config: GLiClassServeConfig, + blocking: bool = False, +) -> Any: + import ray + + if not ray.is_initialized(): + ray.init(address=config.ray_address, ignore_reinit_error=True) + + serve.start(detached=True, http_options={"port": config.http_port}) + + app = _build_deployment(config) + handle = serve.run(app, name="gliclass", route_prefix=config.route_prefix) + + logger.info("GLiClass server running at http://localhost:%d%s", config.http_port, config.route_prefix) + + if blocking: + import time + import signal + + shutdown_event = False + + def handle_signal(_signum, _frame): + nonlocal shutdown_event + shutdown_event = True + + signal.signal(signal.SIGINT, handle_signal) + signal.signal(signal.SIGTERM, handle_signal) + + while not shutdown_event: + time.sleep(1) + + serve.shutdown() + + return handle + + +def shutdown() -> None: + serve.shutdown() + + +class GLiClassFactory: + """Synchronous facade: config → deploy → predict → shutdown in one object. + + Pass list of texts to preserve dynamic batching - Ray Serve accumulates + concurrent requests into single forward pass. + + Example: + >>> from serve import GLiClassFactory + >>> llm = GLiClassFactory(model="knowledgator/gliclass-edge-v3.0") + >>> outputs = llm.predict( + ... ["Great product!", "Terrible service"], + ... labels=["positive", "negative", "neutral"], + ... ) + >>> llm.shutdown() + + Or as context manager: + >>> with GLiClassFactory(model="knowledgator/gliclass-edge-v3.0") as llm: + ... out = llm.predict("Great product!", ["positive", "negative"]) + """ + + def __init__( + self, + model: str | None = None, + *, + config: GLiClassServeConfig | None = None, + **kwargs, + ): + """Pass either `config` or `model`/kwargs, not both.""" + if config is not None: + if model is not None or kwargs: + raise ValueError("Pass either `config` or `model`/kwargs, not both.") + else: + if model is None: + raise ValueError("Must provide either `model` or `config`.") + config = GLiClassServeConfig(model=model, **kwargs) + + self.config = config + self._handle = serve_gliclass(config, blocking=False) + self._closed = False + + @property + def handle(self): + """Underlying Ray Serve deployment handle for async/advanced use.""" + return self._handle + + def predict( + self, + texts: str | list[str], + labels: list[str], + threshold: float | None = None, + multi_label: bool = True, + examples: list[dict[str, Any]] | None = None, + prompt: str | list[str] | None = None, + ) -> dict[str, Any] | list[dict[str, Any]]: + """Blocking prediction. Returns dict for str input, list for list input.""" + single = isinstance(texts, str) + items = [texts] if single else list(texts) + + refs = [ + self._handle.predict.remote( + t, + labels, + threshold, + multi_label, + examples, + prompt, + ) + for t in items + ] + results = [ref.result() for ref in refs] + return results[0] if single else results + + async def predict_async( + self, + texts: str | list[str], + labels: list[str], + threshold: float | None = None, + multi_label: bool = True, + examples: list[dict[str, Any]] | None = None, + prompt: str | list[str] | None = None, + ) -> dict[str, Any] | list[dict[str, Any]]: + """Async prediction. Concurrent calls accumulate into one batch.""" + import asyncio + + single = isinstance(texts, str) + items = [texts] if single else list(texts) + + refs = [ + self._handle.predict.remote( + t, + labels, + threshold, + multi_label, + examples, + prompt, + ) + for t in items + ] + results = list(await asyncio.gather(*refs)) + return results[0] if single else results + + def shutdown(self) -> None: + """Tear down Ray Serve deployment and Ray runtime. + + Idempotent. Shutting down Ray after Serve avoids leaving driver + attached to detached Serve instance. + """ + if self._closed: + return + import ray + + serve.shutdown() + if ray.is_initialized(): + ray.shutdown() + self._closed = True + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.shutdown() + return False + + def __del__(self): + try: + self.shutdown() + except Exception: + pass diff --git a/gliclass/serve/server_model.py b/gliclass/serve/server_model.py new file mode 100644 index 0000000..194d5cd --- /dev/null +++ b/gliclass/serve/server_model.py @@ -0,0 +1,336 @@ +"""Server-optimized model wrapper for GLiClass with low-level batch processing.""" + +import logging +from typing import Any + +import torch +from transformers import AutoTokenizer + +from gliclass.model import GLiClassModel +from gliclass.pipeline import format_examples_prompt + +logger = logging.getLogger(__name__) + + +class GLiClassServerModel: + """Server-optimized wrapper combining pipeline logic and model. + + Splits inference into low-level stages for CPU/GPU parallelism: + - prepare_batch (CPU): text formatting, label flattening + - tokenize_batch (CPU): tokenization, input creation + - collate_batch (CPU): tensor creation, padding + - run_batch (GPU): forward pass, encoding + - decode_batch (CPU/GPU): sigmoid, threshold filtering, result formatting + + Avoids DataLoader overhead by keeping tokenizers/collators alive + and calling model methods directly. + """ + + def __init__( + self, + model_name: str, + device: torch.device, + dtype: torch.dtype, + max_length: int = 2048, + max_classes: int = 25, + max_labels_alloc: str | int = "dynamic", # "dynamic" | int | "fixed" + ): + self.device = device + self.dtype = dtype + self.max_length = max_length + self.max_classes = max_classes + self.max_labels_alloc = max_labels_alloc + + logger.info(f"Loading model: {model_name}") + self.model = GLiClassModel.from_pretrained(model_name) + self.model.config.max_labels_alloc = max_labels_alloc + self.model.to(device=device, dtype=dtype) + self.model.eval() + + self.architecture_type = self.model.config.architecture_type + self.prompt_first = getattr(self.model.config, "prompt_first", False) + + logger.info(f"Loading tokenizers for {self.architecture_type}") + # Use model_name, not encoder_model_name, to get tokenizer with special tokens + self.tokenizer = AutoTokenizer.from_pretrained(model_name) + + if self.architecture_type in {"bi-encoder", "bi-encoder-fused"}: + # For bi-encoder, load labels tokenizer from model (not label_model_name) + # to ensure special tokens are available + self.labels_tokenizer = AutoTokenizer.from_pretrained(model_name) + else: + self.labels_tokenizer = None + + self.label_token = "<