From 42f8bdd84f7a0bbcf788a5dfb1bcea8d37472167 Mon Sep 17 00:00:00 2001 From: Ingvar Date: Sun, 31 May 2026 17:28:09 +0300 Subject: [PATCH 1/3] implement multi-lora serving --- gliclass/model.py | 30 +++++-- gliclass/pipeline.py | 22 ++++- gliclass/serve/__main__.py | 5 +- gliclass/serve/client.py | 17 +++- gliclass/serve/config.py | 24 +++++- gliclass/serve/server.py | 171 ++++++++++++++++++++++++++++++++++--- 6 files changed, 242 insertions(+), 27 deletions(-) diff --git a/gliclass/model.py b/gliclass/model.py index 5857717..bc0dd97 100644 --- a/gliclass/model.py +++ b/gliclass/model.py @@ -795,8 +795,15 @@ def pool_outputs(self, encoder_outputs): text_embeddings = nn.functional.normalize(text_embeddings, p=2, dim=-1, eps=self.epsilon) return text_embeddings - def encode_text(self, input_ids, attention_mask): - outputs = self.encoder_model(input_ids.squeeze(1), attention_mask=attention_mask.squeeze(1)) + def encode_text(self, input_ids, attention_mask, adapter_ids=None): + encoder_kwargs = {} + if adapter_ids is not None: + encoder_kwargs["adapter_ids"] = adapter_ids + outputs = self.encoder_model( + input_ids.squeeze(1), + attention_mask=attention_mask.squeeze(1), + **encoder_kwargs, + ) text_embeddings = self.pool_outputs(outputs) return text_embeddings @@ -843,11 +850,12 @@ def forward( output_text_embeddings: bool | None = None, output_class_embeddings: bool | None = None, return_dict: bool | None = None, + adapter_ids: list[str] | None = None, **kwargs, ) -> Tuple | SequenceClassifierOutput: return_dict = return_dict if return_dict is not None else self.config.use_return_dict - text_embeddings = self.encode_text(input_ids, attention_mask) + text_embeddings = self.encode_text(input_ids, attention_mask, adapter_ids=adapter_ids) class_embeddings = self.encode_classes(class_input_ids, class_attention_mask, labels_mask) logits = self.scorer(text_embeddings, class_embeddings) * self.logit_scale.to(class_embeddings.device) @@ -872,7 +880,7 @@ class GLiClassBiEncoderFused(GLiClassBiEncoder): def __init__(self, config: GLiClassModelConfig, from_pretrained=False): super().__init__(config, from_pretrained) - def encode_text(self, input_ids, attention_mask, class_embeddings, labels_mask): + def encode_text(self, input_ids, attention_mask, class_embeddings, labels_mask, adapter_ids=None): embedding_layer = self.encoder_model.get_input_embeddings() inputs_embeds = embedding_layer(input_ids) @@ -884,7 +892,14 @@ def encode_text(self, input_ids, attention_mask, class_embeddings, labels_mask): selected_class_embeddings = class_embeddings[labels_batch_indices, labels_indices] inputs_embeds[batch_indices, class_token_indices] = selected_class_embeddings - encoder_outputs = self.encoder_model(inputs_embeds=inputs_embeds, attention_mask=attention_mask.squeeze(1)) + encoder_kwargs = {} + if adapter_ids is not None: + encoder_kwargs["adapter_ids"] = adapter_ids + encoder_outputs = self.encoder_model( + inputs_embeds=inputs_embeds, + attention_mask=attention_mask.squeeze(1), + **encoder_kwargs, + ) post_class_embeddings = torch.zeros_like(class_embeddings) post_class_embeddings[labels_batch_indices, labels_indices] = encoder_outputs[0][ @@ -903,6 +918,7 @@ def forward( output_text_embeddings: bool | None = None, output_class_embeddings: bool | None = None, return_dict: bool | None = None, + adapter_ids: list[str] | None = None, **kwargs, ) -> Tuple | SequenceClassifierOutput: return_dict = return_dict if return_dict is not None else self.config.use_return_dict @@ -910,7 +926,7 @@ def forward( raw_class_embeddings = self.encode_classes(class_input_ids, class_attention_mask, labels_mask) encoder_outputs, class_embeddings = self.encode_text( - input_ids, attention_mask, raw_class_embeddings, labels_mask + input_ids, attention_mask, raw_class_embeddings, labels_mask, adapter_ids=adapter_ids ) text_embeddings = self.pool_outputs(encoder_outputs) @@ -1037,5 +1053,7 @@ def resize_token_embeddings(self, new_num_tokens: int | None = None, pad_to_mult return model_embeds def forward(self, *args, **kwargs): + if kwargs.get("adapter_ids") is None: + kwargs.pop("adapter_ids", None) outputs = self.model(*args, **kwargs) return outputs diff --git a/gliclass/pipeline.py b/gliclass/pipeline.py index 44743f2..6d05cf0 100644 --- a/gliclass/pipeline.py +++ b/gliclass/pipeline.py @@ -249,6 +249,15 @@ def _normalize_classification_types( normalized = self._normalize_classification_type(classification_type) return [normalized] * num_texts + def _normalize_adapter_ids(self, adapter_ids: str | List[str] | None, num_texts: int) -> List[str] | None: + if adapter_ids is None: + return None + if isinstance(adapter_ids, str): + return [adapter_ids] * num_texts + if len(adapter_ids) != num_texts: + raise ValueError("Length of adapter_ids list must match number of texts.") + return adapter_ids + def _process_labels( self, labels: List[str] | Dict[str, Any] | List[List[str]] | List[Dict[str, Any]] ) -> List[str] | List[List[str]]: @@ -405,6 +414,7 @@ def __call__( examples: List[Dict[str, Any]] | None = None, prompt: str | List[str] | None = None, return_hierarchical: bool = False, + adapter_ids: str | List[str] | None = None, ): """ Perform zero-shot classification. @@ -421,6 +431,7 @@ def __call__( examples: Few-shot examples with 'text' and 'labels'/'true_labels' keys prompt: Task description - string (same for all) or list (per-text) return_hierarchical: If True, return hierarchical structure with all scores + adapter_ids: Optional LoRA adapter id for all texts or one adapter id per text. Returns: List of classification results or hierarchical dict structure. @@ -430,6 +441,7 @@ def __call__( texts = self._normalize_texts(texts) thresholds = self._normalize_thresholds(threshold, len(texts)) classification_types = self._normalize_classification_types(classification_type, len(texts)) + adapter_ids = self._normalize_adapter_ids(adapter_ids, len(texts)) if rac_examples: if len(texts) == 1 and not isinstance(rac_examples[0], list): @@ -460,12 +472,17 @@ def __call__( batch_examples = self._get_batch_examples(examples, idx, len(batch_texts)) batch_prompt = self._get_batch_prompt(prompt, idx, len(batch_texts)) + batch_adapter_ids = adapter_ids[idx : idx + len(batch_texts)] if adapter_ids is not None else None tokenized_inputs = self.prepare_inputs( batch_texts, batch_labels, same_labels, examples=batch_examples, prompt=batch_prompt ) max_num_classes = self._resolve_max_num_classes(batch_labels, same_labels) - model_output = self.model(**tokenized_inputs, max_num_classes=max_num_classes) + model_output = self.model( + **tokenized_inputs, + max_num_classes=max_num_classes, + adapter_ids=batch_adapter_ids, + ) logits = model_output.logits probs = torch.sigmoid(logits) @@ -894,6 +911,7 @@ def __call__( examples: List[Dict[str, Any]] | None = None, prompt: str | List[str] | None = None, return_hierarchical: bool = False, + adapter_ids: str | List[str] | None = None, ): """ Perform zero-shot classification. @@ -913,6 +931,7 @@ def __call__( examples: Few-shot examples, each with 'text' and 'labels' keys prompt: Task description - string or list of strings (per-text) return_hierarchical: If True, return structure matching input labels + adapter_ids: Optional LoRA adapter id for all texts or one adapter id per text. Returns: List of predictions (flat) or hierarchical dicts with all scores @@ -927,6 +946,7 @@ def __call__( examples=examples, prompt=prompt, return_hierarchical=return_hierarchical, + adapter_ids=adapter_ids, ) diff --git a/gliclass/serve/__main__.py b/gliclass/serve/__main__.py index 1e64c3b..133ce5c 100644 --- a/gliclass/serve/__main__.py +++ b/gliclass/serve/__main__.py @@ -211,11 +211,8 @@ def main(): 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 + _app = serve_gliclass(config, blocking=False, host=args.host) # 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") diff --git a/gliclass/serve/client.py b/gliclass/serve/client.py index 65d91e4..c6aa0b0 100644 --- a/gliclass/serve/client.py +++ b/gliclass/serve/client.py @@ -22,6 +22,7 @@ def __call__( multi_label: bool = True, examples: list[dict] | None = None, prompt: str | list[str] | None = None, + adapter_id: str | None = None, ) -> list[list[dict]]: """Classify text(s) - same interface as pipeline. @@ -32,6 +33,7 @@ def __call__( multi_label: Whether to enable multi-label classification examples: Optional list of example classifications prompt: Optional task description prompt (string or list) + adapter_id: Optional LoRA adapter id Returns: List of results, one per text. Each result is a list of {"label": ..., "score": ...} @@ -46,6 +48,8 @@ def __call__( payload["examples"] = examples if prompt is not None: payload["prompt"] = prompt + if adapter_id is not None: + payload["adapter_id"] = adapter_id response = requests.post(self.url, json=payload, timeout=30) response.raise_for_status() @@ -60,6 +64,7 @@ def classify( multi_label: bool = True, examples: list[dict] | None = None, prompt: str | None = None, + adapter_id: str | None = None, ) -> list[dict]: """Classify a single text (convenience method). @@ -70,13 +75,23 @@ def classify( multi_label: Whether to enable multi-label classification examples: Optional list of example classifications prompt: Optional task description prompt + adapter_id: Optional LoRA adapter id Returns: List of predictions: [{"label": ..., "score": ...}, ...] """ - results = self(text, labels, threshold, multi_label, examples, prompt) + results = self(text, labels, threshold, multi_label, examples, prompt, adapter_id) return results[0] + def adapter_cache_status(self, adapter_id: str | None = None) -> dict: + params = {"adapter_id": adapter_id} if adapter_id is not None else None + response = requests.get(f"{self.url}/adapter-cache", params=params, timeout=30) + response.raise_for_status() + return response.json() + + def is_adapter_cached(self, adapter_id: str) -> bool: + return bool(self.adapter_cache_status(adapter_id).get("cached")) + def health_check(self) -> bool: """Check if the server is healthy. diff --git a/gliclass/serve/config.py b/gliclass/serve/config.py index 382394b..836f20b 100644 --- a/gliclass/serve/config.py +++ b/gliclass/serve/config.py @@ -39,9 +39,9 @@ class GLiClassServeConfig: tokenizer_threads: int = 4 - enable_compilation: bool = True + enable_compilation: bool = False calibrate_on_startup: bool = False - precompile_on_startup: bool = True + precompile_on_startup: bool = False use_memory_aware_batching: bool = False precompiled_batch_sizes: list[int] = field(default_factory=lambda: [1, 2, 4, 8, 16, 32]) @@ -60,6 +60,17 @@ class GLiClassServeConfig: ray_address: str | None = None + enable_mlora: bool = False + mlora_adapter_weight_modules: list[str] | None = None + mlora_max_rank: int = 16 + mlora_max_gpu_adapters: int = 8 + mlora_max_cpu_adapters: int | None = 128 + mlora_disk_cache_dir: str | None = None + mlora_max_disk_adapters: int | None = None + mlora_base_adapter_id: str = "__base__" + mlora_use_triton_kernels: bool = True + mlora_adapter_id_pattern: str = r"^[A-Za-z0-9_.-]{1,128}$" + 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}) @@ -84,7 +95,14 @@ def from_yaml(cls, config_path: str | Path) -> "GLiClassServeConfig": """ config_path = Path(config_path) with config_path.open("r") as f: - config_dict = yaml.safe_load(f) + config_dict = yaml.safe_load(f) or {} + config_dict.pop("mlora_lazy_load_adapters", None) + config_dict.pop("mlora_allow_runtime_adapter_loading", None) + config_dict.pop("mlora_allow_unsafe_bin_adapters", None) + if "mlora_target_modules" in config_dict and "mlora_adapter_weight_modules" not in config_dict: + config_dict["mlora_adapter_weight_modules"] = config_dict.pop("mlora_target_modules") + else: + config_dict.pop("mlora_target_modules", None) return cls(**config_dict) def to_yaml(self, config_path: str | Path) -> None: diff --git a/gliclass/serve/server.py b/gliclass/serve/server.py index 28e12fd..18bf272 100644 --- a/gliclass/serve/server.py +++ b/gliclass/serve/server.py @@ -1,11 +1,15 @@ """Ray Serve deployment for GLiClass with dynamic batching.""" import os +import re import logging from typing import Any import torch from ray import serve +from ray.serve._private import api as serve_private_api +from ray.serve._private.build_app import build_app +from starlette.responses import JSONResponse from transformers import AutoTokenizer from gliclass.model import GLiClassModel @@ -27,6 +31,8 @@ def __init__(self, config: GLiClassServeConfig): config: Server configuration with model and serving parameters """ self.config = config + self._mlora_model = None + self._adapter_id_re = re.compile(config.mlora_adapter_id_pattern) env_vars = config.to_env_vars() for key, value in env_vars.items(): @@ -60,6 +66,8 @@ def __init__(self, config: GLiClassServeConfig): logger.info("Loading model: %s", config.model) self.model = GLiClassModel.from_pretrained(config.model) + if config.enable_mlora: + self._initialize_mlora() self.model.config.max_labels_alloc = config.max_labels_alloc self.model.to(device=self.device, dtype=self.torch_dtype) self.model.eval() @@ -89,6 +97,96 @@ def __init__(self, config: GLiClassServeConfig): logger.info("GLiClass server initialized successfully") + def _initialize_mlora(self) -> None: + try: + from mlora import CustomLoraConfig, CustomPeftModel + except ImportError as exc: + raise ImportError("enable_mlora=True requires the mlora package to be importable") from exc + + target_model = self._get_mlora_target_model() + mlora_config = CustomLoraConfig( + max_gpu_adapters=self.config.mlora_max_gpu_adapters, + max_cpu_adapters=self.config.mlora_max_cpu_adapters, + disk_cache_dir=self.config.mlora_disk_cache_dir, + max_disk_adapters=self.config.mlora_max_disk_adapters, + max_rank=self.config.mlora_max_rank, + target_modules=self.config.mlora_adapter_weight_modules, + base_adapter_id=self.config.mlora_base_adapter_id, + use_triton_kernels=self.config.mlora_use_triton_kernels, + ) + self._mlora_model = CustomPeftModel(target_model, mlora_config) + self._set_mlora_target_model(self._mlora_model) + + logger.info( + "mLoRA enabled with %d GPU slots, max rank %d, disk cache %s", + self.config.mlora_max_gpu_adapters, + self.config.mlora_max_rank, + self.config.mlora_disk_cache_dir or "disabled", + ) + + def _get_mlora_target_model(self): + architecture = self.model.config.architecture_type + if architecture in {"uni-encoder", "bi-encoder", "bi-encoder-fused"}: + return self.model.model.encoder_model + if architecture in {"encoder-decoder", "encoder-decoder-cls"}: + return self.model.model.encoder_decoder_model + raise NotImplementedError(f"mLoRA is not implemented for architecture {architecture!r}") + + def _set_mlora_target_model(self, wrapped_model) -> None: + architecture = self.model.config.architecture_type + if architecture in {"uni-encoder", "bi-encoder", "bi-encoder-fused"}: + self.model.model.encoder_model = wrapped_model + elif architecture in {"encoder-decoder", "encoder-decoder-cls"}: + self.model.model.encoder_decoder_model = wrapped_model + else: + raise NotImplementedError(f"mLoRA is not implemented for architecture {architecture!r}") + + def _validate_adapter_id(self, adapter_id: str) -> None: + if not isinstance(adapter_id, str) or not self._adapter_id_re.fullmatch(adapter_id): + raise ValueError("adapter_id must match mlora_adapter_id_pattern") + if adapter_id == self.config.mlora_base_adapter_id: + raise ValueError(f"{adapter_id!r} is reserved for base-only inference") + + def adapter_cache_status(self, adapter_id: str | None = None) -> dict[str, Any]: + if not self.config.enable_mlora or self._mlora_model is None: + return {"enabled": False, "base_adapter_id": self.config.mlora_base_adapter_id} + store = self._mlora_model.adapter_store + disk_cache = getattr(store, "disk_cache", None) + response: dict[str, Any] = { + "enabled": True, + "base_adapter_id": self.config.mlora_base_adapter_id, + "loaded": sorted(store.adapters.keys()), + "disk_cached": sorted(disk_cache.entries.keys()) if disk_cache is not None else [], + "disk_cache_dir": str(disk_cache.cache_dir) if disk_cache is not None else None, + "max_disk_adapters": disk_cache.max_adapters if disk_cache is not None else None, + "gpu_slots": list(self._mlora_model.adapter_cache.slot_to_adapter), + } + if adapter_id is not None: + if adapter_id == self.config.mlora_base_adapter_id: + response["adapter_id"] = adapter_id + response["cached"] = True + response["cpu_resident"] = False + response["gpu_resident"] = True + return response + self._validate_adapter_id(adapter_id) + response["adapter_id"] = adapter_id + response["cached"] = disk_cache is not None and adapter_id in disk_cache + response["cpu_resident"] = adapter_id in store.adapters + response["gpu_resident"] = adapter_id in self._mlora_model.adapter_cache.adapter_to_slot + return response + + def ensure_adapter_loaded(self, adapter_id: str | None) -> str | None: + if adapter_id is None: + return self.config.mlora_base_adapter_id if self.config.enable_mlora else None + if adapter_id == self.config.mlora_base_adapter_id: + return adapter_id + if not self.config.enable_mlora or self._mlora_model is None: + raise KeyError(f"Unknown LoRA adapter id: {adapter_id}") + self._validate_adapter_id(adapter_id) + if adapter_id in self._mlora_model.adapter_store: + return adapter_id + raise KeyError(f"Unknown LoRA adapter id: {adapter_id}") + def _precompile(self) -> None: logger.info("Precompiling model for batch sizes: %s", self.config.precompiled_batch_sizes) @@ -188,6 +286,7 @@ def _run_batch_internal( multi_label: bool | list[bool] = True, examples: list[dict[str, Any]] | list[list[dict[str, Any]] | None] | None = None, prompt: str | list[str] | None = None, + adapter_ids: str | list[str] | None = None, ) -> list[list[dict[str, Any]]]: """Run batch inference using the shared zero-shot pipeline. @@ -198,6 +297,7 @@ def _run_batch_internal( multi_label: Shared mode or one mode flag per text examples: Shared examples or one example set per text prompt: Shared prompt or one prompt per text + adapter_ids: Shared LoRA adapter id or one adapter id per text Returns: List of prediction dicts @@ -207,6 +307,11 @@ def _run_batch_internal( else: classification_type = "multi-label" if multi_label else "single-label" + if isinstance(adapter_ids, list): + adapter_ids = [self.ensure_adapter_loaded(adapter_id) for adapter_id in adapter_ids] + elif adapter_ids is not None: + adapter_ids = self.ensure_adapter_loaded(adapter_ids) + return self.pipeline( texts, labels, @@ -215,6 +320,7 @@ def _run_batch_internal( classification_type=classification_type, examples=examples, prompt=prompt, + adapter_ids=adapter_ids, ) def predict( @@ -225,6 +331,7 @@ def predict( multi_label: bool = True, examples: list[dict[str, Any]] | None = None, prompt: str | list[str] | None = None, + adapter_id: str | None = None, ) -> list[list[dict[str, Any]]]: if isinstance(texts, str): texts = [texts] @@ -241,6 +348,7 @@ def predict( multi_label=multi_label, examples=examples, prompt=prompt, + adapter_ids=adapter_id, ) return results @@ -281,6 +389,7 @@ async def _infer_batch( multi_label_list: list[bool], examples_list: list[list[dict[str, Any]] | None], prompts_list: list[str | None], + adapter_ids: list[str | None], ) -> list[list[dict[str, Any]]]: """Single forward pass over the Ray-accumulated batch. @@ -309,6 +418,7 @@ async def _infer_batch( multi_label=multi_label_list, examples=examples_list, prompt=prompts_list, + adapter_ids=adapter_ids, ) return results @@ -321,6 +431,7 @@ async def predict( multi_label: bool = True, examples: list[dict[str, Any]] | None = None, prompt: str | None = None, + adapter_id: str | None = None, ) -> list[dict[str, Any]]: """Single prediction endpoint - one text per request.""" if threshold is None: @@ -334,24 +445,39 @@ async def predict( multi_label, examples, prompt, + adapter_id, ) return results async def __call__(self, request) -> list[dict[str, Any]]: """HTTP endpoint - accepts single text per request.""" + path = request.url.path.rstrip("/") + if path.endswith("/adapter-cache"): + adapter_id = request.query_params.get("adapter_id") + try: + return self.server.adapter_cache_status(adapter_id) + except ValueError as exc: + return JSONResponse({"error": str(exc)}, status_code=400) + 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"), - ) + try: + 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"), + adapter_id=payload.get("adapter_id"), + ) + except KeyError as exc: + return JSONResponse({"error": str(exc)}, status_code=404) + except ValueError as exc: + return JSONResponse({"error": str(exc)}, status_code=400) return GLiClassDeployment.bind(config) @@ -359,18 +485,35 @@ async def __call__(self, request) -> list[dict[str, Any]]: def serve_gliclass( config: GLiClassServeConfig, blocking: bool = False, + host: str = "127.0.0.1", ) -> 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}) + serve.start( + detached=True, + http_options={"host": host, "port": config.http_port}, + ) + serve_client = serve_private_api._get_global_client(_health_check_controller=True) 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) + built_app = build_app( + app, + name="gliclass", + route_prefix=config.route_prefix, + default_runtime_env=ray.get_runtime_context().runtime_env, + ) + handle = serve_client.deploy_applications([built_app])[0] + serve_client.wait_for_proxies_serving() + + logger.info( + "GLiClass server running at http://%s:%d%s", + host, + config.http_port, + config.route_prefix, + ) if blocking: import time @@ -450,6 +593,7 @@ def predict( multi_label: bool = True, examples: list[dict[str, Any]] | None = None, prompt: str | list[str] | None = None, + adapter_id: 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) @@ -463,6 +607,7 @@ def predict( multi_label, examples, prompt, + adapter_id, ) for t in items ] @@ -477,6 +622,7 @@ async def predict_async( multi_label: bool = True, examples: list[dict[str, Any]] | None = None, prompt: str | list[str] | None = None, + adapter_id: str | None = None, ) -> dict[str, Any] | list[dict[str, Any]]: """Async prediction. Concurrent calls accumulate into one batch.""" import asyncio @@ -492,6 +638,7 @@ async def predict_async( multi_label, examples, prompt, + adapter_id, ) for t in items ] From e594485f5b76f48147479db7fba68c383364cd0b Mon Sep 17 00:00:00 2001 From: Ingvar Date: Mon, 1 Jun 2026 08:20:40 +0300 Subject: [PATCH 2/3] update polylora naming --- gliclass/serve/config.py | 47 ++++++++++++++++------- gliclass/serve/server.py | 82 ++++++++++++++++++++-------------------- 2 files changed, 75 insertions(+), 54 deletions(-) diff --git a/gliclass/serve/config.py b/gliclass/serve/config.py index 836f20b..1468725 100644 --- a/gliclass/serve/config.py +++ b/gliclass/serve/config.py @@ -60,16 +60,16 @@ class GLiClassServeConfig: ray_address: str | None = None - enable_mlora: bool = False - mlora_adapter_weight_modules: list[str] | None = None - mlora_max_rank: int = 16 - mlora_max_gpu_adapters: int = 8 - mlora_max_cpu_adapters: int | None = 128 - mlora_disk_cache_dir: str | None = None - mlora_max_disk_adapters: int | None = None - mlora_base_adapter_id: str = "__base__" - mlora_use_triton_kernels: bool = True - mlora_adapter_id_pattern: str = r"^[A-Za-z0-9_.-]{1,128}$" + enable_polylora: bool = False + polylora_adapter_weight_modules: list[str] | None = None + polylora_max_rank: int = 16 + polylora_max_gpu_adapters: int = 8 + polylora_max_cpu_adapters: int | None = 128 + polylora_disk_cache_dir: str | None = None + polylora_max_disk_adapters: int | None = None + polylora_base_adapter_id: str = "__base__" + polylora_use_triton_kernels: bool = True + polylora_adapter_id_pattern: str = r"^[A-Za-z0-9_.-]{1,128}$" def __post_init__(self): if self.max_batch_size not in self.precompiled_batch_sizes: @@ -96,13 +96,34 @@ def from_yaml(cls, config_path: str | Path) -> "GLiClassServeConfig": config_path = Path(config_path) with config_path.open("r") as f: config_dict = yaml.safe_load(f) or {} + legacy_polylora_keys = { + "enable_mlora": "enable_polylora", + "mlora_adapter_weight_modules": "polylora_adapter_weight_modules", + "mlora_max_rank": "polylora_max_rank", + "mlora_max_gpu_adapters": "polylora_max_gpu_adapters", + "mlora_max_cpu_adapters": "polylora_max_cpu_adapters", + "mlora_disk_cache_dir": "polylora_disk_cache_dir", + "mlora_max_disk_adapters": "polylora_max_disk_adapters", + "mlora_base_adapter_id": "polylora_base_adapter_id", + "mlora_use_triton_kernels": "polylora_use_triton_kernels", + "mlora_adapter_id_pattern": "polylora_adapter_id_pattern", + "mlora_target_modules": "polylora_target_modules", + } + for old_key, new_key in legacy_polylora_keys.items(): + if old_key in config_dict and new_key not in config_dict: + config_dict[new_key] = config_dict.pop(old_key) + else: + config_dict.pop(old_key, None) config_dict.pop("mlora_lazy_load_adapters", None) config_dict.pop("mlora_allow_runtime_adapter_loading", None) config_dict.pop("mlora_allow_unsafe_bin_adapters", None) - if "mlora_target_modules" in config_dict and "mlora_adapter_weight_modules" not in config_dict: - config_dict["mlora_adapter_weight_modules"] = config_dict.pop("mlora_target_modules") + config_dict.pop("polylora_lazy_load_adapters", None) + config_dict.pop("polylora_allow_runtime_adapter_loading", None) + config_dict.pop("polylora_allow_unsafe_bin_adapters", None) + if "polylora_target_modules" in config_dict and "polylora_adapter_weight_modules" not in config_dict: + config_dict["polylora_adapter_weight_modules"] = config_dict.pop("polylora_target_modules") else: - config_dict.pop("mlora_target_modules", None) + config_dict.pop("polylora_target_modules", None) return cls(**config_dict) def to_yaml(self, config_path: str | Path) -> None: diff --git a/gliclass/serve/server.py b/gliclass/serve/server.py index 18bf272..63e0fd6 100644 --- a/gliclass/serve/server.py +++ b/gliclass/serve/server.py @@ -31,8 +31,8 @@ def __init__(self, config: GLiClassServeConfig): config: Server configuration with model and serving parameters """ self.config = config - self._mlora_model = None - self._adapter_id_re = re.compile(config.mlora_adapter_id_pattern) + self._polylora_model = None + self._adapter_id_re = re.compile(config.polylora_adapter_id_pattern) env_vars = config.to_env_vars() for key, value in env_vars.items(): @@ -66,8 +66,8 @@ def __init__(self, config: GLiClassServeConfig): logger.info("Loading model: %s", config.model) self.model = GLiClassModel.from_pretrained(config.model) - if config.enable_mlora: - self._initialize_mlora() + if config.enable_polylora: + self._initialize_polylora() self.model.config.max_labels_alloc = config.max_labels_alloc self.model.to(device=self.device, dtype=self.torch_dtype) self.model.eval() @@ -97,72 +97,72 @@ def __init__(self, config: GLiClassServeConfig): logger.info("GLiClass server initialized successfully") - def _initialize_mlora(self) -> None: + def _initialize_polylora(self) -> None: try: - from mlora import CustomLoraConfig, CustomPeftModel + from polylora import PolyLoraConfig, PolyLoraModel except ImportError as exc: - raise ImportError("enable_mlora=True requires the mlora package to be importable") from exc - - target_model = self._get_mlora_target_model() - mlora_config = CustomLoraConfig( - max_gpu_adapters=self.config.mlora_max_gpu_adapters, - max_cpu_adapters=self.config.mlora_max_cpu_adapters, - disk_cache_dir=self.config.mlora_disk_cache_dir, - max_disk_adapters=self.config.mlora_max_disk_adapters, - max_rank=self.config.mlora_max_rank, - target_modules=self.config.mlora_adapter_weight_modules, - base_adapter_id=self.config.mlora_base_adapter_id, - use_triton_kernels=self.config.mlora_use_triton_kernels, + raise ImportError("enable_polylora=True requires the polylora package to be importable") from exc + + target_model = self._get_polylora_target_model() + polylora_config = PolyLoraConfig( + max_gpu_adapters=self.config.polylora_max_gpu_adapters, + max_cpu_adapters=self.config.polylora_max_cpu_adapters, + disk_cache_dir=self.config.polylora_disk_cache_dir, + max_disk_adapters=self.config.polylora_max_disk_adapters, + max_rank=self.config.polylora_max_rank, + target_modules=self.config.polylora_adapter_weight_modules, + base_adapter_id=self.config.polylora_base_adapter_id, + use_triton_kernels=self.config.polylora_use_triton_kernels, ) - self._mlora_model = CustomPeftModel(target_model, mlora_config) - self._set_mlora_target_model(self._mlora_model) + self._polylora_model = PolyLoraModel(target_model, polylora_config) + self._set_polylora_target_model(self._polylora_model) logger.info( - "mLoRA enabled with %d GPU slots, max rank %d, disk cache %s", - self.config.mlora_max_gpu_adapters, - self.config.mlora_max_rank, - self.config.mlora_disk_cache_dir or "disabled", + "PolyLoRA enabled with %d GPU slots, max rank %d, disk cache %s", + self.config.polylora_max_gpu_adapters, + self.config.polylora_max_rank, + self.config.polylora_disk_cache_dir or "disabled", ) - def _get_mlora_target_model(self): + def _get_polylora_target_model(self): architecture = self.model.config.architecture_type if architecture in {"uni-encoder", "bi-encoder", "bi-encoder-fused"}: return self.model.model.encoder_model if architecture in {"encoder-decoder", "encoder-decoder-cls"}: return self.model.model.encoder_decoder_model - raise NotImplementedError(f"mLoRA is not implemented for architecture {architecture!r}") + raise NotImplementedError(f"PolyLoRA is not implemented for architecture {architecture!r}") - def _set_mlora_target_model(self, wrapped_model) -> None: + def _set_polylora_target_model(self, wrapped_model) -> None: architecture = self.model.config.architecture_type if architecture in {"uni-encoder", "bi-encoder", "bi-encoder-fused"}: self.model.model.encoder_model = wrapped_model elif architecture in {"encoder-decoder", "encoder-decoder-cls"}: self.model.model.encoder_decoder_model = wrapped_model else: - raise NotImplementedError(f"mLoRA is not implemented for architecture {architecture!r}") + raise NotImplementedError(f"PolyLoRA is not implemented for architecture {architecture!r}") def _validate_adapter_id(self, adapter_id: str) -> None: if not isinstance(adapter_id, str) or not self._adapter_id_re.fullmatch(adapter_id): - raise ValueError("adapter_id must match mlora_adapter_id_pattern") - if adapter_id == self.config.mlora_base_adapter_id: + raise ValueError("adapter_id must match polylora_adapter_id_pattern") + if adapter_id == self.config.polylora_base_adapter_id: raise ValueError(f"{adapter_id!r} is reserved for base-only inference") def adapter_cache_status(self, adapter_id: str | None = None) -> dict[str, Any]: - if not self.config.enable_mlora or self._mlora_model is None: - return {"enabled": False, "base_adapter_id": self.config.mlora_base_adapter_id} - store = self._mlora_model.adapter_store + if not self.config.enable_polylora or self._polylora_model is None: + return {"enabled": False, "base_adapter_id": self.config.polylora_base_adapter_id} + store = self._polylora_model.adapter_store disk_cache = getattr(store, "disk_cache", None) response: dict[str, Any] = { "enabled": True, - "base_adapter_id": self.config.mlora_base_adapter_id, + "base_adapter_id": self.config.polylora_base_adapter_id, "loaded": sorted(store.adapters.keys()), "disk_cached": sorted(disk_cache.entries.keys()) if disk_cache is not None else [], "disk_cache_dir": str(disk_cache.cache_dir) if disk_cache is not None else None, "max_disk_adapters": disk_cache.max_adapters if disk_cache is not None else None, - "gpu_slots": list(self._mlora_model.adapter_cache.slot_to_adapter), + "gpu_slots": list(self._polylora_model.adapter_cache.slot_to_adapter), } if adapter_id is not None: - if adapter_id == self.config.mlora_base_adapter_id: + if adapter_id == self.config.polylora_base_adapter_id: response["adapter_id"] = adapter_id response["cached"] = True response["cpu_resident"] = False @@ -172,18 +172,18 @@ def adapter_cache_status(self, adapter_id: str | None = None) -> dict[str, Any]: response["adapter_id"] = adapter_id response["cached"] = disk_cache is not None and adapter_id in disk_cache response["cpu_resident"] = adapter_id in store.adapters - response["gpu_resident"] = adapter_id in self._mlora_model.adapter_cache.adapter_to_slot + response["gpu_resident"] = adapter_id in self._polylora_model.adapter_cache.adapter_to_slot return response def ensure_adapter_loaded(self, adapter_id: str | None) -> str | None: if adapter_id is None: - return self.config.mlora_base_adapter_id if self.config.enable_mlora else None - if adapter_id == self.config.mlora_base_adapter_id: + return self.config.polylora_base_adapter_id if self.config.enable_polylora else None + if adapter_id == self.config.polylora_base_adapter_id: return adapter_id - if not self.config.enable_mlora or self._mlora_model is None: + if not self.config.enable_polylora or self._polylora_model is None: raise KeyError(f"Unknown LoRA adapter id: {adapter_id}") self._validate_adapter_id(adapter_id) - if adapter_id in self._mlora_model.adapter_store: + if adapter_id in self._polylora_model.adapter_store: return adapter_id raise KeyError(f"Unknown LoRA adapter id: {adapter_id}") From 4e0ef1f99eb04c07220483eb2f57c86ea92b20d8 Mon Sep 17 00:00:00 2001 From: Ingvar Date: Tue, 2 Jun 2026 10:40:15 +0300 Subject: [PATCH 3/3] fix ruff errors --- gliclass/serve/__main__.py | 336 +++++++++++++++++++++++++++---------- gliclass/serve/server.py | 30 +++- 2 files changed, 265 insertions(+), 101 deletions(-) diff --git a/gliclass/serve/__main__.py b/gliclass/serve/__main__.py index 133ce5c..1e4c0c0 100644 --- a/gliclass/serve/__main__.py +++ b/gliclass/serve/__main__.py @@ -18,14 +18,20 @@ logger = logging.getLogger(__name__) -def main(): - """Main entry point for GLiClass serving.""" +def _parse_int_list(value: str) -> list[int]: + return [int(item.strip()) for item in value.split(",") if item.strip()] + + +def _parse_str_list(value: str) -> list[str]: + return [item.strip() for item in value.split(",") if item.strip()] + + +def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( description="GLiClass Ray Serve deployment", formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) - # Config file parser.add_argument( "--config", type=str, @@ -33,138 +39,248 @@ def main(): 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( + model_group = parser.add_argument_group("Model Configuration") + model_group.add_argument("--model", type=str, default=None, help="Model name or path") + model_group.add_argument("--device", type=str, default=None, help="Device to run on (cuda or cpu)") + model_group.add_argument( "--dtype", type=str, default=None, - choices=["float32", "float16", "bfloat16"], + choices=["float32", "float16", "fp16", "bfloat16", "bf16"], 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, + model_group.add_argument( + "--quantization", + type=str, default=None, - help="Maximum batch size", + help="Model quantization mode, if supported by the GLiClass model", ) - parser.add_argument( + + limits_group = parser.add_argument_group("Model Limits") + limits_group.add_argument("--max-model-len", type=int, default=None, help="Maximum sequence length") + limits_group.add_argument("--max-batch-size", type=int, default=None, help="Maximum batch size") + limits_group.add_argument( "--max-labels", type=int, default=None, help="Maximum number of labels (-1 for unlimited)", ) - parser.add_argument( + limits_group.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", + threshold_group = parser.add_argument_group("Thresholds") + threshold_group.add_argument( + "--default-threshold", + type=float, + default=None, + help="Default confidence threshold", ) - parser.add_argument( - "--port", + + replica_group = parser.add_argument_group("Replica Configuration") + replica_group.add_argument("--num-replicas", type=int, default=None, help="Number of model replicas") + replica_group.add_argument( + "--num-gpus-per-replica", + type=float, + default=None, + help="Number of GPUs per replica", + ) + replica_group.add_argument( + "--num-cpus-per-replica", + type=float, + default=None, + help="Number of CPUs per replica", + ) + + batch_group = parser.add_argument_group("Batching Configuration") + batch_group.add_argument( + "--batch-wait-timeout-ms", + type=float, + default=None, + help="Batch wait timeout in milliseconds", + ) + batch_group.add_argument( + "--request-timeout-s", + type=float, + default=None, + help="Request timeout in seconds", + ) + batch_group.add_argument( + "--max-ongoing-requests", type=int, default=None, - help="Port to bind to", + help="Maximum number of ongoing requests", ) - parser.add_argument( - "--route-prefix", + batch_group.add_argument( + "--queue-capacity", + type=int, + default=None, + help="Request queue capacity", + ) + batch_group.add_argument( + "--precompiled-batch-sizes", + type=_parse_int_list, + default=None, + help="Comma-separated list of batch sizes to precompile", + ) + + server_group = parser.add_argument_group("Server Configuration") + server_group.add_argument("--host", type=str, default="0.0.0.0", help="Host to bind to") + server_group.add_argument("--port", type=int, default=None, help="HTTP port for Ray Serve") + server_group.add_argument("--route-prefix", type=str, default=None, help="HTTP route prefix") + server_group.add_argument( + "--ray-address", type=str, default=None, - help="HTTP route prefix", + help="Ray cluster address (default: local)", ) - parser.add_argument( - "--num-replicas", + + perf_group = parser.add_argument_group("Performance Options") + perf_group.add_argument( + "--tokenizer-threads", type=int, default=None, - help="Number of model replicas", + help="Number of tokenizer threads", ) - - # Performance configuration - parser.add_argument( - "--calibrate-on-startup", - action="store_true", + perf_group.add_argument( + "--enable-compilation", + action=argparse.BooleanOptionalAction, default=None, - help="Run memory calibration on startup", + help="Enable torch.compile", ) - parser.add_argument( + perf_group.add_argument( "--precompile-on-startup", action="store_true", default=None, - help="Precompile model on startup", + help="Run warmup inference for configured precompiled batch sizes", ) - parser.add_argument( - "--use-memory-aware-batching", + perf_group.add_argument( + "--calibrate-on-startup", action="store_true", default=None, - help="Use memory-aware dynamic batching", + help="Run memory calibration on startup", ) - parser.add_argument( - "--enable-compilation", + perf_group.add_argument( + "--use-memory-aware-batching", action="store_true", default=None, - help="Enable torch.compile", + help="Use memory-aware dynamic batching", ) - parser.add_argument( - "--tokenizer-threads", + perf_group.add_argument( + "--warmup-iterations", type=int, default=None, - help="Number of tokenizer threads", + help="Number of warmup iterations per batch size", ) - # Calibration configuration - parser.add_argument( + memory_group = parser.add_argument_group("Memory Configuration") + memory_group.add_argument( + "--target-memory-fraction", + type=float, + default=None, + help="Target GPU memory fraction (0.0-1.0)", + ) + memory_group.add_argument( + "--memory-overhead-factor", + type=float, + default=None, + help="Memory overhead factor for safety margin", + ) + memory_group.add_argument( "--calibration-min-batch-size", type=int, default=None, help="Minimum batch size for calibration", ) - parser.add_argument( + memory_group.add_argument( "--calibration-max-batch-size", type=int, default=None, help="Maximum batch size for calibration", ) - parser.add_argument( + memory_group.add_argument( "--calibration-min-seq-len", type=int, default=None, help="Minimum sequence length for calibration", ) + memory_group.add_argument( + "--calibration-probe-batch-size", + type=int, + default=None, + help="Probe batch size for memory calibration", + ) - args = parser.parse_args() + polylora_group = parser.add_argument_group("PolyLoRA Configuration") + polylora_group.add_argument( + "--enable-polylora", + action="store_true", + default=None, + help="Enable PolyLoRA adapter serving", + ) + polylora_group.add_argument( + "--polylora-adapter-weight-modules", + type=_parse_str_list, + default=None, + help="Comma-separated target module names for adapter weights", + ) + polylora_group.add_argument("--polylora-max-rank", type=int, default=None, help="Maximum LoRA rank") + polylora_group.add_argument( + "--polylora-max-gpu-adapters", + type=int, + default=None, + help="Maximum PolyLoRA GPU adapter slots", + ) + polylora_group.add_argument( + "--polylora-max-cpu-adapters", + type=int, + default=None, + help="Maximum PolyLoRA CPU adapters", + ) + polylora_group.add_argument( + "--polylora-disk-cache-dir", + type=str, + default=None, + help="PolyLoRA disk cache directory", + ) + polylora_group.add_argument( + "--polylora-max-disk-adapters", + type=int, + default=None, + help="Maximum PolyLoRA disk-cached adapters", + ) + polylora_group.add_argument( + "--polylora-base-adapter-id", + type=str, + default=None, + help="Adapter id reserved for base-only inference", + ) + polylora_group.add_argument( + "--polylora-use-triton-kernels", + action=argparse.BooleanOptionalAction, + default=None, + help="Use Triton kernels for PolyLoRA", + ) + polylora_group.add_argument( + "--polylora-adapter-id-pattern", + type=str, + default=None, + help="Regular expression for valid adapter ids", + ) + + return parser + +def config_from_args(args: argparse.Namespace) -> GLiClassServeConfig: if args.config: - logger.info(f"Loading config from: {args.config}") + logger.info("Loading config from: %s", 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) @@ -173,48 +289,84 @@ def main(): "model": args.model, "device": args.device, "dtype": args.dtype, + "quantization": args.quantization, "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, + "default_threshold": args.default_threshold, + "num_replicas": args.num_replicas, + "num_gpus_per_replica": args.num_gpus_per_replica, + "num_cpus_per_replica": args.num_cpus_per_replica, + "batch_wait_timeout_ms": args.batch_wait_timeout_ms, + "request_timeout_s": args.request_timeout_s, + "max_ongoing_requests": args.max_ongoing_requests, + "queue_capacity": args.queue_capacity, + "precompiled_batch_sizes": args.precompiled_batch_sizes, "http_port": args.port, "route_prefix": args.route_prefix, - "num_replicas": args.num_replicas, - "calibrate_on_startup": args.calibrate_on_startup, + "ray_address": args.ray_address, + "tokenizer_threads": args.tokenizer_threads, + "enable_compilation": args.enable_compilation, "precompile_on_startup": args.precompile_on_startup, + "calibrate_on_startup": args.calibrate_on_startup, "use_memory_aware_batching": args.use_memory_aware_batching, - "enable_compilation": args.enable_compilation, - "tokenizer_threads": args.tokenizer_threads, + "warmup_iterations": args.warmup_iterations, + "target_memory_fraction": args.target_memory_fraction, + "memory_overhead_factor": args.memory_overhead_factor, "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, + "calibration_probe_batch_size": args.calibration_probe_batch_size, + "enable_polylora": args.enable_polylora, + "polylora_adapter_weight_modules": args.polylora_adapter_weight_modules, + "polylora_max_rank": args.polylora_max_rank, + "polylora_max_gpu_adapters": args.polylora_max_gpu_adapters, + "polylora_max_cpu_adapters": args.polylora_max_cpu_adapters, + "polylora_disk_cache_dir": args.polylora_disk_cache_dir, + "polylora_max_disk_adapters": args.polylora_max_disk_adapters, + "polylora_base_adapter_id": args.polylora_base_adapter_id, + "polylora_use_triton_kernels": args.polylora_use_triton_kernels, + "polylora_adapter_id_pattern": args.polylora_adapter_id_pattern, } config.update(**cli_overrides) + config.__post_init__() + return config + + +def main(): + """Main entry point for GLiClass serving.""" + parser = build_parser() + args = parser.parse_args() + config = config_from_args(args) 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(" Model: %s", config.model) + logger.info(" Device: %s", config.device) + logger.info(" Dtype: %s", config.dtype) + logger.info(" Quantization: %s", config.quantization or "disabled") + logger.info(" Max model length: %s", config.max_model_len) + logger.info(" Max batch size: %s", config.max_batch_size) + logger.info(" Precompiled batch sizes: %s", config.precompiled_batch_sizes) + logger.info(" Max labels: %s", config.max_labels) + logger.info(" Max labels alloc: %s", config.max_labels_alloc) + logger.info(" HTTP port: %s", config.http_port) + logger.info(" Route prefix: %s", config.route_prefix) + logger.info(" Num replicas: %s", config.num_replicas) + logger.info(" GPUs per replica: %s", config.num_gpus_per_replica) + logger.info(" CPUs per replica: %s", config.num_cpus_per_replica) + logger.info(" Compile: %s", config.enable_compilation) + logger.info(" Precompile on startup: %s", config.precompile_on_startup) + logger.info(" Calibrate on startup: %s", config.calibrate_on_startup) + logger.info(" Memory-aware batching: %s", config.use_memory_aware_batching) + logger.info(" PolyLoRA: %s", config.enable_polylora) logger.info("=" * 60) - logger.info("Initializing Ray...") - ray.init(ignore_reinit_error=True) - - logger.info(f"Deploying GLiClass with model: {config.model}") - _app = serve_gliclass(config, blocking=False, host=args.host) # Keep reference to prevent GC + logger.info("Deploying GLiClass with model: %s", config.model) + _app = serve_gliclass(config, blocking=False, host=args.host) - logger.info(f"GLiClass server running at http://{args.host}:{config.http_port}{config.route_prefix}") + logger.info("GLiClass server running at http://%s:%d%s", args.host, config.http_port, config.route_prefix) logger.info("Press Ctrl+C to stop the server") def signal_handler(_sig, _frame): diff --git a/gliclass/serve/server.py b/gliclass/serve/server.py index 63e0fd6..375a1f6 100644 --- a/gliclass/serve/server.py +++ b/gliclass/serve/server.py @@ -7,10 +7,10 @@ import torch from ray import serve +from transformers import AutoTokenizer from ray.serve._private import api as serve_private_api -from ray.serve._private.build_app import build_app from starlette.responses import JSONResponse -from transformers import AutoTokenizer +from ray.serve._private.build_app import build_app from gliclass.model import GLiClassModel from gliclass.pipeline import ZeroShotClassificationPipeline @@ -72,6 +72,12 @@ def __init__(self, config: GLiClassServeConfig): self.model.to(device=self.device, dtype=self.torch_dtype) self.model.eval() + if config.quantization: + if not hasattr(self.model, "quantize"): + raise ValueError("quantization is configured, but this GLiClass model does not support quantize()") + logger.info("Applying quantization: %s", config.quantization) + self.model.quantize(config.quantization) + self.tokenizer = AutoTokenizer.from_pretrained(config.model) pipeline_kwargs = { "model": self.model, @@ -89,17 +95,21 @@ def __init__(self, config: GLiClassServeConfig): if torch.cuda.is_available(): self.memory_estimator.measure_model_memory() - if config.enable_compilation: + if config.enable_compilation and hasattr(self.model, "compile"): + logger.info("Compiling GLiClass model") + self.model.compile() + + if config.precompile_on_startup: self._precompile() - if torch.cuda.is_available(): + if config.calibrate_on_startup and torch.cuda.is_available(): self._calibrate_memory() logger.info("GLiClass server initialized successfully") def _initialize_polylora(self) -> None: try: - from polylora import PolyLoraConfig, PolyLoraModel + from polylora import PolyLoraModel, PolyLoraConfig except ImportError as exc: raise ImportError("enable_polylora=True requires the polylora package to be importable") from exc @@ -190,9 +200,6 @@ def ensure_adapter_loaded(self, adapter_id: str | None) -> str | None: def _precompile(self) -> None: logger.info("Precompiling model for batch sizes: %s", self.config.precompiled_batch_sizes) - if hasattr(self.model, "compile"): - self.model.compile() - dummy_labels = ["person", "organization", "location"] for batch_size in self.config.precompiled_batch_sizes: @@ -233,7 +240,11 @@ def batch_size_fn(self, seq_len: int | None = None) -> int: Returns: Optimal batch size from precompiled sizes """ - if not torch.cuda.is_available(): + if not self.config.use_memory_aware_batching or not torch.cuda.is_available(): + return self.config.precompiled_batch_sizes[-1] + + if not self.memory_estimator.per_sample_table: + logger.warning("Memory-aware batching is enabled, but memory estimator is not calibrated") return self.config.precompiled_batch_sizes[-1] if seq_len is None: @@ -365,6 +376,7 @@ def _build_deployment(config: GLiClassServeConfig): "num_cpus": config.num_cpus_per_replica, }, max_ongoing_requests=config.max_ongoing_requests, + max_queued_requests=config.queue_capacity, ) class GLiClassDeployment: def __init__(self, serve_config: GLiClassServeConfig):