truss 0.10.13__py3-none-any.whl → 0.11.1rc1__py3-none-any.whl

truss/cli/chains_commands.py

@@ -359,7 +359,7 @@ def push_chain(
     "--name",
     type=str,
     required=False,
-    help="Name of the chain to be deployed, if not given, the entrypoint name is used.",
+    help="Name of the chain to be watched. If not given, the entrypoint name is used.",
 )
 @click.option(
     "--remote",

truss/cli/train/core.py

@@ -1,4 +1,5 @@
 import json
+import os
 import tarfile
 import tempfile
 from dataclasses import dataclass
@@ -16,6 +17,11 @@ from truss.cli.train.metrics_watcher import MetricsWatcher
 from truss.cli.train.types import PrepareCheckpointArgs, PrepareCheckpointResult
 from truss.cli.utils import common as cli_common
 from truss.cli.utils.output import console
+from truss.remote.baseten.custom_types import (
+    FileSummary,
+    FileSummaryWithTotalSize,
+    GetCacheSummaryResponseV1,
+)
 from truss.remote.baseten.remote import BasetenRemote
 from truss_train import loader
 from truss_train.definitions import DeployCheckpointsConfig
@@ -446,6 +452,44 @@ def fetch_project_by_name_or_id(
         raise click.ClickException(f"Error fetching project: {str(e)}")
 
 
+def create_file_summary_with_directory_sizes(
+    files: list[FileSummary],
+) -> list[FileSummaryWithTotalSize]:
+    directory_sizes = calculate_directory_sizes(files)
+    return [
+        FileSummaryWithTotalSize(
+            file_summary=file_info,
+            total_size=directory_sizes.get(file_info.path, file_info.size_bytes),
+        )
+        for file_info in files
+    ]
+
+
+def calculate_directory_sizes(
+    files: list[FileSummary], max_depth: int = 100
+) -> dict[str, int]:
+    directory_sizes = {}
+
+    for file_info in files:
+        if file_info.file_type == "directory":
+            directory_sizes[file_info.path] = 0
+
+    for file_info in files:
+        current_path = file_info.path
+        for i in range(max_depth):
+            if current_path is None:
+                break
+            if current_path in directory_sizes:
+                directory_sizes[current_path] += file_info.size_bytes
+            # Move to parent directory
+            parent = os.path.dirname(current_path)
+            if parent == current_path:  # Reached root
+                break
+            current_path = parent
+
+    return directory_sizes
+
+
 def view_cache_summary(
     remote_provider: BasetenRemote,
     project_id: str,
@@ -454,12 +498,14 @@ def view_cache_summary(
 ):
     """View cache summary for a training project."""
     try:
-        cache_data = remote_provider.api.get_cache_summary(project_id)
+        raw_cache_data = remote_provider.api.get_cache_summary(project_id)
 
-        if not cache_data:
+        if not raw_cache_data:
             console.print("No cache summary found for this project.", style="yellow")
             return
 
+        cache_data = GetCacheSummaryResponseV1.model_validate(raw_cache_data)
+
         table = rich.table.Table(title=f"Cache summary for project: {project_id}")
         table.add_column("File Path", style="cyan")
         table.add_column("Size", style="green")
@@ -467,58 +513,48 @@ def view_cache_summary(
         table.add_column("Type")
         table.add_column("Permissions", style="magenta")
 
-        files = cache_data.get("file_summaries", [])
+        files = cache_data.file_summaries
         if not files:
             console.print("No files found in cache.", style="yellow")
             return
 
-        reverse = order == SORT_ORDER_DESC
+        files_with_total_sizes = create_file_summary_with_directory_sizes(files)
 
-        if sort_by == SORT_BY_FILEPATH:
-            files.sort(key=lambda x: x.get("path", ""), reverse=reverse)
-        elif sort_by == SORT_BY_SIZE:
-            files.sort(key=lambda x: x.get("size_bytes", 0), reverse=reverse)
-        elif sort_by == SORT_BY_MODIFIED:
-            files.sort(key=lambda x: x.get("modified", ""), reverse=reverse)
-        elif sort_by == SORT_BY_TYPE:
-            files.sort(key=lambda x: x.get("file_type", ""), reverse=reverse)
-        elif sort_by == SORT_BY_PERMISSIONS:
-            files.sort(key=lambda x: x.get("permissions", ""), reverse=reverse)
-
-        total_size = 0
-        for file_info in files:
-            total_size += file_info.get("size_bytes", 0)
+        reverse = order == SORT_ORDER_DESC
+        sort_key = _get_sort_key(sort_by)
+        files_with_total_sizes.sort(key=sort_key, reverse=reverse)
 
+        total_size = sum(
+            file_info.file_summary.size_bytes for file_info in files_with_total_sizes
+        )
         total_size_str = common.format_bytes_to_human_readable(total_size)
 
         console.print(
-            f"📅 Cache captured at: {cache_data.get('timestamp', 'Unknown')}",
-            style="bold blue",
+            f"📅 Cache captured at: {cache_data.timestamp}", style="bold blue"
         )
+        console.print(f"📁 Project ID: {cache_data.project_id}", style="bold blue")
+        console.print()
         console.print(
-            f"📁 Project ID: {cache_data.get('project_id', 'Unknown')}",
-            style="bold blue",
+            f"📊 Total files: {len(files_with_total_sizes)}", style="bold green"
         )
-        console.print()
-        console.print(f"📊 Total files: {len(files)}", style="bold green")
         console.print(f"💾 Total size: {total_size_str}", style="bold green")
         console.print()
 
-        for file_info in files:
-            size_bytes = file_info.get("size_bytes", 0)
+        for file_info in files_with_total_sizes:
+            total_size = file_info.total_size
 
-            size_str = cli_common.format_bytes_to_human_readable(int(size_bytes))
+            size_str = cli_common.format_bytes_to_human_readable(int(total_size))
 
             modified_str = cli_common.format_localized_time(
-                file_info.get("modified", "Unknown")
+                file_info.file_summary.modified
             )
 
             table.add_row(
-                file_info.get("path", "Unknown"),
+                file_info.file_summary.path,
                 size_str,
                 modified_str,
-                file_info.get("file_type", "Unknown"),
-                file_info.get("permissions", "Unknown"),
+                file_info.file_summary.file_type or "Unknown",
+                file_info.file_summary.permissions or "Unknown",
             )
 
         console.print(table)
@@ -528,6 +564,21 @@ def view_cache_summary(
         raise
 
 
+def _get_sort_key(sort_by: str) -> Callable[[FileSummaryWithTotalSize], Any]:
+    if sort_by == SORT_BY_FILEPATH:
+        return lambda x: x.file_summary.path
+    elif sort_by == SORT_BY_SIZE:
+        return lambda x: x.total_size
+    elif sort_by == SORT_BY_MODIFIED:
+        return lambda x: x.file_summary.modified
+    elif sort_by == SORT_BY_TYPE:
+        return lambda x: x.file_summary.file_type or ""
+    elif sort_by == SORT_BY_PERMISSIONS:
+        return lambda x: x.file_summary.permissions or ""
+    else:
+        raise ValueError(f"Invalid --sort argument: {sort_by}")
+
+
 def view_cache_summary_by_project(
     remote_provider: BasetenRemote,
     project_identifier: str,

truss/contexts/image_builder/serving_image_builder.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import json
 import logging
+import os
 import re
 import shutil
 from abc import ABC, abstractmethod
@@ -783,6 +784,10 @@ class ServingImageBuilder(ImageBuilder):
             config
         )
 
+        non_root_user = os.getenv("BT_USE_NON_ROOT_USER", False)
+        enable_model_container_admin_commands = os.getenv(
+            "BT_ENABLE_MODEL_CONTAINER_ADMIN_CMDS"
+        )
         dockerfile_contents = dockerfile_template.render(
             should_install_server_requirements=should_install_server_requirements,
             base_image_name_and_tag=base_image_name_and_tag,
@@ -816,6 +821,8 @@ class ServingImageBuilder(ImageBuilder):
             build_commands=build_commands,
             use_local_src=config.use_local_src,
             passthrough_environment_variables=passthrough_environment_variables,
+            non_root_user=non_root_user,
+            enable_model_container_admin_commands=enable_model_container_admin_commands,
             **FILENAME_CONSTANTS_MAP,
         )
         # Consolidate repeated empty lines to single empty lines.

truss/contexts/local_loader/docker_build_emulator.py

@@ -1,3 +1,4 @@
+import re
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Dict, List
@@ -31,12 +32,32 @@ class DockerBuildEmulator:
         self._context_dir = context_dir
 
     def run(self, fs_root_dir: Path) -> DockerBuildEmulatorResult:
-        def _resolve_env(key: str) -> str:
-            if key.startswith("$"):
-                key = key.replace("$", "", 1)
-                v = result.env[key]
-                return v
-            return key
+        def _resolve_env(in_value: str) -> str:
+            # Valid environment variable name pattern
+            var_name_pattern = r"[A-Za-z_][A-Za-z0-9_]*"
+
+            # Handle ${VAR} syntax
+            def replace_braced_var(match):
+                var_name = match.group(1)
+                return result.env.get(
+                    var_name, match.group(0)
+                )  # Return original if not found
+
+            # Handle $VAR syntax (word boundary ensures we don't match parts of other vars)
+            def replace_simple_var(match):
+                var_name = match.group(1)
+                return result.env.get(
+                    var_name, match.group(0)
+                )  # Return original if not found
+
+            # Replace ${VAR} patterns first, using % substitution to avoid additional braces noise with f-strings
+            value = re.sub(
+                r"\$\{(%s)\}" % var_name_pattern, replace_braced_var, in_value
+            )
+            # Then replace remaining $VAR patterns (only at word boundaries)
+            value = re.sub(r"\$(%s)\b" % var_name_pattern, replace_simple_var, value)
+
+            return value
 
         def _resolve_values(keys: List[str]) -> List[str]:
             return list(map(_resolve_env, keys))
@@ -53,11 +74,14 @@ class DockerBuildEmulator:
             if cmd.instruction == DockerInstruction.ENTRYPOINT:
                 result.entrypoint = list(values)
             if cmd.instruction == DockerInstruction.COPY:
+                # Filter out --chown flags
+                filtered_values = [v for v in values if not v.startswith("--chown")]
+
                 # NB(nikhil): Skip COPY commands with --from flag (multi-stage builds)
-                if len(values) != 2:
+                if len(filtered_values) != 2:
                     continue
 
-                src, dst = values
+                src, dst = filtered_values
                 src = src.replace("./", "", 1)
                 dst = dst.replace("/", "", 1)
                 copy_tree_or_file(self._context_dir / src, fs_root_dir / dst)

truss/remote/baseten/custom_types.py

@@ -138,6 +138,13 @@ class FileSummary(pydantic.BaseModel):
     )
 
 
+class FileSummaryWithTotalSize(pydantic.BaseModel):
+    file_summary: FileSummary
+    total_size: int = pydantic.Field(
+        description="Total size of the file and all its subdirectories"
+    )
+
+
 class GetCacheSummaryResponseV1(pydantic.BaseModel):
     """Response for getting cache summary."""
 

truss/templates/base.Dockerfile.jinja

@@ -8,6 +8,35 @@ FROM {{ base_image_name_and_tag }} AS truss_server
 {%- set python_executable = config.base_image.python_executable_path or 'python3' %}
 ENV PYTHON_EXECUTABLE="{{ python_executable }}"
 
+{%- set app_username = "app" %}  {# needed later for USER directive#}
+{% block user_setup %}
+{%- set app_user_uid = 60000 %}
+{%- set control_server_dir = "/control" %}
+{%- set default_owner = "root:root" %}
+{# The non-root user's home directory. #}
+{# uv will use $HOME to install packages. #}
+ENV HOME=/home/{{ app_username }}
+{# Directory containing inference server code. #}
+ENV APP_HOME=/{{ app_username }}
+RUN mkdir -p ${APP_HOME} {{ control_server_dir }}
+{# Create a non-root user to run model containers. #}
+RUN useradd -u {{ app_user_uid }} -ms /bin/bash {{ app_username }}
+{% endblock %} {#- endblock user_setup #}
+
+{#- at the very beginning, set non-interactive mode for apt #}
+ENV DEBIAN_FRONTEND=noninteractive
+
+{# If non-root user is enabled and model container admin commands are enabled, install sudo #}
+{# to allow the non-root user to install packages. #}
+{%- if non_root_user and enable_model_container_admin_commands %}
+RUN apt update && apt install -y sudo
+{%- set allowed_admin_commands = ["/usr/bin/apt install *", "/usr/bin/apt update"] %}
+RUN echo "Defaults:{{ app_username }} passwd_tries=0\n{{ app_username }} ALL=(root) NOPASSWD: {{ allowed_admin_commands | join(", ") }}" > /etc/sudoers.d/app-packages
+RUN chmod 0440 /etc/sudoers.d/app-packages
+{#- optional but good practice: check if the sudoers file is valid #}
+RUN visudo -c
+{%- endif %} {#- endif non_root_user and enable_model_container_admin_commands #}
+
 {%- set UV_VERSION = "0.7.19" %}
 {#
 NB(nikhil): We use a semi-complex uv installation command across the board:
@@ -39,7 +68,8 @@ RUN if ! command -v uv >/dev/null 2>&1; then \
     command -v curl >/dev/null 2>&1 || (apt update && apt install -y curl) && \
     curl -LsSf --retry 5 --retry-delay 5 https://astral.sh/uv/{{ UV_VERSION }}/install.sh | sh; \
 fi
-ENV PATH="/root/.local/bin:$PATH"
+{# Add the user's local bin to the path, used by uv. #}
+ENV PATH=${PATH}:${HOME}/.local/bin
 {% endblock %}
 
 {% block base_image_patch %}
@@ -57,7 +87,7 @@ RUN {{ sys_pip_install_command }} install mkl
 
 {% block install_system_requirements %}
     {%- if should_install_system_requirements %}
-COPY ./{{ system_packages_filename }} {{ system_packages_filename }}
+COPY --chown={{ default_owner }} ./{{ system_packages_filename }} {{ system_packages_filename }}
 RUN apt-get update && apt-get install --yes --no-install-recommends $(cat {{ system_packages_filename }}) \
     && apt-get autoremove -y \
     && apt-get clean -y \
@@ -68,11 +98,11 @@ RUN apt-get update && apt-get install --yes --no-install-recommends $(cat {{ sys
 
 {% block install_requirements %}
     {%- if should_install_user_requirements_file %}
-COPY ./{{ user_supplied_requirements_filename }} {{ user_supplied_requirements_filename }}
+COPY --chown={{ default_owner }} ./{{ user_supplied_requirements_filename }} {{ user_supplied_requirements_filename }}
 RUN {{ sys_pip_install_command }} -r {{ user_supplied_requirements_filename }} --no-cache-dir
     {%- endif %}
     {%- if should_install_requirements %}
-COPY ./{{ config_requirements_filename }} {{ config_requirements_filename }}
+COPY --chown={{ default_owner }} ./{{ config_requirements_filename }} {{ config_requirements_filename }}
 RUN {{ sys_pip_install_command }} -r {{ config_requirements_filename }} --no-cache-dir
 {%- endif %}
 {% endblock %}
@@ -80,7 +110,6 @@ RUN {{ sys_pip_install_command }} -r {{ config_requirements_filename }} --no-cac
 
 
 {%- if not config.docker_server %}
-ENV APP_HOME="/app"
 WORKDIR $APP_HOME
 {%- endif %}
 
@@ -90,7 +119,7 @@ WORKDIR $APP_HOME
 
 {% block bundled_packages_copy %}
     {%- if bundled_packages_dir_exists %}
-COPY ./{{ config.bundled_packages_dir }} /packages
+COPY --chown={{ default_owner }} ./{{ config.bundled_packages_dir }} /packages
 {%- endif %}
 {% endblock %}
 

truss/templates/cache.Dockerfile.jinja

@@ -1,20 +1,21 @@
 FROM python:3.11-slim AS cache_warmer
 
-RUN mkdir -p /app/model_cache
-WORKDIR /app
+ENV APP_HOME=/app
+RUN mkdir -p ${APP_HOME}/model_cache
+WORKDIR ${APP_HOME}
 
 {% if hf_access_token %}
 ENV HUGGING_FACE_HUB_TOKEN="{{hf_access_token}}"
 {% endif %}
 
 RUN apt-get -y update; apt-get -y install curl; curl -s https://baseten-public.s3.us-west-2.amazonaws.com/bin/b10cp-5fe8dc7da-linux-amd64 -o /app/b10cp; chmod +x /app/b10cp
-ENV B10CP_PATH_TRUSS="/app/b10cp"
-COPY ./cache_requirements.txt /app/cache_requirements.txt
-RUN pip install -r /app/cache_requirements.txt --no-cache-dir && rm -rf /root/.cache/pip
-COPY ./cache_warmer.py /cache_warmer.py
+ENV B10CP_PATH_TRUSS="${APP_HOME}/b10cp"
+COPY --chown={{ default_owner }} ./cache_requirements.txt ${APP_HOME}/cache_requirements.txt
+RUN pip install -r ${APP_HOME}/cache_requirements.txt --no-cache-dir && rm -rf /root/.cache/pip
+COPY --chown={{ default_owner }} ./cache_warmer.py /cache_warmer.py
 
 {% for credential in credentials_to_cache  %}
-COPY ./{{credential}} /app/{{credential}}
+COPY ./{{credential}} ${APP_HOME}/{{credential}}
 {% endfor %}
 
 {% for repo, hf_dir in models.items() %}

truss/templates/control/control/endpoints.py

@@ -1,14 +1,15 @@
 import asyncio
 import logging
-from typing import Any, Callable, Dict
+from typing import Any, Callable, Dict, Optional, Protocol
 
 import httpx
 from fastapi import APIRouter, WebSocket
 from fastapi.responses import JSONResponse, StreamingResponse
+from httpx_ws import AsyncWebSocketSession, WebSocketDisconnect, aconnect_ws
 from httpx_ws import _exceptions as httpx_ws_exceptions
-from httpx_ws import aconnect_ws
 from starlette.requests import ClientDisconnect, Request
 from starlette.responses import Response
+from starlette.websockets import WebSocketDisconnect as StartletteWebSocketDisconnect
 from tenacity import RetryCallState, Retrying, retry_if_exception_type, wait_fixed
 from wsproto.events import BytesMessage, TextMessage
 
@@ -30,6 +31,10 @@ BASE_RETRY_EXCEPTIONS = (
 control_app = APIRouter()
 
 
+class CloseableWebsocket(Protocol):
+    async def close(self, code: int = 1000, reason: Optional[str] = None) -> None: ...
+
+
 @control_app.get("/")
 def index():
     return {}
@@ -118,13 +123,75 @@ def inference_retries(
         yield attempt
 
 
-async def _safe_close_ws(ws: WebSocket, logger: logging.Logger):
+async def _safe_close_ws(
+    ws: CloseableWebsocket,
+    logger: logging.Logger,
+    code: int = 1000,
+    reason: Optional[str] = None,
+):
     try:
-        await ws.close()
+        await ws.close(code, reason)
     except RuntimeError as close_error:
         logger.debug(f"Duplicate close of websocket: `{close_error}`.")
 
 
+async def forward_to_server(
+    client_ws: WebSocket, server_ws: AsyncWebSocketSession
+) -> None:
+    while True:
+        message = await client_ws.receive()
+        if message.get("type") == "websocket.disconnect":
+            raise StartletteWebSocketDisconnect(
+                message.get("code", 1000), message.get("reason")
+            )
+        if "text" in message:
+            await server_ws.send_text(message["text"])
+        elif "bytes" in message:
+            await server_ws.send_bytes(message["bytes"])
+
+
+async def forward_to_client(client_ws: WebSocket, server_ws: AsyncWebSocketSession):
+    while True:
+        message = await server_ws.receive()
+        if message is None:
+            break
+        if isinstance(message, TextMessage):
+            await client_ws.send_text(message.data)
+        elif isinstance(message, BytesMessage):
+            await client_ws.send_bytes(message.data)
+
+
+# NB(nikhil): _handle_websocket_forwarding uses some py311 specific syntax, but in newer
+# versions of truss we're guaranteed to be running the control server with at least that version.
+async def _handle_websocket_forwarding(
+    client_ws: WebSocket, server_ws: AsyncWebSocketSession
+):
+    logger = client_ws.app.state.logger
+    try:
+        async with asyncio.TaskGroup() as tg:  # type: ignore[attr-defined]
+            tg.create_task(forward_to_client(client_ws, server_ws))
+            tg.create_task(forward_to_server(client_ws, server_ws))
+    except ExceptionGroup as eg:  # type: ignore[name-defined] # noqa: F821
+        exc = eg.exceptions[0]  # NB(nikhil): Only care about the first one.
+        if isinstance(exc, WebSocketDisconnect):
+            await _safe_close_ws(client_ws, logger, exc.code, exc.reason)
+        elif isinstance(exc, StartletteWebSocketDisconnect):
+            await _safe_close_ws(server_ws, logger, exc.code, exc.reason)
+        else:
+            logger.warning(f"Ungraceful websocket close: {exc}")
+    finally:
+        await _safe_close_ws(client_ws, logger)
+        await _safe_close_ws(server_ws, logger)
+
+
+async def _attempt_websocket_proxy(
+    client_ws: WebSocket, proxy_client: httpx.AsyncClient, logger
+):
+    async with aconnect_ws("/v1/websocket", proxy_client) as server_ws:  # type: ignore
+        await client_ws.accept()
+        await _handle_websocket_forwarding(client_ws, server_ws)
+
+
 async def proxy_ws(client_ws: WebSocket):
     proxy_client: httpx.AsyncClient = client_ws.app.state.proxy_client
     logger = client_ws.app.state.logger
@@ -132,34 +199,7 @@ async def proxy_ws(client_ws: WebSocket):
     for attempt in inference_retries():
         with attempt:
             try:
-                async with aconnect_ws("/v1/websocket", proxy_client) as server_ws:  # type: ignore
-                    # Unfortunate, but FastAPI and httpx-ws have slightly different abstractions
-                    # for sending data, so it's not easy to create a unified wrapper.
-                    async def forward_to_server():
-                        while True:
-                            message = await client_ws.receive()
-                            if message.get("type") == "websocket.disconnect":
-                                break
-                            if "text" in message:
-                                await server_ws.send_text(message["text"])
-                            elif "bytes" in message:
-                                await server_ws.send_bytes(message["bytes"])
-
-                    async def forward_to_client():
-                        while True:
-                            message = await server_ws.receive()
-                            if message is None:
-                                break
-                            if isinstance(message, TextMessage):
-                                await client_ws.send_text(message.data)
-                            elif isinstance(message, BytesMessage):
-                                await client_ws.send_bytes(message.data)
-
-                    await client_ws.accept()
-                    try:
-                        await asyncio.gather(forward_to_client(), forward_to_server())
-                    finally:
-                        await _safe_close_ws(client_ws, logger)
+                await _attempt_websocket_proxy(client_ws, proxy_client, logger)
             except httpx_ws_exceptions.HTTPXWSException as e:
                 logger.warning(f"WebSocket connection rejected: {e}")
                 await _safe_close_ws(client_ws, logger)

truss/templates/copy_cache_files.Dockerfile.jinja

@@ -1,3 +1,3 @@
 {% for file in cached_files %}
-COPY --from=cache_warmer {{file.source}} {{file.dst}}
+COPY --chown={{ default_owner }} --from=cache_warmer {{file.source}} {{file.dst}}
 {% endfor %}

truss/templates/docker_server/supervisord.conf.jinja

@@ -1,4 +1,5 @@
 [supervisord]
+pidfile=/tmp/supervisord.pid ; Set PID file location to /tmp to be writable by the non-root user
 nodaemon=true                ; Run supervisord in the foreground (useful for containers)
 logfile=/dev/null            ; Disable logging to file (send logs to /dev/null)
 logfile_maxbytes=0           ; No size limit on logfile (since logging is disabled)

truss/templates/server/truss_server.py

@@ -76,7 +76,7 @@ async def parse_body(request: Request) -> bytes:
 
 
 async def _safe_close_websocket(
-    ws: WebSocket, reason: Optional[str], status_code: int = 1000
+    ws: WebSocket, status_code: int = 1000, reason: Optional[str] = None
 ) -> None:
     try:
         await ws.close(code=status_code, reason=reason)
@@ -257,14 +257,14 @@ class BasetenEndpoints:
                 try:
                     await ws.accept()
                     await self._model.websocket(ws)
-                    await _safe_close_websocket(ws, None, status_code=1000)
+                    await _safe_close_websocket(ws, status_code=1000, reason=None)
                 except WebSocketDisconnect as ws_error:
                     logging.info(
                         f"Client terminated websocket connection: `{ws_error}`."
                     )
                 except Exception:
                     await _safe_close_websocket(
-                        ws, errors.MODEL_ERROR_MESSAGE, status_code=1011
+                        ws, status_code=1011, reason=errors.MODEL_ERROR_MESSAGE
                     )
                     raise  # Re raise to let `intercept_exceptions` deal with it.