Make local code upload size limit configurable (#2673)

Author: colinjc

docs/docs/reference/environment-variables.md

@@ -12,7 +12,7 @@ tasks, and services:
      ```yaml
      type: task
      name: vscode
-     
+
      commands:
        - echo $DSTACK_RUN_NAME
      ```
@@ -25,11 +25,11 @@ tasks, and services:
 - `DSTACK_GPUS_NUM`{ #DSTACK_GPUS_NUM } – The total number of GPUs in the run.
 
      Example:
-     
+
      ```yaml
      type: service
      name: llama31
-     
+
      env:
        - HF_TOKEN
      commands:
@@ -39,7 +39,7 @@ tasks, and services:
          --tensor-parallel-size $DSTACK_GPUS_NUM
      port: 8000
      model: meta-llama/Meta-Llama-3.1-8B-Instruct
-     
+
      resources:
        gpu: 24GB
      ```
@@ -51,7 +51,7 @@ tasks, and services:
 
      Below is an example of using `DSTACK_NODES_NUM`, `DSTACK_GPUS_PER_NODE`, `DSTACK_NODE_RANK`, and `DSTACK_MASTER_NODE_IP`
      for distributed training:
-     
+
      ```yaml
       type: task
       name: train-distrib
@@ -83,18 +83,18 @@ tasks, and services:
 The following environment variables are supported by the `dstack` server and can be specified whether the server is run
 via `dstack server` or deployed using Docker.
 
-For more details on the options below, refer to the [server deployment](../guides/server-deployment.md) guide. 
+For more details on the options below, refer to the [server deployment](../guides/server-deployment.md) guide.
 
 - `DSTACK_SERVER_LOG_LEVEL`{ #DSTACK_SERVER_LOG_LEVEL } – Has the same effect as `--log-level`. Defaults to `INFO`.
 
      Example:
-     
+
      <div class="termy">
-     
+
      ```shell
      $ DSTACK_SERVER_LOG_LEVEL=debug dstack server
      ```
-     
+
      </div>
 
 - `DSTACK_SERVER_LOG_FORMAT`{ #DSTACK_SERVER_LOG_FORMAT } – Sets format of log output. Can be `rich`, `standard`, `json`. Defaults to `rich`.
@@ -110,9 +110,10 @@ For more details on the options below, refer to the [server deployment](../guide
 - `DSTACK_ENABLE_PROMETHEUS_METRICS`{ #DSTACK_ENABLE_PROMETHEUS_METRICS } — Enables Prometheus metrics collection and export.
 - `DSTACK_DEFAULT_SERVICE_CLIENT_MAX_BODY_SIZE`{ #DSTACK_DEFAULT_SERVICE_CLIENT_MAX_BODY_SIZE } – Request body size limit for services running with a gateway, in bytes. Defaults to 64 MiB.
 - `DSTACK_FORBID_SERVICES_WITHOUT_GATEWAY`{ #DSTACK_FORBID_SERVICES_WITHOUT_GATEWAY } – Forbids registering new services without a gateway if set to any value.
+- `DSTACK_SERVER_CODE_UPLOAD_LIMIT`{ #DSTACK_SERVER_CODE_UPLOAD_LIMIT } - The repo size limit when uploading diffs or local repos, in bytes. Set to 0 to disable size limits. Defaults to 2MiB.
 
 ??? info "Internal environment variables"
-     The following environment variables are intended for development purposes: 
+     The following environment variables are intended for development purposes:
 
      * `DSTACK_SERVER_ROOT_LOG_LEVEL` – Sets root logger log level. Defaults to `ERROR`.
      * `DSTACK_SERVER_UVICORN_LOG_LEVEL` – Sets uvicorn logger log level. Defaults to `ERROR`.

pyproject.toml

@@ -34,6 +34,7 @@ dependencies = [
     "psutil",
     "gpuhunt==0.1.5",
     "argcomplete>=3.5.0",
+    "humanize>=4.12.3",
 ]
 
 [project.urls]

src/dstack/_internal/server/routers/repos.py

@@ -1,6 +1,7 @@
 from typing import List, Tuple
 
 from fastapi import APIRouter, Depends, Request, UploadFile
+from humanize import naturalsize
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from dstack._internal.core.errors import ResourceNotExistsError, ServerClientError
@@ -14,9 +15,10 @@
 )
 from dstack._internal.server.security.permissions import ProjectMember
 from dstack._internal.server.services import repos
+from dstack._internal.server.settings import SERVER_CODE_UPLOAD_LIMIT
 from dstack._internal.server.utils.routers import (
     get_base_api_additional_responses,
-    request_size_exceeded,
+    get_request_size,
 )
 
 router = APIRouter(
@@ -94,10 +96,12 @@ async def upload_code(
     session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()),
 ):
-    if request_size_exceeded(request, limit=2 * 2**20):
+    request_size = get_request_size(request)
+    if SERVER_CODE_UPLOAD_LIMIT > 0 and request_size > SERVER_CODE_UPLOAD_LIMIT:
         raise ServerClientError(
-            "Repo diff size exceeds the limit of 2MB. "
-            "Use .gitignore to exclude large files from the repo."
+            f"Repo diff size is {naturalsize(request_size)}, which exceeds the limit of "
+            f"{naturalsize(SERVER_CODE_UPLOAD_LIMIT)}. Use .gitignore to exclude large files from the repo. This "
+            f"limit can be modified by setting the DSTACK_SERVER_CODE_UPLOAD_LIMIT_BYTES environment variable"
         )
     _, project = user_project
     await repos.upload_code(

src/dstack/_internal/server/settings.py

@@ -85,6 +85,7 @@
 USER_PROJECT_DEFAULT_QUOTA = int(os.getenv("DSTACK_USER_PROJECT_DEFAULT_QUOTA", 10))
 FORBID_SERVICES_WITHOUT_GATEWAY = os.getenv("DSTACK_FORBID_SERVICES_WITHOUT_GATEWAY") is not None
 
+SERVER_CODE_UPLOAD_LIMIT = int(os.getenv("DSTACK_SERVER_CODE_UPLOAD_LIMIT", 2 * 2**20))
 
 # Development settings
 

src/dstack/_internal/server/utils/routers.py

@@ -93,13 +93,10 @@ def get_server_client_error_details(error: ServerClientError) -> List[Dict]:
     return details
 
 
-def request_size_exceeded(request: Request, limit: int) -> bool:
+def get_request_size(request: Request) -> int:
     if "content-length" not in request.headers:
-        return True
-    content_length = int(request.headers["content-length"])
-    if content_length > limit:
-        return True
-    return False
+        return 0
+    return int(request.headers["content-length"])
 
 
 def check_client_server_compatibility(