Rename handler.py.bak to handler.py

handler.py

@@ -0,0 +1,237 @@
+# handler.py
+# Hugging Face Inference Endpoints - Custom Handler
+#
+# This handler starts an internal SGLang server (OpenAI-compatible) and proxies
+# requests to it. It supports both:
+#  - HF "inputs": str (single prompt)
+#  - HF "inputs": list[{"role": "...", "content": "..."}] (chat style)
+#
+# Expected request body patterns (common in HF endpoints):
+#  - {"inputs": "Hello", "parameters": {"max_new_tokens": 256, "temperature": 0.7}}
+#  - {"inputs": [{"role":"user","content":"Hello"}], "parameters": {...}}
+
+from __future__ import annotations
+
+import json
+import os
+import socket
+import subprocess
+import time
+from typing import Any, Dict, List, Optional, Union
+
+import requests
+
+
+def _is_port_open(host: str, port: int, timeout_s: float = 0.5) -> bool:
+    try:
+        with socket.create_connection((host, port), timeout=timeout_s):
+            return True
+    except OSError:
+        return False
+
+
+def _wait_for_server(host: str, port: int, health_url: str, timeout_s: int = 300) -> None:
+    start = time.time()
+    # 1) Wait for TCP port
+    while time.time() - start < timeout_s:
+        if _is_port_open(host, port):
+            break
+        time.sleep(0.5)
+
+    # 2) Wait for /health (preferred)
+    while time.time() - start < timeout_s:
+        try:
+            r = requests.get(health_url, timeout=2)
+            if r.status_code == 200:
+                return
+        except requests.RequestException:
+            pass
+        time.sleep(0.5)
+
+    raise RuntimeError(
+        f"SGLang server did not become ready within {timeout_s}s "
+        f"(host={host}, port={port}, health={health_url})."
+    )
+
+
+def _coerce_messages(inputs: Any) -> List[Dict[str, str]]:
+    """
+    Convert HF inputs into OpenAI chat messages.
+    """
+    if isinstance(inputs, str):
+        return [{"role": "user", "content": inputs}]
+
+    if isinstance(inputs, list):
+        # Already messages?
+        # We accept list of dicts with role/content, or list of strings.
+        if all(isinstance(x, dict) for x in inputs):
+            msgs: List[Dict[str, str]] = []
+            for m in inputs:
+                role = str(m.get("role", "user"))
+                content = m.get("content", "")
+                if content is None:
+                    content = ""
+                msgs.append({"role": role, "content": str(content)})
+            return msgs
+        if all(isinstance(x, str) for x in inputs):
+            # Treat as a multi-line user prompt
+            return [{"role": "user", "content": "\n".join(inputs)}]
+
+    # Fallback: stringify
+    return [{"role": "user", "content": json.dumps(inputs, ensure_ascii=False)}]
+
+
+def _map_generation_params(hf_params: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Map typical HF params to OpenAI-compatible chat completion params.
+    Keep pass-through for unknown keys where it is safe.
+    """
+    if hf_params is None:
+        hf_params = {}
+
+    # Common HF keys: max_new_tokens, temperature, top_p, repetition_penalty, stop, seed
+    out: Dict[str, Any] = {}
+
+    max_new_tokens = hf_params.get("max_new_tokens", hf_params.get("max_tokens"))
+    if max_new_tokens is not None:
+        out["max_tokens"] = int(max_new_tokens)
+
+    for k in ("temperature", "top_p", "seed"):
+        if k in hf_params and hf_params[k] is not None:
+            out[k] = hf_params[k]
+
+    # HF sometimes uses "stop" (str or list[str])
+    if "stop" in hf_params and hf_params["stop"] is not None:
+        out["stop"] = hf_params["stop"]
+
+    # OpenAI-compatible streaming flag; HF toolkit generally expects non-streaming response
+    if "stream" in hf_params:
+        out["stream"] = bool(hf_params["stream"])
+    else:
+        out["stream"] = False
+
+    # Best-effort pass-through for presence/frequency penalty if provided
+    for k in ("presence_penalty", "frequency_penalty"):
+        if k in hf_params and hf_params[k] is not None:
+            out[k] = hf_params[k]
+
+    return out
+
+
+class EndpointHandler:
+    """
+    Hugging Face Inference Endpoints custom handler:
+      - __init__(model_dir): load/init anything
+      - __call__(data): run inference
+    """
+
+    def __init__(self, model_dir: str, **_: Any) -> None:
+        # HF mounts the repo under model_dir (typically /repository)
+        self.model_dir = model_dir
+
+        # Where SGLang will listen
+        self.host = os.getenv("SGLANG_HOST", "127.0.0.1")
+        self.port = int(os.getenv("SGLANG_PORT", "30000"))
+
+        # Model identifier/path
+        # For Inference Endpoints, weights/artifacts are available under model_dir.
+        # Using local path avoids extra hub downloads.
+        self.model_path = os.getenv("SGLANG_MODEL_PATH", model_dir)
+
+        # Optional: tokenizer path (defaults to model path)
+        self.tokenizer_path = os.getenv("SGLANG_TOKENIZER_PATH", self.model_path)
+
+        # Optional: tensor parallel size, chunked prefill, etc. (SGLang server args)
+        self.tp_size = int(os.getenv("SGLANG_TP_SIZE", "1"))
+        self.chunked_prefill_size = os.getenv("SGLANG_CHUNKED_PREFILL_SIZE", "")  # e.g. "4096"
+        self.max_running_requests = os.getenv("SGLANG_MAX_RUNNING_REQUESTS", "")  # e.g. "64"
+
+        # If you already have a command you want to run, you can override entirely:
+        #   SGLANG_LAUNCH_CMD='python -m sglang.launch_server --model-path ...'
+        launch_cmd = os.getenv("SGLANG_LAUNCH_CMD", "").strip()
+        if launch_cmd:
+            cmd = launch_cmd.split()
+        else:
+            # Default launch command (SGLang OpenAI-compatible server)
+            cmd = [
+                "python",
+                "-m",
+                "sglang.launch_server",
+                "--model-path",
+                self.model_path,
+                "--tokenizer-path",
+                self.tokenizer_path,
+                "--host",
+                "0.0.0.0",
+                "--port",
+                str(self.port),
+                "--tp-size",
+                str(self.tp_size),
+            ]
+
+            if self.chunked_prefill_size:
+                cmd += ["--chunked-prefill-size", str(self.chunked_prefill_size)]
+            if self.max_running_requests:
+                cmd += ["--max-running-requests", str(self.max_running_requests)]
+
+        self.health_url = f"http://{self.host}:{self.port}/health"
+        self.chat_url = f"http://{self.host}:{self.port}/v1/chat/completions"
+
+        # Start SGLang server if not already up
+        if not _is_port_open(self.host, self.port):
+            # Important: do NOT use stdout=PIPE in production unless you drain it (deadlocks).
+            self.proc = subprocess.Popen(
+                cmd,
+                env=os.environ.copy(),
+            )
+        else:
+            self.proc = None
+
+        _wait_for_server(self.host, self.port, self.health_url, timeout_s=int(os.getenv("SGLANG_STARTUP_TIMEOUT", "600")))
+
+        # Model name presented to OpenAI-compatible API (some servers accept "model" as optional)
+        self.served_model_name = os.getenv("SGLANG_SERVED_MODEL_NAME", "ALIA-40b-instruct-nvfp4")
+
+    def __call__(self, data: Dict[str, Any]) -> Union[str, Dict[str, Any]]:
+        inputs = data.get("inputs", data)  # sometimes HF passes the full payload as inputs
+        params = data.get("parameters", {}) or {}
+
+        messages = _coerce_messages(inputs)
+        gen = _map_generation_params(params)
+
+        payload: Dict[str, Any] = {
+            "model": self.served_model_name,
+            "messages": messages,
+            **gen,
+        }
+
+        # Optional: allow user to set response_format / tools, etc. via "parameters"
+        # We pass through a small allowlist safely.
+        for k in ("response_format", "tools", "tool_choice"):
+            if k in params and params[k] is not None:
+                payload[k] = params[k]
+
+        try:
+            r = requests.post(self.chat_url, json=payload, timeout=float(os.getenv("SGLANG_REQUEST_TIMEOUT", "300")))
+            r.raise_for_status()
+            out = r.json()
+        except requests.RequestException as e:
+            raise RuntimeError(f"SGLang request failed: {e}") from e
+
+        # Normalize return to what HF widgets commonly expect:
+        # either a raw string or a dict with generated_text
+        try:
+            text = out["choices"][0]["message"]["content"]
+        except Exception:
+            # Fallback: return the full response
+            return out
+
+        # If caller asked for "details", return full payload
+        if bool(params.get("return_full_text")) or bool(params.get("details")):
+            return {
+                "generated_text": text,
+                "raw": out,
+            }
+
+        return text
+

handler.py.bak

@@ -1,296 +0,0 @@
-# handler.py
-# handler.py
-# Hugging Face Inference Endpoints "custom handler" for TensorRT-LLM (trtllm-serve),
-# including NVFP4-quantized engines.
-#
-# Expected by HF Inference Toolkit:
-#   - file name: handler.py (repo root)
-#   - class: EndpointHandler with __init__(path) and __call__(data)
-#
-# This handler:
-#   1) starts `trtllm-serve <model_dir>` once (lazy init)
-#   2) forwards requests to the local OpenAI-compatible HTTP API
-#
-# Environment variables (optional):
-#   TRTLLM_HOST           default: 127.0.0.1
-#   TRTLLM_PORT           default: 8000
-#   TRTLLM_START_CMD      default: "trtllm-serve"
-#   TRTLLM_START_ARGS     default: ""   (extra args appended verbatim)
-#   TRTLLM_HEALTH_PATH    default: "/health"
-#   TRTLLM_READY_TIMEOUT  default: 180  (seconds)
-#   TRTLLM_VERBOSE        default: "0"
-#
-# Notes:
-# - If your container uses a different binary or endpoints, set TRTLLM_START_CMD
-#   and/or adjust _chat/_completion URLs below.
-# - HF will call __call__ with a dict similar to:
-#     {"inputs": "...", "parameters": {...}}
-#   or for chat:
-#     {"messages": [...], "parameters": {...}}
-
-from __future__ import annotations
-
-import json
-import os
-import subprocess
-import time
-import threading
-from typing import Any, Dict, Optional
-
-try:
-    import requests
-except Exception:  # pragma: no cover
-    requests = None  # type: ignore
-
-
-class EndpointHandler:
-    _lock = threading.Lock()
-    _server_proc: Optional[subprocess.Popen] = None
-    _server_started: bool = False
-
-    def __init__(self, path: str):
-        # HF passes the model directory path (repo checkout) here.
-        self.model_dir = path
-
-        self.host = os.getenv("TRTLLM_HOST", "127.0.0.1")
-        self.port = int(os.getenv("TRTLLM_PORT", "8000"))
-        self.base_url = f"http://{self.host}:{self.port}"
-
-        self.health_path = os.getenv("TRTLLM_HEALTH_PATH", "/health")
-        self.ready_timeout = int(os.getenv("TRTLLM_READY_TIMEOUT", "180"))
-
-        self.start_cmd = os.getenv("TRTLLM_START_CMD", "trtllm-serve")
-        self.start_args = os.getenv("TRTLLM_START_ARGS", "").strip()
-
-        self.verbose = os.getenv("TRTLLM_VERBOSE", "0").strip() in ("1", "true", "TRUE", "yes", "YES")
-
-    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        self._ensure_server()
-
-        # HF commonly uses:
-        #   - data["inputs"] + data["parameters"]
-        # For chat-like:
-        #   - data["messages"] + data["parameters"]
-        parameters = data.get("parameters") or {}
-        if not isinstance(parameters, dict):
-            parameters = {}
-
-        # If the caller provides "messages", treat it as chat.
-        if "messages" in data and isinstance(data["messages"], list):
-            return self._handle_chat(data["messages"], parameters)
-
-        # Otherwise treat as completion.
-        inputs = data.get("inputs")
-        if inputs is None:
-            # Some clients use "prompt"
-            inputs = data.get("prompt")
-
-        if isinstance(inputs, list):
-            # Batch prompts: run sequentially (simple + robust).
-            outputs = [self._handle_completion(prompt, parameters) for prompt in inputs]
-            return {"results": outputs}
-
-        if not isinstance(inputs, str):
-            raise ValueError("Expected 'inputs' (or 'prompt') to be a string or list of strings.")
-
-        return self._handle_completion(inputs, parameters)
-
-    # -------------------------
-    # TensorRT-LLM server start
-    # -------------------------
-    def _ensure_server(self) -> None:
-        with self._lock:
-            if self._server_started:
-                return
-
-            # If server already reachable (e.g., started by container entrypoint), skip spawning.
-            if self._is_healthy():
-                self._server_started = True
-                return
-
-            cmd = [self.start_cmd, self.model_dir]
-
-            if self.start_args:
-                # Append extra args verbatim, allowing the user to pass things like:
-                #   "--backend pytorch --max_batch_size 4 --port 8000"
-                cmd.extend(self.start_args.split())
-
-            # Ensure server binds to desired port if user didn't specify it.
-            # If you already pass "--port" in TRTLLM_START_ARGS, this is redundant but harmless.
-            if "--port" not in cmd:
-                cmd.extend(["--port", str(self.port)])
-
-            if self.verbose:
-                print(f"[handler] Starting TensorRT-LLM server: {' '.join(cmd)}", flush=True)
-
-            # Start server process.
-            # Important: do not use shell=True.
-            self._server_proc = subprocess.Popen(
-                cmd,
-                stdout=subprocess.PIPE,
-                stderr=subprocess.STDOUT,
-                env=os.environ.copy(),
-                text=True,
-                bufsize=1,
-            )
-
-            # Wait until healthy
-            self._wait_until_ready()
-
-            self._server_started = True
-
-    def _wait_until_ready(self) -> None:
-        deadline = time.time() + self.ready_timeout
-        last_line = None
-
-        while time.time() < deadline:
-            if self._server_proc is not None:
-                # If process exited early, surface logs.
-                code = self._server_proc.poll()
-                if code is not None:
-                    logs = self._drain_logs(max_lines=2000)
-                    raise RuntimeError(
-                        f"TensorRT-LLM server exited with code {code} before becoming ready.\n"
-                        f"Last logs:\n{logs}"
-                    )
-
-            if self._is_healthy():
-                if self.verbose:
-                    print("[handler] TensorRT-LLM server is healthy.", flush=True)
-                return
-
-            # Optionally peek at logs to help debugging (non-blocking-ish).
-            if self.verbose:
-                line = self._read_one_log_line()
-                if line:
-                    last_line = line.strip()
-                    print(f"[trtllm] {last_line}", flush=True)
-
-            time.sleep(0.5)
-
-        logs = self._drain_logs(max_lines=500)
-        raise TimeoutError(
-            f"TensorRT-LLM server not ready after {self.ready_timeout}s. "
-            f"Health endpoint: {self.base_url}{self.health_path}\n"
-            f"Recent logs:\n{logs}"
-        )
-
-    def _is_healthy(self) -> bool:
-        try:
-            if requests is None:
-                return False
-            r = requests.get(f"{self.base_url}{self.health_path}", timeout=1.5)
-            return 200 <= r.status_code < 300
-        except Exception:
-            return False
-
-    def _read_one_log_line(self) -> Optional[str]:
-        try:
-            if self._server_proc and self._server_proc.stdout:
-                return self._server_proc.stdout.readline()
-        except Exception:
-            return None
-        return None
-
-    def _drain_logs(self, max_lines: int = 500) -> str:
-        if not self._server_proc or not self._server_proc.stdout:
-            return ""
-        lines = []
-        try:
-            for _ in range(max_lines):
-                line = self._server_proc.stdout.readline()
-                if not line:
-                    break
-                lines.append(line.rstrip("\n"))
-        except Exception:
-            pass
-        return "\n".join(lines)
-
-    # -------------------------
-    # Request forwarding
-    # -------------------------
-    def _handle_chat(self, messages: list, parameters: Dict[str, Any]) -> Dict[str, Any]:
-        payload = {
-            "model": parameters.pop("model", "trtllm"),
-            "messages": messages,
-        }
-        payload.update(self._map_parameters(parameters))
-
-        # TensorRT-LLM OpenAI-compatible chat endpoint
-        url = f"{self.base_url}/v1/chat/completions"
-        resp = self._post_json(url, payload)
-
-        # Normalize output for HF consumers
-        # Prefer returning OpenAI-like response, but also provide HF-style "generated_text".
-        generated_text = None
-        try:
-            generated_text = resp["choices"][0]["message"]["content"]
-        except Exception:
-            pass
-
-        return {
-            "generated_text": generated_text,
-            "raw": resp,
-        }
-
-    def _handle_completion(self, prompt: str, parameters: Dict[str, Any]) -> Dict[str, Any]:
-        payload = {
-            "model": parameters.pop("model", "trtllm"),
-            "prompt": prompt,
-        }
-        payload.update(self._map_parameters(parameters))
-
-        # TensorRT-LLM OpenAI-compatible completions endpoint
-        url = f"{self.base_url}/v1/completions"
-        resp = self._post_json(url, payload)
-
-        generated_text = None
-        try:
-            generated_text = resp["choices"][0]["text"]
-        except Exception:
-            pass
-
-        return {
-            "generated_text": generated_text,
-            "raw": resp,
-        }
-
-    def _post_json(self, url: str, payload: Dict[str, Any]) -> Dict[str, Any]:
-        if requests is None:
-            raise RuntimeError(
-                "The 'requests' package is not available in the container. "
-                "Install it or replace _post_json with urllib."
-            )
-        headers = {"Content-Type": "application/json"}
-        r = requests.post(url, headers=headers, data=json.dumps(payload), timeout=600)
-        if r.status_code >= 400:
-            raise RuntimeError(f"Upstream TRTLLM error {r.status_code}: {r.text}")
-        return r.json()
-
-    def _map_parameters(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
-        """
-        Map common HF generation parameters to OpenAI-compatible fields.
-        TensorRT-LLM may ignore unsupported fields; this mapping is conservative.
-        """
-        out: Dict[str, Any] = {}
-
-        # Common parameters
-        if "max_new_tokens" in parameters and "max_tokens" not in parameters:
-            out["max_tokens"] = parameters["max_new_tokens"]
-        if "max_tokens" in parameters:
-            out["max_tokens"] = parameters["max_tokens"]
-
-        for k in ("temperature", "top_p", "seed", "stop"):
-            if k in parameters:
-                out[k] = parameters[k]
-
-        # HF sometimes uses repetition_penalty; OpenAI doesn't have it.
-        # TensorRT-LLM may support presence/frequency penalties; pass through if provided.
-        for k in ("presence_penalty", "frequency_penalty"):
-            if k in parameters:
-                out[k] = parameters[k]
-
-        # Streaming is not supported by this handler (HF expects a single response).
-        # Ignore "stream" if present.
-        return out
-