Add /set command and context length configuration

Add a /set command to the interactive mode to allow users to configure
parameters like num_ctx during runtime. Also implement environment
variable support for default context length (DMR_CONTEXT_LENGTH) and
ensure the scheduler respects these configurations when setting up
backend runners.

Signed-off-by: Eric Curtin <eric.curtin@docker.com>

Author: ericcurtin

cmd/cli/commands/run.go

@@ -8,13 +8,16 @@ import (
 	"io"
 	"os"
 	"os/signal"
+	"strconv"
 	"strings"
 	"syscall"
 
 	"github.com/charmbracelet/glamour"
 	"github.com/docker/model-runner/cmd/cli/commands/completion"
 	"github.com/docker/model-runner/cmd/cli/desktop"
 	"github.com/docker/model-runner/cmd/cli/readline"
+	"github.com/docker/model-runner/pkg/inference"
+	"github.com/docker/model-runner/pkg/inference/scheduling"
 	"github.com/fatih/color"
 	"github.com/muesli/termenv"
 	"github.com/spf13/cobra"
@@ -92,6 +95,7 @@ func generateInteractiveWithReadline(cmd *cobra.Command, desktopClient *desktop.
 		fmt.Fprintln(os.Stderr, "Available Commands:")
 		fmt.Fprintln(os.Stderr, "  /set system     Set or update the system message")
 		fmt.Fprintln(os.Stderr, "  /bye            Exit")
+		fmt.Fprintln(os.Stderr, "  /set            Set a parameter (e.g., /set parameter num_ctx 4096)")
 		fmt.Fprintln(os.Stderr, "  /?, /help       Help for a command")
 		fmt.Fprintln(os.Stderr, "  /? shortcuts    Help for keyboard shortcuts")
 		fmt.Fprintln(os.Stderr, "  /? files        Help for file inclusion with @ symbol")
@@ -231,6 +235,36 @@ func generateInteractiveWithReadline(cmd *cobra.Command, desktopClient *desktop.
 			continue
 		case strings.HasPrefix(line, "/exit"), strings.HasPrefix(line, "/bye"):
 			return nil
+		case strings.HasPrefix(line, "/set"):
+			args := strings.Fields(line)
+			if len(args) < 4 || args[1] != "parameter" {
+				fmt.Fprintln(os.Stderr, "Usage: /set parameter <name> <value>")
+				fmt.Fprintln(os.Stderr, "Available parameters: num_ctx")
+				continue
+			}
+			paramName, paramValue := args[2], args[3]
+			switch paramName {
+			case "num_ctx":
+				if val, err := strconv.ParseInt(paramValue, 10, 32); err == nil && val > 0 {
+					ctx := int32(val)
+					if err := desktopClient.ConfigureBackend(scheduling.ConfigureRequest{
+						Model: model,
+						BackendConfiguration: inference.BackendConfiguration{
+							ContextSize: &ctx,
+						},
+					}); err != nil {
+						fmt.Fprintf(os.Stderr, "Failed to set num_ctx: %v\n", err)
+					} else {
+						fmt.Fprintf(os.Stderr, "Set num_ctx to %d\n", val)
+					}
+				} else {
+					fmt.Fprintf(os.Stderr, "Invalid value for num_ctx: %s (must be a positive integer)\n", paramValue)
+				}
+			default:
+				fmt.Fprintf(os.Stderr, "Unknown parameter: %s\n", paramName)
+				fmt.Fprintln(os.Stderr, "Available parameters: num_ctx")
+			}
+			continue
 		case strings.HasPrefix(line, "/"):
 			fmt.Printf("Unknown command '%s'. Type /? for help\n", strings.Fields(line)[0])
 			continue

main.go

@@ -7,6 +7,7 @@ import (
 	"os"
 	"os/signal"
 	"path/filepath"
+	"strconv"
 	"strings"
 	"syscall"
 	"time"
@@ -77,6 +78,18 @@ func main() {
 	sglangServerPath := os.Getenv("SGLANG_SERVER_PATH")
 	mlxServerPath := os.Getenv("MLX_SERVER_PATH")
 
+	// Parse default context length from environment
+	var defaultContextLength *int32
+	if ctxStr := os.Getenv("DMR_CONTEXT_LENGTH"); ctxStr != "" {
+		if parsed, err := strconv.ParseInt(ctxStr, 10, 32); err == nil && parsed > 0 {
+			ctx := int32(parsed)
+			defaultContextLength = &ctx
+			log.Infof("DMR_CONTEXT_LENGTH: %d", ctx)
+		} else {
+			log.Warnf("Invalid DMR_CONTEXT_LENGTH: %s (must be a positive integer)", ctxStr)
+		}
+	}
+
 	// Create a proxy-aware HTTP transport
 	// Use a safe type assertion with fallback, and explicitly set Proxy to http.ProxyFromEnvironment
 	var baseTransport *http.Transport
@@ -175,6 +188,7 @@ func main() {
 			"",
 			false,
 		),
+		defaultContextLength,
 	)
 
 	// Create the HTTP handler for the scheduler

pkg/inference/scheduling/scheduler.go

@@ -40,6 +40,8 @@ type Scheduler struct {
 	tracker *metrics.Tracker
 	// openAIRecorder is used to record OpenAI API inference requests and responses.
 	openAIRecorder *metrics.OpenAIRecorder
+	// defaultContextLength is the default context length from environment variable.
+	defaultContextLength *int32
 }
 
 // NewScheduler creates a new inference scheduler.
@@ -50,19 +52,21 @@ func NewScheduler(
 	modelManager *models.Manager,
 	httpClient *http.Client,
 	tracker *metrics.Tracker,
+	defaultContextLength *int32,
 ) *Scheduler {
 	openAIRecorder := metrics.NewOpenAIRecorder(log.WithField("component", "openai-recorder"), modelManager)
 
 	// Create the scheduler.
 	s := &Scheduler{
-		log:            log,
-		backends:       backends,
-		defaultBackend: defaultBackend,
-		modelManager:   modelManager,
-		installer:      newInstaller(log, backends, httpClient),
-		loader:         newLoader(log, backends, modelManager, openAIRecorder),
-		tracker:        tracker,
-		openAIRecorder: openAIRecorder,
+		log:                  log,
+		backends:             backends,
+		defaultBackend:       defaultBackend,
+		modelManager:         modelManager,
+		installer:            newInstaller(log, backends, httpClient),
+		loader:               newLoader(log, backends, modelManager, openAIRecorder),
+		tracker:              tracker,
+		openAIRecorder:       openAIRecorder,
+		defaultContextLength: defaultContextLength,
 	}
 
 	// Scheduler successfully initialized.
@@ -253,7 +257,12 @@ func (s *Scheduler) ConfigureRunner(ctx context.Context, backend inference.Backe
 
 	// Build runner configuration with shared settings
 	var runnerConfig inference.BackendConfiguration
-	runnerConfig.ContextSize = req.ContextSize
+	// Use request context size if provided, otherwise fall back to default from env var
+	if req.ContextSize != nil {
+		runnerConfig.ContextSize = req.ContextSize
+	} else if s.defaultContextLength != nil {
+		runnerConfig.ContextSize = s.defaultContextLength
+	}
 	runnerConfig.Speculative = req.Speculative
 	runnerConfig.RuntimeFlags = runtimeFlags
 

pkg/inference/scheduling/scheduler_test.go

@@ -33,7 +33,7 @@ func TestCors(t *testing.T) {
 			discard := logrus.New()
 			discard.SetOutput(io.Discard)
 			log := logrus.NewEntry(discard)
-			s := NewScheduler(log, nil, nil, nil, nil, nil)
+			s := NewScheduler(log, nil, nil, nil, nil, nil, nil)
 			httpHandler := NewHTTPHandler(s, nil, []string{"*"})
 			req := httptest.NewRequest(http.MethodOptions, "http://model-runner.docker.internal"+tt.path, http.NoBody)
 			req.Header.Set("Origin", "docker.com")