Aegis

chutes/_version.py

@@ -1 +1 @@
-version = "0.5.4.rc15"
+version = "0.5.5.rc42"

chutes/cfsv_wrapper.py

@@ -1,4 +1,3 @@
-import os
 import ctypes
 import asyncio
 from functools import lru_cache
@@ -7,9 +6,14 @@
 
 class CFSVWrapper:
     def __init__(
-        self, lib_path=os.path.join(os.path.dirname(os.path.abspath(__file__)), "chutes-cfsv.so")
+        self,
+        lib_path=None,
     ):
-        self.lib = ctypes.CDLL(lib_path)
+        if lib_path:
+            self.lib = ctypes.CDLL(lib_path)
+        else:
+            # CFSV is compiled into chutes-aegis.so (loaded via LD_PRELOAD).
+            self.lib = ctypes.CDLL(None)
 
         # cfsv_challenge(base_path, salt, sparse, index_file, exclude_path, result_buf, result_buf_size)
         self.lib.cfsv_challenge.argtypes = [

chutes/chute/base.py

@@ -74,6 +74,8 @@ def __init__(
         self.allow_external_egress = allow_external_egress
         self.encrypted_fs = encrypted_fs
         self.passthrough_headers = passthrough_headers
+        self.passthrough_ssl_context = None
+        self._wrong_ssl_context = None
         self.docs_url = None
         self.redoc_url = None
         self.tee = tee

chutes/chute/template/embedding.py

@@ -13,7 +13,18 @@
 from chutes.image import Image
 from chutes.image.standard.vllm import VLLM
 from chutes.chute import Chute, ChutePack, NodeSelector
-from chutes.chute.template.helpers import set_default_cache_dirs, set_nccl_flags, monitor_engine
+from chutes.chute.template.helpers import (
+    set_default_cache_dirs,
+    set_nccl_flags,
+    monitor_engine,
+    generate_mtls_certs,
+    build_client_ssl_context,
+    build_wrong_client_ssl_context,
+    validate_mtls,
+    force_exit,
+    mtls_enabled,
+    set_encrypted_env_var,
+)
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
@@ -172,7 +183,10 @@ async def initialize_vllm_embedding(self):
         # Verify the cache contents.
         await verify_cache(repo_id=model_name, revision=revision)
 
-        set_default_cache_dirs(download_path)
+        set_default_cache_dirs(
+            download_path,
+            cache_version=getattr(self, "_source_hash", None),
+        )
 
         torch.cuda.empty_cache()
         torch.cuda.init()
@@ -212,6 +226,31 @@ async def initialize_vllm_embedding(self):
         if "--api-key" in engine_args:
             raise ValueError("You may not override api key!")
 
+        use_mtls = mtls_enabled()
+        ssl_ctx = None
+        wrong_ssl_ctx = None
+
+        if use_mtls:
+            # Generate ephemeral mTLS certificates.
+            certs = generate_mtls_certs()
+            ssl_ctx = build_client_ssl_context(
+                certs["ca_cert_file"],
+                certs["client_cert_file"],
+                certs["client_key_file"],
+                certs["password"],
+            )
+            wrong_ssl_ctx = build_wrong_client_ssl_context(
+                certs["ca_cert_file"],
+                certs["wrong_client_cert_file"],
+                certs["wrong_client_key_file"],
+                certs["password"],
+            )
+            self.passthrough_ssl_context = ssl_ctx
+            self._wrong_ssl_context = wrong_ssl_ctx
+            logger.info("mTLS enabled for vLLM embedding engine communication")
+        else:
+            logger.warning("mTLS disabled (LLM_ENGINE_MTLS_ENABLE not set)")
+
         env = os.environ.copy()
         env["PYTHONDONTWRITEBYTECODE"] = "1"
         if enable_chunked_processing:
@@ -220,13 +259,19 @@ async def initialize_vllm_embedding(self):
         env["HF_HUB_OFFLINE"] = "1"
         env["SGL_MODEL_NAME"] = self.name
         env["SGL_REVISION"] = revision
+        if use_mtls:
+            env["VLLM_SSL_KEYFILE_PEM"] = certs["server_key_pem"].decode()
+            env["VLLM_SSL_CERTFILE_PEM"] = certs["server_cert_pem"].decode()
+            env["VLLM_SSL_CA_CERTS_PEM"] = certs["ca_cert_pem"].decode()
+            set_encrypted_env_var(env, "VLLM_SSL_KEYFILE_PASSWORD", certs["password"])
 
+        ssl_args = " --ssl-cert-reqs 2" if use_mtls else ""
         pooler_config_arg = shlex.quote(json.dumps(pooler_config))
         startup_command = (
             f"{sys.executable} -m vllm.entrypoints.openai.api_server "
             f"--model {model_name} --served-model-name {self.name} "
             f"--revision {revision} --pooler-config {pooler_config_arg} "
-            f"--port 10101 --host 127.0.0.1 --api-key {api_key} {engine_args}"
+            f"--port 10101 --host 127.0.0.1 --api-key {api_key}{ssl_args} {engine_args}"
         )
         display_cmd = startup_command.replace(api_key, "*" * len(api_key))
         parts = shlex.split(startup_command)
@@ -237,14 +282,38 @@ async def initialize_vllm_embedding(self):
 
         server_ready = asyncio.Event()
         self._monitor_task = asyncio.create_task(
-            monitor_engine(self._vllm_process, api_key, server_ready, model_name=self.name)
+            monitor_engine(
+                self._vllm_process,
+                api_key,
+                server_ready,
+                model_name=self.name,
+                ssl_context=ssl_ctx,
+                wrong_ssl_context=wrong_ssl_ctx,
+            )
         )
 
+        def _on_monitor_done(t):
+            if t.cancelled():
+                return
+            exc = t.exception()
+            if exc:
+                logger.error("Embedding vLLM monitor task failed, terminating: {}", exc)
+                force_exit(1)
+
+        self._monitor_task.add_done_callback(_on_monitor_done)
+
+        base_url = "https://127.0.0.1:10101" if use_mtls else "http://127.0.0.1:10101"
         while True:
+            if self._vllm_process.poll() is not None:
+                raise RuntimeError(
+                    "Embedding vLLM subprocess exited before readiness check "
+                    f"(exit={self._vllm_process.returncode})"
+                )
             try:
-                async with aiohttp.ClientSession() as session:
+                connector = aiohttp.TCPConnector(ssl=ssl_ctx) if ssl_ctx else None
+                async with aiohttp.ClientSession(connector=connector) as session:
                     async with session.get(
-                        "http://127.0.0.1:10101/v1/models",
+                        f"{base_url}/v1/models",
                         headers={"Authorization": f"Bearer {api_key}"},
                     ) as resp:
                         if resp.status == 200:
@@ -255,8 +324,10 @@ async def initialize_vllm_embedding(self):
             await asyncio.sleep(1)
 
         self.passthrough_headers["Authorization"] = f"Bearer {api_key}"
+        if use_mtls:
+            await validate_mtls(self.name, api_key, ssl_ctx, wrong_ssl_ctx)
         server_ready.set()
-        logger.info("✅ Embedding server initialized successfully!")
+        logger.info("Embedding server initialized successfully!")
 
     @chute.cord(
         passthrough_path="/v1/embeddings",