chutes 0.3.61rc22__py3-none-any.whl → 0.5.3rc1__py3-none-any.whl

chutes/_version.py

@@ -1 +1 @@
-version = "0.3.61.rc22"
+version = "0.5.3.rc1"

chutes/cfsv_wrapper.py

@@ -1,16 +1,140 @@
 import os
-import sys
-import stat
-import subprocess
-from pathlib import Path
+import ctypes
+import asyncio
+from functools import lru_cache
+from fastapi import Request
 
 
-def main():
-    binary_path = Path(__file__).parent / "cfsv"
-    os.chmod(binary_path, os.stat(binary_path).st_mode | stat.S_IEXEC)
-    result = subprocess.run([str(binary_path)] + sys.argv[1:])
-    sys.exit(result.returncode)
+class CFSVWrapper:
+    def __init__(
+        self, lib_path=os.path.join(os.path.dirname(os.path.abspath(__file__)), "chutes-cfsv.so")
+    ):
+        self.lib = ctypes.CDLL(lib_path)
 
+        # cfsv_challenge(base_path, salt, sparse, index_file, exclude_path, result_buf, result_buf_size)
+        self.lib.cfsv_challenge.argtypes = [
+            ctypes.c_char_p,  # base_path
+            ctypes.c_char_p,  # salt
+            ctypes.c_int,     # sparse
+            ctypes.c_char_p,  # index_file
+            ctypes.c_char_p,  # exclude_path
+            ctypes.c_char_p,  # result_buf
+            ctypes.c_size_t,  # result_buf_size
+        ]
+        self.lib.cfsv_challenge.restype = ctypes.c_int
 
-if __name__ == "__main__":
-    main()
+        # cfsv_sizetest(test_dir, size_gib)
+        self.lib.cfsv_sizetest.argtypes = [
+            ctypes.c_char_p,  # test_dir
+            ctypes.c_size_t,  # size_gib
+        ]
+        self.lib.cfsv_sizetest.restype = ctypes.c_int
+
+        # cfsv_version()
+        self.lib.cfsv_version.argtypes = []
+        self.lib.cfsv_version.restype = ctypes.c_char_p
+
+    def challenge(
+        self,
+        salt,
+        mode="full",
+        base_path="/",
+        index_file="/etc/chutesfs.index",
+        exclude_path="/app/chute.py",
+    ):
+        """
+        Compute filesystem challenge hash.
+
+        Args:
+            salt: Challenge salt from validator
+            mode: "sparse" or "full" (default: "full")
+            base_path: Root path for file scanning (default: "/")
+            index_file: Path to encrypted index file (default: "/etc/chutesfs.index")
+            exclude_path: Path to exclude from hashing (default: "/app/chute.py")
+
+        Returns:
+            Hex string hash on success, None on failure
+        """
+        sparse = 1 if mode == "sparse" else 0
+        result_buf = ctypes.create_string_buffer(65)
+        ret = self.lib.cfsv_challenge(
+            base_path.encode() if isinstance(base_path, str) else base_path,
+            salt.encode() if isinstance(salt, str) else salt,
+            sparse,
+            index_file.encode() if isinstance(index_file, str) else index_file,
+            exclude_path.encode() if isinstance(exclude_path, str) else exclude_path,
+            result_buf,
+            65,
+        )
+        if ret == 0:
+            return result_buf.value.decode("utf-8")
+        return None
+
+    def sizetest(self, test_dir, size_gib):
+        """
+        Test filesystem capacity and integrity.
+
+        Args:
+            test_dir: Directory to test in
+            size_gib: Size in GiB to test
+
+        Returns:
+            True on success, False on failure
+        """
+        ret = self.lib.cfsv_sizetest(
+            test_dir.encode() if isinstance(test_dir, str) else test_dir,
+            size_gib,
+        )
+        return ret == 0
+
+    def version(self):
+        """
+        Get library version.
+
+        Returns:
+            Version string
+        """
+        result = self.lib.cfsv_version()
+        if result:
+            return result.decode("utf-8")
+        return None
+
+
+@lru_cache(maxsize=1)
+def get_cfsv():
+    """Lazily initialize CFSV wrapper (only works on Linux)."""
+    return CFSVWrapper()
+
+
+async def handle_challenge(request: Request):
+    loop = asyncio.get_event_loop()
+    salt = request.state.decrypted["salt"]
+    mode = request.state.decrypted.get("mode", "full")
+    exclude_path = request.state.decrypted.get("exclude_path", "/app/chute.py")
+    result = await loop.run_in_executor(
+        None,
+        get_cfsv().challenge,
+        salt,
+        mode,
+        "/",
+        "/etc/chutesfs.index",
+        exclude_path,
+    )
+    return {"result": result}
+
+
+async def handle_sizetest(request: Request):
+    loop = asyncio.get_event_loop()
+    test_dir = request.state.decrypted.get("test_dir", "/tmp")
+    size_gib = request.state.decrypted.get("size_gib", 10)
+    result = await loop.run_in_executor(
+        None,
+        get_cfsv().sizetest,
+        test_dir,
+        size_gib,
+    )
+    return {"result": result}
+
+
+async def handle_version(request: Request):
+    return {"result": get_cfsv().version()}

chutes/chute/base.py

@@ -45,6 +45,8 @@ class Chute(FastAPI):
         scaling_threshold: float = 0.75,
         allow_external_egress: bool = False,
         encrypted_fs: bool = False,
+        passthrough_headers: dict = {},
+        tee: bool = False,
         **kwargs,
     ):
         from chutes.chute.cord import Cord
@@ -71,8 +73,10 @@ class Chute(FastAPI):
         self.shutdown_after_seconds = shutdown_after_seconds
         self.allow_external_egress = allow_external_egress
         self.encrypted_fs = encrypted_fs
+        self.passthrough_headers = passthrough_headers
         self.docs_url = None
         self.redoc_url = None
+        self.tee = tee
 
     @property
     def name(self):

chutes/chute/cord.py

@@ -1,13 +1,15 @@
 import os
-import aiohttp
 import re
-import backoff
+import uuid
 import gzip
 import time
-import orjson as json
-import fickling
 import pickle
 import base64
+import aiohttp
+import asyncio
+import backoff
+import fickling
+import orjson as json
 from pydantic import ValidationError
 from typing import Optional, Dict, Any
 from fastapi import Request, HTTPException, status
@@ -44,6 +46,7 @@ class Cord:
         minimal_input_schema: Optional[Any] = None,
         output_content_type: Optional[str] = None,
         output_schema: Optional[Dict] = None,
+        sglang_passthrough: bool = False,
         **session_kwargs,
     ):
         """
@@ -67,6 +70,7 @@ class Cord:
         self._session_kwargs = session_kwargs
         self._provision_timeout = provision_timeout
         self._config = None
+        self._sglang_passthrough = sglang_passthrough
         self.input_models = (
             [input_schema] if input_schema and hasattr(input_schema, "__fields__") else None
         )
@@ -156,6 +160,9 @@ class Cord:
             raise InvalidPath(path)
         self._public_api_path = path
 
+    def _is_sglang_passthrough(self) -> bool:
+        return self._passthrough and self._sglang_passthrough
+
     @asynccontextmanager
     async def _local_call_base(self, *args, **kwargs):
         """
@@ -268,15 +275,28 @@ class Cord:
                         yield data["result"]
 
     @asynccontextmanager
-    async def _passthrough_call(self, **kwargs):
+    async def _passthrough_call(self, request: Request, **kwargs):
         """
         Call a passthrough endpoint.
         """
         logger.debug(
             f"Received passthrough call, passing along to {self.passthrough_path} via {self._method}"
         )
+        headers = kwargs.pop("headers", {}) or {}
+        if self._app.passthrough_headers:
+            headers.update(self._app.passthrough_headers)
+        kwargs["headers"] = headers
+
+        # Set (if needed) timeout.
+        timeout = None
+        if self._is_sglang_passthrough():
+            timeout = aiohttp.ClientTimeout(connect=5.0, total=None)
+        else:
+            total_timeout = kwargs.pop("timeout", 1800)
+            timeout = aiohttp.ClientTimeout(connect=5.0, total=total_timeout)
+
         async with aiohttp.ClientSession(
-            timeout=aiohttp.ClientTimeout(connect=5.0, total=900.0),
+            timeout=timeout,
             read_bufsize=8 * 1024 * 1024,
             base_url=f"http://127.0.0.1:{self._passthrough_port or 8000}",
         ) as session:
@@ -300,9 +320,88 @@ class Cord:
             function=self._func.__name__,
         ).set_to_current_time()
         encrypt = getattr(request.state, "_encrypt", None)
+
         try:
             if self._passthrough:
-                async with self._passthrough_call(**kwargs) as response:
+                rid = getattr(request.state, "sglang_rid", None)
+
+                # SGLang passthrough: run upstream call and disconnect watcher in parallel
+                if self._is_sglang_passthrough():
+
+                    async def call_upstream():
+                        async with self._passthrough_call(request, **kwargs) as response:
+                            if not 200 <= response.status < 300:
+                                try:
+                                    error_detail = await response.json()
+                                except Exception:
+                                    error_detail = await response.text()
+                                logger.error(
+                                    f"Failed to generate response from func={self._func.__name__}: {response.status=} -> {error_detail}"
+                                )
+                                raise HTTPException(
+                                    status_code=response.status,
+                                    detail=error_detail,
+                                )
+                            if encrypt:
+                                raw = await response.read()
+                                return {"json": encrypt(raw)}
+                            return await response.json()
+
+                    async def watch_disconnect():
+                        try:
+                            while True:
+                                message = await request._receive()
+                                if message.get("type") == "http.disconnect":
+                                    logger.info(
+                                        f"[{self._func.__name__}] Received http.disconnect, "
+                                        f"aborting upstream SGLang request (rid={rid})"
+                                    )
+                                    try:
+                                        await self._abort_sglang_request(rid)
+                                    except Exception as exc:
+                                        logger.warning(
+                                            f"Error while sending abort_request for rid={rid}: {exc}"
+                                        )
+                                    raise HTTPException(
+                                        status_code=499,
+                                        detail="Client disconnected during SGLang request",
+                                    )
+                        except HTTPException:
+                            raise
+                        except Exception as exc:
+                            logger.warning(f"watch_disconnect error: {exc}")
+                            raise HTTPException(
+                                status_code=499,
+                                detail="Client disconnected during SGLang request",
+                            )
+
+                    upstream_task = asyncio.create_task(call_upstream())
+                    watcher_task = asyncio.create_task(watch_disconnect())
+
+                    done, pending = await asyncio.wait(
+                        {upstream_task, watcher_task},
+                        return_when=asyncio.FIRST_COMPLETED,
+                    )
+                    for task in pending:
+                        task.cancel()
+
+                    if watcher_task in done:
+                        exc = watcher_task.exception()
+
+                        if exc:
+                            raise exc
+                        raise HTTPException(
+                            status_code=499,
+                            detail="Client disconnected during SGLang request",
+                        )
+                    result = upstream_task.result()
+                    logger.success(
+                        f"Completed request [{self._func.__name__} passthrough={self._passthrough}] "
+                        f"in {time.time() - started_at} seconds"
+                    )
+                    return result
+
+                async with self._passthrough_call(request, **kwargs) as response:
                     if not 200 <= response.status < 300:
                         try:
                             error_detail = await response.json()
@@ -316,15 +415,18 @@ class Cord:
                             detail=error_detail,
                         )
                     logger.success(
-                        f"Completed request [{self._func.__name__} passthrough={self._passthrough}] in {time.time() - started_at} seconds"
+                        f"Completed request [{self._func.__name__} passthrough={self._passthrough}] "
+                        f"in {time.time() - started_at} seconds"
                     )
                     if encrypt:
                         return {"json": encrypt(await response.read())}
                     return await response.json()
 
-            response = await self._func(self._app, *args, **kwargs)
+            # Non-passthrough call (local Python function)
+            response = await asyncio.wait_for(self._func(self._app, *args, **kwargs), 1800)
             logger.success(
-                f"Completed request [{self._func.__name__} passthrough={self._passthrough}] in {time.time() - started_at} seconds"
+                f"Completed request [{self._func.__name__} passthrough={self._passthrough}] "
+                f"in {time.time() - started_at} seconds"
             )
             if hasattr(response, "body"):
                 if encrypt:
@@ -340,6 +442,19 @@ class Cord:
             if encrypt:
                 return {"json": encrypt(json.dumps(response))}
             return response
+        except asyncio.CancelledError:
+            rid = getattr(request.state, "sglang_rid", None)
+            if self._is_sglang_passthrough():
+                logger.info(
+                    f"Non-stream request for {self._func.__name__} cancelled "
+                    f"(likely client disconnect), aborting SGLang rid={rid}"
+                )
+                try:
+                    await self._abort_sglang_request(rid)
+                except Exception as exc:
+                    logger.warning(f"Error while sending abort_request for rid={rid}: {exc}")
+            status = 499
+            raise
         except Exception as exc:
             logger.error(f"Error performing non-streamed call: {exc}")
             status = 500
@@ -371,7 +486,9 @@ class Cord:
         encrypt = getattr(request.state, "_encrypt", None)
         try:
             if self._passthrough:
-                async with self._passthrough_call(**kwargs) as response:
+                rid = getattr(request.state, "sglang_rid", None)
+
+                async with self._passthrough_call(request, **kwargs) as response:
                     if not 200 <= response.status < 300:
                         try:
                             error_detail = await response.json()
@@ -385,6 +502,13 @@ class Cord:
                             detail=error_detail,
                         )
                     async for content in response.content:
+                        if await request.is_disconnected():
+                            logger.info(
+                                f"Client disconnected for {self._func.__name__}, aborting upstream (rid={rid})"
+                            )
+                            await self._abort_sglang_request(rid)
+                            break
+
                         if encrypt:
                             yield encrypt(content) + "\n"
                         else:
@@ -402,6 +526,20 @@ class Cord:
             logger.success(
                 f"Completed request [{self._func.__name__}] in {time.time() - started_at} seconds"
             )
+        except asyncio.CancelledError:
+            rid = getattr(request.state, "sglang_rid", None)
+            if self._is_sglang_passthrough():
+                logger.info(
+                    f"Streaming cancelled for {self._func.__name__} "
+                    f"(likely client disconnect), aborting SGLang rid={rid}"
+                )
+                try:
+                    await self._abort_sglang_request(rid)
+                except Exception as exc:
+                    logger.warning(f"Error while sending abort_request for rid={rid}: {exc}")
+            status = 499
+            raise
+
         except Exception as exc:
             logger.error(f"Error performing stream call: {exc}")
             status = 500
@@ -418,6 +556,21 @@ class Cord:
                 status=status,
             ).observe(time.time() - started_at)
 
+    async def _abort_sglang_request(self, rid: Optional[str]):
+        if not rid or not self._is_sglang_passthrough():
+            return
+        try:
+            async with aiohttp.ClientSession(
+                timeout=aiohttp.ClientTimeout(connect=5.0, total=15.0),
+                base_url=f"http://127.0.0.1:{self._passthrough_port or 8000}",
+                headers=self._app.passthrough_headers or {},
+            ) as session:
+                logger.warning(f"Aborting SGLang request {rid=}")
+                await session.post("/abort_request", json={"rid": rid})
+                logger.success(f"Sent SGLang abort_request for rid={rid}")
+        except Exception as exc:
+            logger.warning(f"Failed to send abort_request for rid={rid}: {exc}")
+
     async def _request_handler(self, request: Request):
         """
         Decode/deserialize incoming request and call the appropriate function.
@@ -448,6 +601,13 @@ class Cord:
             else:
                 args = []
                 kwargs = {"json": request.state.decrypted} if request.state.decrypted else {}
+
+        # Set a custom request ID for SGLang passthroughs.
+        if self._is_sglang_passthrough() and isinstance(kwargs.get("json"), dict):
+            rid = uuid.uuid4().hex
+            kwargs["json"].setdefault("rid", rid)
+            request.state.sglang_rid = rid
+
         if not self._passthrough:
             if self.input_models and all([isinstance(args[idx], dict) for idx in range(len(args))]):
                 try:

chutes/chute/template/diffusion.py

@@ -152,6 +152,7 @@ def build_diffusion_chute(
     max_instances: int = 1,
     scaling_threshold: float = 0.75,
     shutdown_after_seconds: int = 300,
+    tee: bool = False,
 ):
     chute = Chute(
         username=username,
@@ -165,6 +166,7 @@ def build_diffusion_chute(
         shutdown_after_seconds=shutdown_after_seconds,
         max_instances=max_instances,
         scaling_threshold=scaling_threshold,
+        tee=tee,
     )
 
     @chute.on_startup()