Brunobkr/llama.cpp_AlgMor24_github

1#!/usr/bin/env python32# -*- coding: utf-8 -*-3 4# type: ignore[reportUnusedImport]5 6import subprocess7import os8 9TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tmp")10import re11import json12from json import JSONDecodeError13import sys14import requests15import time16from concurrent.futures import ThreadPoolExecutor, as_completed17from typing import (18    Any,19    Callable,20    ContextManager,21    Iterable,22    Iterator,23    List,24    Literal,25    Tuple,26    Set,27)28from re import RegexFlag29import wget30 31 32DEFAULT_HTTP_TIMEOUT = 6033 34# per-request timeout, a hung server fails the test instead of stalling the CI for hours35DEFAULT_REQUEST_TIMEOUT = 60036 37 38class ServerResponse:39    headers: dict40    status_code: int41    body: dict | Any42 43 44class ServerError(Exception):45    def __init__(self, code, body):46        self.code = code47        self.body = body48 49 50class ServerProcess:51    # default options52    debug: bool = False53    server_port: int = 808054    server_host: str = "127.0.0.1"55    model_hf_repo: str | None = "ggml-org/models"56    model_hf_file: str | None = "tinyllamas/stories260K.gguf"57    model_alias: str = "tinyllama-2"58    temperature: float = 0.859    seed: int = 4260    offline: bool = False61 62    # custom options63    model_alias: str | None = None64    model_tags: str | None = None65    model_url: str | None = None66    model_file: str | None = None67    model_draft: str | None = None68    n_threads: int | None = None69    n_gpu_layer: int | None = None70    n_batch: int | None = None71    n_ubatch: int | None = None72    n_ctx: int | None = None73    n_ga: int | None = None74    n_ga_w: int | None = None75    n_predict: int | None = None76    n_prompts: int | None = 077    slot_save_path: str | None = None78    id_slot: int | None = None79    cache_prompt: bool | None = None80    n_slots: int | None = None81    ctk: str | None = None82    ctv: str | None = None83    fa: str | None = None84    server_continuous_batching: bool | None = False85    server_embeddings: bool | None = False86    server_reranking: bool | None = False87    server_metrics: bool | None = False88    kv_unified: bool | None = False89    server_slots: bool | None = False90    pooling: str | None = None91    api_key: str | None = None92    models_dir: str | None = None93    models_max: int | None = None94    models_preset: str | None = None95    no_models_autoload: bool | None = None96    lora_files: List[str] | None = None97    enable_ctx_shift: int | None = False98    spec_type: str | None = None99    spec_draft_n_min: int | None = None100    spec_draft_n_max: int | None = None101    no_ui: bool | None = None102    jinja: bool | None = None103    reasoning_format: Literal['deepseek', 'none', 'nothink'] | None = None104    reasoning: Literal['on', 'off', 'auto'] | None = None105    chat_template: str | None = None106    chat_template_file: str | None = None107    server_path: str | None = None108    mmproj_url: str | None = None109    media_path: str | None = None110    sleep_idle_seconds: int | None = None111    cache_ram: int | None = None112    no_cache_idle_slots: bool = False113    log_path: str | None = None114    ui_mcp_proxy: bool = False115    backend_sampling: bool = False116    gcp_compat: bool = False117    server_tools: str | None = None118    server_tools_runtime: str | None = None119    mcp_servers_config: str | None = None120    mcp_servers_json: str | None = None121    cors_origins: str | None = None122 123    # session variables124    process: subprocess.Popen | None = None125 126    def __init__(self):127        if "N_GPU_LAYERS" in os.environ:128            self.n_gpu_layer = int(os.environ["N_GPU_LAYERS"])129        if "DEBUG" in os.environ:130            self.debug = True131        if "PORT" in os.environ:132            self.server_port = int(os.environ["PORT"])133        self.external_server = "DEBUG_EXTERNAL" in os.environ134 135    def start(self, timeout_seconds: int = DEFAULT_HTTP_TIMEOUT) -> None:136        env = {137            **os.environ,138            "LLAMA_SERVER_DEBUG_FAKE_TIMING": "1",139        }140        if "LLAMA_CACHE" not in os.environ:141            env["LLAMA_CACHE"] = "tmp"142        if self.external_server:143            print(f"[external_server]: Assuming external server running on {self.server_host}:{self.server_port}")144            return145        if self.server_path is not None:146            server_path = self.server_path147        elif "LLAMA_SERVER_BIN_PATH" in os.environ:148            server_path = os.environ["LLAMA_SERVER_BIN_PATH"]149        elif os.name == "nt":150            server_path = "../../../build/bin/Release/llama-server.exe"151        else:152            server_path = "../../../build/bin/llama-server"153        server_args = [154            "--host",155            self.server_host,156            "--port",157            self.server_port,158            "--temp",159            self.temperature,160            "--seed",161            self.seed,162        ]163        if self.offline:164            server_args.append("--offline")165        if self.model_file:166            server_args.extend(["--model", self.model_file])167        if self.model_url:168            server_args.extend(["--model-url", self.model_url])169        if self.model_draft:170            server_args.extend(["--model-draft", self.model_draft])171        if self.model_hf_repo:172            server_args.extend(["--hf-repo", self.model_hf_repo])173        if self.model_hf_file:174            server_args.extend(["--hf-file", self.model_hf_file])175        if self.models_dir:176            server_args.extend(["--models-dir", self.models_dir])177        if self.models_max is not None:178            server_args.extend(["--models-max", self.models_max])179        if self.models_preset:180            server_args.extend(["--models-preset", self.models_preset])181        if self.cors_origins:182            server_args.extend(["--cors-origins", self.cors_origins])183        if self.n_batch:184            server_args.extend(["--batch-size", self.n_batch])185        if self.n_ubatch:186            server_args.extend(["--ubatch-size", self.n_ubatch])187        if self.n_threads:188            server_args.extend(["--threads", self.n_threads])189        if self.n_gpu_layer:190            server_args.extend(["--n-gpu-layers", self.n_gpu_layer])191        if self.server_continuous_batching:192            server_args.append("--cont-batching")193        if self.server_embeddings:194            server_args.append("--embedding")195        if self.server_reranking:196            server_args.append("--reranking")197        if self.server_metrics:198            server_args.append("--metrics")199        if self.kv_unified:200            server_args.append("--kv-unified")201        if self.server_slots:202            server_args.append("--slots")203        else:204            server_args.append("--no-slots")205        if self.pooling:206            server_args.extend(["--pooling", self.pooling])207        if self.model_alias:208            server_args.extend(["--alias", self.model_alias])209        if self.model_tags:210            server_args.extend(["--tags", self.model_tags])211        if self.n_ctx:212            server_args.extend(["--ctx-size", self.n_ctx])213        if self.n_slots:214            server_args.extend(["--parallel", self.n_slots])215        if self.ctk:216            server_args.extend(["-ctk", self.ctk])217        if self.ctv:218            server_args.extend(["-ctv", self.ctv])219        if self.fa is not None:220            server_args.extend(["-fa", self.fa])221        if self.n_predict:222            server_args.extend(["--n-predict", self.n_predict])223        if self.slot_save_path:224            server_args.extend(["--slot-save-path", self.slot_save_path])225        if self.n_ga:226            server_args.extend(["--grp-attn-n", self.n_ga])227        if self.n_ga_w:228            server_args.extend(["--grp-attn-w", self.n_ga_w])229        if self.debug:230            server_args.append("--verbose")231        if self.lora_files:232            for lora_file in self.lora_files:233                server_args.extend(["--lora", lora_file])234        if self.enable_ctx_shift:235            server_args.append("--context-shift")236        if self.spec_type:237            server_args.extend(["--spec-type", self.spec_type])238        if self.api_key:239            server_args.extend(["--api-key", self.api_key])240        if self.spec_draft_n_max:241            server_args.extend(["--spec-draft-n-max", self.spec_draft_n_max])242        if self.spec_draft_n_min:243            server_args.extend(["--spec-draft-n-min", self.spec_draft_n_min])244        if self.no_ui:245            server_args.append("--no-ui")246        if self.no_models_autoload:247            server_args.append("--no-models-autoload")248        if self.jinja:249            server_args.append("--jinja")250        else:251            server_args.append("--no-jinja")252        if self.reasoning_format is not None:253            server_args.extend(("--reasoning-format", self.reasoning_format))254        if self.reasoning is not None:255            server_args.extend(("--reasoning", self.reasoning))256        if self.chat_template:257            server_args.extend(["--chat-template", self.chat_template])258        if self.chat_template_file:259            server_args.extend(["--chat-template-file", self.chat_template_file])260        if self.mmproj_url:261            server_args.extend(["--mmproj-url", self.mmproj_url])262        if self.media_path:263            server_args.extend(["--media-path", self.media_path])264        if self.sleep_idle_seconds is not None:265            server_args.extend(["--sleep-idle-seconds", self.sleep_idle_seconds])266        if self.cache_ram is not None:267            server_args.extend(["--cache-ram", self.cache_ram])268        if self.no_cache_idle_slots:269            server_args.append("--no-cache-idle-slots")270        if self.ui_mcp_proxy:271            server_args.append("--ui-mcp-proxy")272        if self.server_tools:273            server_args.extend(["--tools", self.server_tools])274        if self.server_tools_runtime:275            server_args.extend(["--tools-runtime", self.server_tools_runtime])276        if self.mcp_servers_config:277            server_args.extend(["--mcp-servers-config", self.mcp_servers_config])278        if self.mcp_servers_json:279            server_args.extend(["--mcp-servers-json", self.mcp_servers_json])280        if self.backend_sampling:281            server_args.append("--backend_sampling")282        if self.gcp_compat:283            env["AIP_MODE"] = "PREDICTION"284 285        args = [str(arg) for arg in [server_path, *server_args]]286        print(f"tests: starting server with: {' '.join(args)}")287 288        flags = 0289        if "nt" == os.name:290            flags |= subprocess.DETACHED_PROCESS291            flags |= subprocess.CREATE_NEW_PROCESS_GROUP292            flags |= subprocess.CREATE_NO_WINDOW293 294        if self.log_path:295            self._log = open(self.log_path, "w")296        else:297            self._log = sys.stdout298 299        self.process = subprocess.Popen(300            [str(arg) for arg in [server_path, *server_args]],301            creationflags=flags,302            stdout=self._log,303            stderr=self._log if self._log != sys.stdout else sys.stdout,304            env=env,305        )306        server_instances.add(self)307 308        print(f"server pid={self.process.pid}, pytest pid={os.getpid()}")309 310        # wait for server to start311        start_time = time.time()312        last_print_time = start_time313        while time.time() - start_time < timeout_seconds:314            try:315                response = self.make_request("GET", "/health", headers={316                    "Authorization": f"Bearer {self.api_key}" if self.api_key else None317                })318                if response.status_code == 200:319                    self.ready = True320                    return  # server is ready321            except Exception as e:322                pass323            # Check if process died324            if self.process.poll() is not None:325                raise RuntimeError(f"Server process died with return code {self.process.returncode}")326 327            if time.time() - last_print_time >= 1.0:328                print(f"Waiting for server to start...")329                last_print_time = time.time()330            time.sleep(0.01)331        raise TimeoutError(f"Server did not start within {timeout_seconds} seconds")332 333    def stop(self) -> None:334        if self.external_server:335            print("[external_server]: Not stopping external server")336            return337        if self in server_instances:338            server_instances.remove(self)339        if self.process:340            print(f"Stopping server with pid={self.process.pid}")341            self.process.terminate()342            try:343                self.process.wait(timeout=5)344            except subprocess.TimeoutExpired:345                print(f"Server pid={self.process.pid} did not terminate in time, killing")346                self.process.kill()347                self.process.wait(timeout=5)348            except Exception as e:349                print(f"Error waiting for server: {e}")350            self.process = None351        if hasattr(self, '_log') and self._log != sys.stdout:352            self._log.close()353 354    def make_request(355        self,356        method: str,357        path: str,358        data: dict | Any | None = None,359        headers: dict | None = None,360        timeout: float | None = DEFAULT_REQUEST_TIMEOUT,361    ) -> ServerResponse:362        url = f"http://{self.server_host}:{self.server_port}{path}"363        parse_body = False364        if method == "GET":365            response = requests.get(url, headers=headers, timeout=timeout)366            parse_body = True367        elif method == "POST":368            response = requests.post(url, headers=headers, json=data, timeout=timeout)369            parse_body = True370        elif method == "DELETE":371            response = requests.delete(url, headers=headers, timeout=timeout)372            parse_body = True373        elif method == "OPTIONS":374            response = requests.options(url, headers=headers, timeout=timeout)375        else:376            raise ValueError(f"Unimplemented method: {method}")377        result = ServerResponse()378        result.headers = dict(response.headers)379        result.status_code = response.status_code380        if parse_body:381            try:382                result.body = response.json()383            except (JSONDecodeError, requests.exceptions.JSONDecodeError):384                result.body = response.text385        else:386            result.body = None387        print("Response from server", json.dumps(result.body, indent=2))388        return result389 390    def make_stream_request(391        self,392        method: str,393        path: str,394        data: dict | None = None,395        headers: dict | None = None,396    ) -> Iterator[dict]:397        url = f"http://{self.server_host}:{self.server_port}{path}"398        if method == "POST":399            response = requests.post(url, headers=headers, json=data, stream=True)400        else:401            raise ValueError(f"Unimplemented method: {method}")402        if response.status_code != 200:403            raise ServerError(response.status_code, response.json())404        for line_bytes in response.iter_lines():405            line = line_bytes.decode("utf-8")406            if '[DONE]' in line:407                break408            elif line.startswith('data: '):409                data = json.loads(line[6:])410                print("Partial response from server", json.dumps(data, indent=2))411                yield data412 413    def make_any_request(414        self,415        method: str,416        path: str,417        data: dict | None = None,418        headers: dict | None = None,419        timeout: float | None = DEFAULT_REQUEST_TIMEOUT,420    ) -> dict:421        stream = data.get('stream', False)422        if stream:423            content: list[str] = []424            reasoning_content: list[str] = []425            tool_calls: list[dict] = []426            finish_reason: Optional[str] = None427 428            content_parts = 0429            reasoning_content_parts = 0430            tool_call_parts = 0431            arguments_parts = 0432 433            for chunk in self.make_stream_request(method, path, data, headers):434                if chunk['choices']:435                    assert len(chunk['choices']) == 1, f'Expected 1 choice, got {len(chunk["choices"])}'436                    choice = chunk['choices'][0]437                    if choice['delta'].get('content') is not None:438                        assert len(choice['delta']['content']) > 0, f'Expected non empty content delta!'439                        content.append(choice['delta']['content'])440                        content_parts += 1441                    if choice['delta'].get('reasoning_content') is not None:442                        assert len(choice['delta']['reasoning_content']) > 0, f'Expected non empty reasoning_content delta!'443                        reasoning_content.append(choice['delta']['reasoning_content'])444                        reasoning_content_parts += 1445                    if choice['delta'].get('finish_reason') is not None:446                        finish_reason = choice['delta']['finish_reason']447                    for tc in choice['delta'].get('tool_calls', []):448                        if 'function' not in tc:449                            raise ValueError(f"Expected function type, got {tc['type']}")450                        if tc['index'] >= len(tool_calls):451                            assert 'id' in tc452                            assert tc.get('type') == 'function'453                            assert 'function' in tc and 'name' in tc['function'] and len(tc['function']['name']) > 0, \454                                f"Expected function call with name, got {tc.get('function')}"455                            tool_calls.append(dict(456                                id="",457                                type="function",458                                function=dict(459                                    name="",460                                    arguments="",461                                )462                            ))463                        tool_call = tool_calls[tc['index']]464                        if tc.get('id') is not None:465                            tool_call['id'] = tc['id']466                        fct = tc['function']467                        assert 'id' not in fct, f"Function call should not have id: {fct}"468                        if fct.get('name') is not None:469                            tool_call['function']['name'] = tool_call['function'].get('name', '') + fct['name']470                        if fct.get('arguments') is not None:471                            tool_call['function']['arguments'] += fct['arguments']472                            arguments_parts += 1473                        tool_call_parts += 1474                else:475                    # When `include_usage` is True (the default), we expect the last chunk of the stream476                    # immediately preceding the `data: [DONE]` message to contain a `choices` field with an empty array477                    # and a `usage` field containing the usage statistics (n.b., llama-server also returns `timings` in478                    # the last chunk)479                    assert 'usage' in chunk, f"Expected finish_reason in chunk: {chunk}"480                    assert 'timings' in chunk, f"Expected finish_reason in chunk: {chunk}"481            print(f'Streamed response had {content_parts} content parts, {reasoning_content_parts} reasoning_content parts, {tool_call_parts} tool call parts incl. {arguments_parts} arguments parts')482            result = dict(483                choices=[484                    dict(485                        index=0,486                        finish_reason=finish_reason,487                        message=dict(488                            role='assistant',489                            content=''.join(content) if content else None,490                            reasoning_content=''.join(reasoning_content) if reasoning_content else None,491                            tool_calls=tool_calls if tool_calls else None,492                        ),493                    )494                ],495            )496            print("Final response from server", json.dumps(result, indent=2))497            return result498        else:499            response = self.make_request(method, path, data, headers, timeout=timeout)500            assert response.status_code == 200, f"Server returned error: {response.status_code}"501            return response.body502 503 504 505server_instances: Set[ServerProcess] = set()506 507 508class ServerPreset:509    @staticmethod510    def load_all() -> None:511        """ Load all server presets to ensure model files are cached. """512        servers: List[ServerProcess] = [513            method()514            for name, method in ServerPreset.__dict__.items()515            if callable(method) and name != "load_all"516        ]517        for server in servers:518            server.offline = False519            server.start()520            server.stop()521 522    @staticmethod523    def tinyllama2() -> ServerProcess:524        server = ServerProcess()525        server.offline = True # will be downloaded by load_all()526        server.model_hf_repo = "ggml-org/test-model-stories260K"527        server.model_hf_file = None528        server.model_alias = "tinyllama-2"529        server.n_ctx = 512530        server.n_batch = 32531        server.n_slots = 2532        server.n_predict = 64533        server.seed = 42534        return server535 536    @staticmethod537    def bert_bge_small() -> ServerProcess:538        server = ServerProcess()539        server.offline = True # will be downloaded by load_all()540        server.model_hf_repo = "ggml-org/models"541        server.model_hf_file = "bert-bge-small/ggml-model-f16.gguf"542        server.model_alias = "bert-bge-small"543        server.n_ctx = 512544        server.n_batch = 128545        server.n_ubatch = 128546        server.n_slots = 2547        server.seed = 42548        server.server_embeddings = True549        return server550 551    @staticmethod552    def bert_bge_small_with_fa() -> ServerProcess:553        server = ServerProcess()554        server.offline = True # will be downloaded by load_all()555        server.model_hf_repo = "ggml-org/models"556        server.model_hf_file = "bert-bge-small/ggml-model-f16.gguf"557        server.model_alias = "bert-bge-small"558        server.n_ctx = 1024559        server.n_batch = 300560        server.n_ubatch = 300561        server.n_slots = 2562        server.fa = "on"563        server.seed = 42564        server.server_embeddings = True565        return server566 567    @staticmethod568    def tinyllama_infill() -> ServerProcess:569        server = ServerProcess()570        server.offline = True # will be downloaded by load_all()571        server.model_hf_repo = "ggml-org/test-model-stories260K-infill"572        server.model_hf_file = None573        server.model_alias = "tinyllama-infill"574        server.n_ctx = 2048575        server.n_batch = 1024576        server.n_slots = 1577        server.n_predict = 64578        server.temperature = 0.0579        server.seed = 42580        return server581 582    @staticmethod583    def stories15m_moe() -> ServerProcess:584        server = ServerProcess()585        server.offline = True # will be downloaded by load_all()586        server.model_hf_repo = "ggml-org/stories15M_MOE"587        server.model_hf_file = "stories15M_MOE-F16.gguf"588        server.model_alias = "stories15m-moe"589        server.n_ctx = 2048590        server.n_batch = 1024591        server.n_slots = 1592        server.n_predict = 64593        server.temperature = 0.0594        server.seed = 42595        return server596 597    @staticmethod598    def jina_reranker_tiny() -> ServerProcess:599        server = ServerProcess()600        server.offline = True # will be downloaded by load_all()601        server.model_hf_repo = "ggml-org/models"602        server.model_hf_file = "jina-reranker-v1-tiny-en/ggml-model-f16.gguf"603        server.model_alias = "jina-reranker"604        server.n_ctx = 512605        server.n_batch = 512606        server.n_slots = 1607        server.seed = 42608        server.server_reranking = True609        return server610 611    @staticmethod612    def tinygemma3() -> ServerProcess:613        server = ServerProcess()614        server.offline = True # will be downloaded by load_all()615        # mmproj is already provided by HF registry API616        server.model_hf_file = None617        server.model_hf_repo = "ggml-org/tinygemma3-GGUF:Q8_0"618        server.model_alias = "tinygemma3"619        server.n_ctx = 1024620        server.n_batch = 32621        server.n_slots = 2622        server.n_predict = 4623        server.seed = 42624        return server625 626    @staticmethod627    def router() -> ServerProcess:628        server = ServerProcess()629        server.offline = True # will be downloaded by load_all()630        # router server has no models631        server.model_file = None632        server.model_alias = None633        server.model_hf_repo = None634        server.model_hf_file = None635        server.n_ctx = 1024636        server.n_batch = 16637        server.n_slots = 1638        server.n_predict = 16639        server.seed = 42640        return server641 642 643def parallel_function_calls(function_list: List[Tuple[Callable[..., Any], Tuple[Any, ...]]]) -> List[Any]:644    """645    Run multiple functions in parallel and return results in the same order as calls. Equivalent to Promise.all in JS.646 647    Example usage:648 649    results = parallel_function_calls([650        (func1, (arg1, arg2)),651        (func2, (arg3, arg4)),652    ])653    """654    results = [None] * len(function_list)655    exceptions = []656 657    def worker(index, func, args):658        try:659            result = func(*args)660            results[index] = result661        except Exception as e:662            exceptions.append((index, str(e)))663 664    with ThreadPoolExecutor() as executor:665        futures = []666        for i, (func, args) in enumerate(function_list):667            future = executor.submit(worker, i, func, args)668            futures.append(future)669 670        # Wait for all futures to complete671        for future in as_completed(futures):672            pass673 674    # Check if there were any exceptions675    if exceptions:676        print("Exceptions occurred:")677        for index, error in exceptions:678            print(f"Function at index {index}: {error}")679 680    return results681 682 683def match_regex(regex: str, text: str) -> bool:684    return (685        re.compile(686            regex, flags=RegexFlag.IGNORECASE | RegexFlag.MULTILINE | RegexFlag.DOTALL687        ).search(text)688        is not None689    )690 691 692def download_file(url: str, output_file_path: str | None = None) -> str:693    """694    Download a file from a URL to a local path. If the file already exists, it will not be downloaded again.695 696    output_file_path is the local path to save the downloaded file. If not provided, the file will be saved in the root directory.697 698    Returns the local path of the downloaded file.699    """700    file_name = url.split('/').pop()701    output_file = f'./tmp/{file_name}' if output_file_path is None else output_file_path702    if not os.path.exists(output_file):703        print(f"Downloading {url} to {output_file}")704        wget.download(url, out=output_file)705        print(f"Done downloading to {output_file}")706    else:707        print(f"File already exists at {output_file}")708    return output_file709 710 711def is_slow_test_allowed():712    return os.environ.get("SLOW_TESTS") == "1" or os.environ.get("SLOW_TESTS") == "ON"713