Fix runner server not respecting max code upload size

Author: colinjc

runner/cmd/runner/cmd.go

@@ -15,6 +15,7 @@ func App() {
 	var httpPort int
 	var sshPort int
 	var logLevel int
+	var codeUploadLimit int64
 
 	app := &cli.App{
 		Name:    "dstack-runner",
@@ -28,6 +29,13 @@ func App() {
 				Usage:       "log verbosity level: 2 (Error), 3 (Warning), 4 (Info), 5 (Debug), 6 (Trace)",
 				Destination: &logLevel,
 			},
+			&cli.Int64Flag{
+				Name:        "code-upload-limit",
+				Usage:       "Set the code upload limit",
+				Value:       10 * 1024 * 1024,
+				Destination: &codeUploadLimit,
+				EnvVars:     []string{"DSTACK_RUNNER_CODE_UPLOAD_LIMIT"},
+			},
 		},
 		Commands: []*cli.Command{
 			{
@@ -66,7 +74,7 @@ func App() {
 					},
 				},
 				Action: func(c *cli.Context) error {
-					err := start(paths.tempDir, paths.homeDir, paths.workingDir, httpPort, sshPort, logLevel, Version)
+					err := start(paths.tempDir, paths.homeDir, paths.workingDir, httpPort, sshPort, logLevel, Version, codeUploadLimit)
 					if err != nil {
 						return cli.Exit(err, 1)
 					}

runner/cmd/runner/main.go

@@ -19,7 +19,7 @@ func main() {
 	App()
 }
 
-func start(tempDir string, homeDir string, workingDir string, httpPort int, sshPort int, logLevel int, version string) error {
+func start(tempDir string, homeDir string, workingDir string, httpPort int, sshPort int, logLevel int, version string, codeUploadLimit int64) error {
 	if err := os.MkdirAll(tempDir, 0o755); err != nil {
 		return tracerr.Errorf("Failed to create temp directory: %w", err)
 	}
@@ -38,7 +38,7 @@ func start(tempDir string, homeDir string, workingDir string, httpPort int, sshP
 	log.DefaultEntry.Logger.SetOutput(io.MultiWriter(os.Stdout, defaultLogFile))
 	log.DefaultEntry.Logger.SetLevel(logrus.Level(logLevel))
 
-	server, err := api.NewServer(tempDir, homeDir, workingDir, fmt.Sprintf(":%d", httpPort), sshPort, version)
+	server, err := api.NewServer(tempDir, homeDir, workingDir, fmt.Sprintf(":%d", httpPort), sshPort, version, codeUploadLimit)
 	if err != nil {
 		return tracerr.Errorf("Failed to create server: %w", err)
 	}

runner/internal/runner/api/http.go

@@ -3,6 +3,7 @@ package api
 import (
 	"context"
 	"io"
+	"math"
 	"net/http"
 	"os"
 	"path/filepath"
@@ -64,7 +65,11 @@ func (s *Server) uploadCodePostHandler(w http.ResponseWriter, r *http.Request) (
 		return nil, &api.Error{Status: http.StatusConflict}
 	}
 
-	r.Body = http.MaxBytesReader(w, r.Body, 10*1024*1024)
+	if s.codeUploadLimit > 0 {
+		r.Body = http.MaxBytesReader(w, r.Body, s.codeUploadLimit)
+	} else {
+		r.Body = http.MaxBytesReader(w, r.Body, math.MaxInt64)
+	}
 	codePath := filepath.Join(s.tempDir, "code") // todo random name?
 	file, err := os.Create(codePath)
 	if err != nil {

runner/internal/runner/api/server.go

@@ -31,9 +31,11 @@ type Server struct {
 	cancelRun context.CancelFunc
 
 	version string
+
+	codeUploadLimit int64
 }
 
-func NewServer(tempDir string, homeDir string, workingDir string, address string, sshPort int, version string) (*Server, error) {
+func NewServer(tempDir string, homeDir string, workingDir string, address string, sshPort int, version string, codeUploadLimit int64) (*Server, error) {
 	r := api.NewRouter()
 	ex, err := executor.NewRunExecutor(tempDir, homeDir, workingDir, sshPort)
 	if err != nil {
@@ -57,6 +59,8 @@ func NewServer(tempDir string, homeDir string, workingDir string, address string
 		executor: ex,
 
 		version: version,
+
+		codeUploadLimit: codeUploadLimit,
 	}
 	r.AddHandler("GET", "/api/healthcheck", s.healthcheckGetHandler)
 	r.AddHandler("GET", "/api/metrics", s.metricsGetHandler)

src/dstack/_internal/core/backends/base/compute.py

@@ -38,6 +38,7 @@
     VolumeProvisioningData,
 )
 from dstack._internal.core.services import is_valid_dstack_resource_name
+from dstack._internal.server.settings import SERVER_CODE_UPLOAD_LIMIT
 from dstack._internal.utils.logging import get_logger
 from dstack._internal.utils.path import PathLike
 
@@ -538,6 +539,7 @@ def get_shim_env(
         "DSTACK_RUNNER_HTTP_PORT": str(DSTACK_RUNNER_HTTP_PORT),
         "DSTACK_RUNNER_SSH_PORT": str(DSTACK_RUNNER_SSH_PORT),
         "DSTACK_RUNNER_LOG_LEVEL": log_level,
+        "DSTACK_RUNNER_CODE_UPLOAD_LIMIT": str(SERVER_CODE_UPLOAD_LIMIT),
         "DSTACK_PUBLIC_SSH_KEY": "\n".join(authorized_keys),
     }
     if backend_shim_env is not None: