Skip to content
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ build-cli:
$(MAKE) -C cmd/cli

build-dmr:
go build -ldflags="-s -w" -o dmr ./cmd/dmr
CGO_ENABLED=1 go build -ldflags="-s -w -X main.Version=$(shell git describe --tags --always --dirty --match 'v*')" -o dmr ./cmd/dmr

build-llamacpp:
git submodule update --init llamacpp/native
Expand Down
147 changes: 43 additions & 104 deletions cmd/dmr/main.go
Original file line number Diff line number Diff line change
@@ -1,130 +1,69 @@
// dmr is a developer convenience wrapper that starts the model-runner server on
// a free port and runs a model-cli command against it in one step.
//
// Usage: dmr <cli-args...>
//
// Example: dmr run qwen3:0.6B-Q4_0 tell me today's news
// Command dmr is the unified Docker Model Runner daemon and client.
package main

import (
"errors"
"context"
"fmt"
"net"
"net/http"
"os"
"os/exec"
"os/signal"
"path/filepath"
"strconv"
"syscall"
"time"
)

func freePort() (int, error) {
l, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
return 0, err
}
defer l.Close()
return l.Addr().(*net.TCPAddr).Port, nil
}
"github.com/docker/cli/cli/command"
"github.com/docker/model-runner/cmd/cli/commands"
"github.com/docker/model-runner/pkg/server"
"github.com/spf13/cobra"
)

func waitForServer(url string, timeout time.Duration) error {
client := &http.Client{Timeout: time.Second}
deadline := time.Now().Add(timeout)
for time.Now().Before(deadline) {
resp, err := client.Get(url)
if err == nil {
resp.Body.Close()
if resp.StatusCode == http.StatusOK {
return nil
}
}
time.Sleep(200 * time.Millisecond)
}
return fmt.Errorf("server not ready after %s", timeout)
}
var Version = "dev"

func checkBinary(path, name, expectedLayout string) error {
if _, err := os.Stat(path); os.IsNotExist(err) {
return fmt.Errorf("missing %s binary at %s\n\nExpected directory layout:\n%s\n\nPlease run 'make build' to build all binaries", name, path, expectedLayout)
}
return nil
}
const defaultHost = "http://localhost:12434"
const defaultPort = "12434"

func main() {
self, err := os.Executable()
if err != nil {
fmt.Fprintf(os.Stderr, "dmr: %v\n", err)
os.Exit(1)
}
dir := filepath.Dir(self)

serverBin := filepath.Join(dir, "model-runner")
cliBin := filepath.Join(dir, "cmd", "cli", "model-cli")

expectedLayout := fmt.Sprintf(`%s/
├── model-runner (server binary)
├── dmr (this wrapper)
└── cmd/
└── cli/
└── model-cli (CLI binary)`, dir)

if err := checkBinary(serverBin, "model-runner", expectedLayout); err != nil {
fmt.Fprintf(os.Stderr, "dmr: %v\n", err)
os.Exit(1)
}
if err := checkBinary(cliBin, "model-cli", expectedLayout); err != nil {
fmt.Fprintf(os.Stderr, "dmr: %v\n", err)
if err := run(); err != nil {
fmt.Fprintln(os.Stderr, err)
os.Exit(1)
}
}

port, err := freePort()
func run() error {
cli, err := command.NewDockerCli()
if err != nil {
fmt.Fprintf(os.Stderr, "dmr: failed to find free port: %v\n", err)
os.Exit(1)
return fmt.Errorf("unable to initialize CLI: %w", err)
}
portStr := strconv.Itoa(port)
serverURL := "http://localhost:" + portStr

fmt.Fprintf(os.Stderr, "dmr: starting model-runner on port %d\n", port)

server := exec.Command(serverBin)
server.Env = append(os.Environ(), "MODEL_RUNNER_PORT="+portStr)
server.Stderr = os.Stderr
server.Stdout = os.Stdout
root := commands.NewRootCmd(cli)
root.Use = "dmr"
root.Short = "Docker Model Runner"

if err := server.Start(); err != nil {
fmt.Fprintf(os.Stderr, "dmr: failed to start model-runner: %v\n", err)
os.Exit(1)
if os.Getenv("MODEL_RUNNER_HOST") == "" {
if err := os.Setenv("MODEL_RUNNER_HOST", defaultHost); err != nil {
return fmt.Errorf("unable to set MODEL_RUNNER_HOST: %w", err)
}
}
defer server.Process.Kill()

sigCh := make(chan os.Signal, 1)
signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
go func() {
<-sigCh
server.Process.Kill()
}()
root.AddCommand(newServeCmd())

if err := waitForServer(serverURL+"/", 30*time.Second); err != nil {
fmt.Fprintf(os.Stderr, "dmr: %v\n", err)
os.Exit(1)
}
return root.Execute()
}

// #nosec G702 - Intentional: dmr is a CLI wrapper that forwards arguments to model-cli
cli := exec.Command(cliBin, os.Args[1:]...)
cli.Env = append(os.Environ(), "MODEL_RUNNER_HOST="+serverURL)
cli.Stdin = os.Stdin
cli.Stdout = os.Stdout
cli.Stderr = os.Stderr
func newServeCmd() *cobra.Command {
return &cobra.Command{
Use: "serve",
Short: "Start the Docker Model Runner daemon",
// skip Docker CLI init; serve runs the daemon directly
PersistentPreRunE: func(*cobra.Command, []string) error { return nil },
RunE: func(cmd *cobra.Command, _ []string) error {
if os.Getenv("MODEL_RUNNER_PORT") == "" {
if err := os.Setenv("MODEL_RUNNER_PORT", defaultPort); err != nil {
return fmt.Errorf("unable to set MODEL_RUNNER_PORT: %w", err)
}
}

if err := cli.Run(); err != nil {
var exitErr *exec.ExitError
if errors.As(err, &exitErr) {
os.Exit(exitErr.ExitCode())
}
fmt.Fprintf(os.Stderr, "dmr: %v\n", err)
os.Exit(1)
ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
defer cancel()

return server.Run(ctx, server.Config{Version: Version})
},
}
}
Loading