Add UI links after benchmarking completion

README.md

@@ -4,21 +4,25 @@ A command-line tool for node provisioning (CUDA toolkit, Nvidia recommended driv
 
 Supports all major NVIDIA GPU architectures including the **CoreSpan 5090 Inference System** (Blackwell / RTX 5090, GB202). The provisioning script automatically detects the RTX 5090, selects CUDA 12.8, and installs driver 570+.
 
-## Building from Source
+## Installation
 
-```bash
-go mod tidy
+### Quick Install (from GitHub Releases)
 
-CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -ldflags="-s -w" -o ai-studio-cli .
+```bash
+curl -sL https://raw.githubusercontent.com/corespan/aistudio-cli/master/install.sh | bash
 ```
 
----
+This downloads the latest release binary for your platform and installs it to `/usr/local/bin`. Requires `curl`, `tar`, and `sudo`.
 
-## Installation
-
-After building, move the binary to `/usr/local/bin` to use it globally:
+### Building from Source
 
 ```bash
+cd ai-studio-cli
+
+go mod tidy
+
+CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -ldflags="-s -w" -o ai-studio-cli .
+
 sudo mv ai-studio-cli /usr/local/bin/
 ```
 

ai-studio-cli/cmd/bench.go

@@ -13,6 +13,7 @@ import (
 	"time"
 
 	"github.com/spf13/cobra"
+	"golang.org/x/term"
 )
 
 // serverConfig groups the fields that control server lifecycle and Docker image selection.
@@ -49,6 +50,8 @@ type benchRunner struct {
 	metadata          []string
 	additionalArgs    []string
 	gpuTag            string
+	openUI            bool
+	uiPort            int
 }
 
 var bench = &benchRunner{}
@@ -71,9 +74,9 @@ func (r *benchRunner) registerFlags(cmd *cobra.Command) {
 	cmd.Flags().StringVar(&r.endpoint, "endpoint", "http://localhost:8010", "vLLM base URL or full endpoint URL")
 	cmd.Flags().StringVar(&r.apiEndpoint, "api-endpoint", "/v1/completions", "API path for vLLM benchmark requests")
 	cmd.Flags().StringVar(&r.model, "model", "", "Model name to request; if empty, auto-detected from /v1/models")
-	cmd.Flags().IntVar(&r.concurrency, "concurrency", 128, "Maximum concurrent benchmark requests (match the server's --max-num-seqs)")
-	cmd.Flags().IntVar(&r.requests, "requests", 1000, "Total benchmark prompts to send")
-	cmd.Flags().IntVar(&r.maxTokens, "max-tokens", 256, "Target output token length")
+	cmd.Flags().IntVar(&r.concurrency, "concurrency", 20, "Maximum concurrent benchmark requests")
+	cmd.Flags().IntVar(&r.requests, "requests", 200, "Total benchmark prompts to send")
+	cmd.Flags().IntVar(&r.maxTokens, "max-tokens", 1024, "Target output token length")
 	cmd.Flags().StringVar(&r.dataset, "dataset", "random", "Dataset name for vLLM bench: random, sonnet, sharegpt, custom, hf, etc.")
 	cmd.Flags().StringVar(&r.datasetPath, "dataset-path", "", "Dataset path for datasets that require one")
 	cmd.Flags().IntVar(&r.warmup, "warmup", 2, "Number of warmup requests")
@@ -97,21 +100,42 @@ func (r *benchRunner) registerFlags(cmd *cobra.Command) {
 	cmd.Flags().StringVar(&r.server.composeFile, "compose-file", "", "Path to a custom docker-compose.yml; uses bundled default if server is unreachable and this is unset")
 	cmd.Flags().BoolVar(&r.server.keepServer, "keep-server", false, "Leave the vLLM server running after the benchmark completes")
 	cmd.Flags().IntVar(&r.server.timeoutSec, "server-timeout", 900, "Maximum seconds to wait for the vLLM server to become ready (default allows for large model load times)")
+	// Dashboard
+	cmd.Flags().BoolVar(&r.openUI, "ui", true, "On an interactive terminal, launch the results dashboard when the run finishes and print a clickable link (blocks until Ctrl+C); auto-skipped for non-interactive/CI runs, or pass --ui=false")
+	cmd.Flags().IntVar(&r.uiPort, "ui-port", 9090, "Port for the dashboard launched by --ui")
 }
 
 // entry point
 func (r *benchRunner) run(cmd *cobra.Command) error {
-	// Check compose file
 	effectiveCompose, cleanup, err := serverLifecycle(r.server.composeFile, r.endpoint, r.server.timeoutSec, r.server.keepServer)
 	if err != nil {
 		return err
 	}
-	defer cleanup()
 
-	return r.runBenchmark(cmd, effectiveCompose)
+	cleanedUp := false
+	teardown := func() {
+		if !cleanedUp {
+			cleanedUp = true
+			cleanup()
+		}
+	}
+	defer teardown()
+
+	if err := r.runBenchmark(cmd, effectiveCompose); err != nil {
+		return err
+	}
+
+	teardown()
+
+	if r.openUI && term.IsTerminal(int(os.Stdout.Fd())) {
+		fmt.Printf("\nLaunching results dashboard (Ctrl+C to stop)...\n")
+		if err := serveBenchUI(r.resultDir, r.uiPort, false); err != nil {
+			fmt.Fprintf(os.Stderr, "dashboard failed to start: %v\n", err)
+		}
+	}
+	return nil
 }
 
-// runBenchmark contains the core benchmarking logic, decoupled from lifecycle management.
 func (r *benchRunner) runBenchmark(cmd *cobra.Command, composePath string) error {
 	apiEndpointChanged := cmd.Flags().Changed("api-endpoint")
 	host, port, apiEndpoint, err := splitBenchmarkEndpoint(r.endpoint, r.apiEndpoint, apiEndpointChanged)
@@ -125,15 +149,15 @@ func (r *benchRunner) runBenchmark(cmd *cobra.Command, composePath string) error
 		return err
 	}
 
+	r.ensureTokenizer(composePath, modelName)
+
 	gpu := r.gpuTag
 	if gpu == "" {
 		gpu = detectGPUTag()
 	}
 
 	timestamp := time.Now().Format("2006-01-02T15-04-05")
 	structuredDir := filepath.Join(r.resultDir, sanitizePath(modelName), sanitizePath(gpu), timestamp)
-	// Absolute so the path matches the Docker bind mount (absDir:absDir); a
-	// relative --result-dir resolves against the container WORKDIR and is lost.
 	if abs, err := filepath.Abs(structuredDir); err == nil {
 		structuredDir = abs
 	}
@@ -165,6 +189,27 @@ func (r *benchRunner) runBenchmark(cmd *cobra.Command, composePath string) error
 	return nil
 }
 
+func (r *benchRunner) ensureTokenizer(composePath, modelName string) {
+	if composePath == "" || hasTokenizerArg(r.additionalArgs) {
+		return
+	}
+	modelPath := ExtractModelPath(composePath)
+	if modelPath == "" || modelPath == modelName {
+		return
+	}
+	r.additionalArgs = append(r.additionalArgs, "--tokenizer="+modelPath)
+	fmt.Printf("Served model %q is an alias; using tokenizer %q from compose file.\n", modelName, modelPath)
+}
+
+func hasTokenizerArg(args []string) bool {
+	for _, a := range args {
+		if a == "--tokenizer" || strings.HasPrefix(a, "--tokenizer=") {
+			return true
+		}
+	}
+	return false
+}
+
 func (r *benchRunner) resolveDatasetName() string {
 	if r.promptStr != "" {
 		return "custom"
@@ -457,4 +502,4 @@ func ensureImagePulled(ctx context.Context, image string) error {
 	pull.Stdout = os.Stdout
 	pull.Stderr = os.Stderr
 	return pull.Run()
-}
+}

ai-studio-cli/cmd/bench_ui.go

@@ -3,6 +3,7 @@ package cmd
 import (
 	"context"
 	"fmt"
+	"net"
 	"net/http"
 	"os"
 	"os/exec"
@@ -45,14 +46,30 @@ func init() {
 }
 
 func runBenchUI() error {
-	router := benchui.NewRouter(uiResultDir)
+	return serveBenchUI(uiResultDir, uiPort, uiAutoOpen)
+}
+
+func serveBenchUI(resultDir string, port int, autoOpen bool) error {
+	router := benchui.NewRouter(resultDir)
+
+	// Get a free port
+	ln, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
+	if err != nil {
+		fmt.Printf("Port %d is in use; picking a free port instead...\n", port)
+		ln, err = net.Listen("tcp", ":0")
+		if err != nil {
+			return fmt.Errorf("could not open a port for the dashboard: %w", err)
+		}
+	}
+	port = ln.Addr().(*net.TCPAddr).Port
 
-	addr := fmt.Sprintf(":%d", uiPort)
-	localURL := fmt.Sprintf("http://localhost:%d", uiPort)
-	networkURL := fmt.Sprintf("http://0.0.0.0:%d", uiPort)
+	localURL := fmt.Sprintf("http://localhost:%d", port)
+	networkURL := ""
+	if ip := outboundIP(); ip != "" {
+		networkURL = fmt.Sprintf("http://%s:%d", ip, port)
+	}
 
 	srv := &http.Server{
-		Addr:         addr,
 		Handler:      router,
 		ReadTimeout:  15 * time.Second,
 		WriteTimeout: 30 * time.Second,
@@ -62,6 +79,7 @@ func runBenchUI() error {
 	// Graceful shutdown
 	stop := make(chan os.Signal, 1)
 	signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM)
+	defer signal.Stop(stop)
 
 	errCh := make(chan error, 1)
 
@@ -70,19 +88,21 @@ func runBenchUI() error {
 		fmt.Printf("  │  AI Studio Benchmark Dashboard          │\n")
 		fmt.Printf("  │                                         │\n")
 		fmt.Printf("  │  ➜  Local:   %-26s│\n", localURL)
-		fmt.Printf("  │  ➜  Network: %-26s│\n", networkURL)
-		fmt.Printf("  │  ➜  Results: %-26s│\n", uiResultDir)
+		if networkURL != "" {
+			fmt.Printf("  │  ➜  Network: %-26s│\n", networkURL)
+		}
+		fmt.Printf("  │  ➜  Results: %-26s│\n", resultDir)
 		fmt.Printf("  │                                         │\n")
 		fmt.Printf("  │  Press Ctrl+C to stop                   │\n")
 		fmt.Printf("  └─────────────────────────────────────────┘\n\n")
 
-		if uiAutoOpen {
+		if autoOpen {
 			openBrowser(localURL)
 		}
 	}()
 
 	go func() {
-		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
+		if err := srv.Serve(ln); err != nil && err != http.ErrServerClosed {
 			errCh <- err
 		}
 	}()
@@ -91,14 +111,26 @@ func runBenchUI() error {
 	case <-stop:
 		fmt.Println("\nShutting down...")
 	case err := <-errCh:
-		return fmt.Errorf("server failed: %w (is port %d already in use?)", err, uiPort)
+		return fmt.Errorf("dashboard server failed: %w", err)
 	}
 
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
 	return srv.Shutdown(ctx)
 }
 
+func outboundIP() string {
+	conn, err := net.Dial("udp", "8.8.8.8:80")
+	if err != nil {
+		return ""
+	}
+	defer conn.Close()
+	if addr, ok := conn.LocalAddr().(*net.UDPAddr); ok {
+		return addr.IP.String()
+	}
+	return ""
+}
+
 func openBrowser(url string) {
 	var cmd *exec.Cmd
 	switch runtime.GOOS {

ai-studio-cli/cmd/compose.go

@@ -31,6 +31,38 @@ type composeConfig struct {
 	Services map[string]composeService `yaml:"services"`
 }
 
+func ExtractModelPath(composePath string) string {
+	if composePath == "" {
+		return ""
+	}
+	data, err := os.ReadFile(composePath)
+	if err != nil {
+		return ""
+	}
+	var cf composeConfig
+	if err := yaml.Unmarshal(data, &cf); err != nil {
+		return ""
+	}
+	svc, _, found := findVLLMService(cf.Services)
+	if !found && len(cf.Services) > 0 {
+		svc = cf.Services[sortedServiceNames(cf.Services)[0]]
+	}
+	tokens := flattenCommand(svc.Command)
+	for i, tok := range tokens {
+		flagName, inlineVal, hasInline := strings.Cut(tok, "=")
+		if flagName != "--model" {
+			continue
+		}
+		if hasInline {
+			return strings.TrimSpace(inlineVal)
+		}
+		if i+1 < len(tokens) {
+			return strings.TrimSpace(tokens[i+1])
+		}
+	}
+	return ""
+}
+
 // ExtractModelConfig reads a docker-compose file and returns model config.
 func ExtractModelConfig(composePath string) (ModelConfig, []string) {
 	cfg := ModelConfig{TensorParallel: 1, PipelineParallel: 1}
@@ -89,27 +121,29 @@ func sortedServiceNames(services map[string]composeService) []string {
 }
 
 func findVLLMService(services map[string]composeService) (composeService, string, bool) {
-	// A service literally named "vllm" is the serving one (a sibling
-	// "benchmark" service often shares the same image but has no command).
-	if svc, ok := services["vllm"]; ok {
+	// Helper: does this service look like vLLM based on image or command content?
+	looksLikeVLLM := func(svc composeService) bool {
+		img := strings.ToLower(svc.Image)
+		cmdParts := flattenCommand(svc.Command)
+		cmd := strings.ToLower(strings.Join(cmdParts, " "))
+		return strings.Contains(img, "vllm") ||
+			strings.Contains(cmd, "vllm") ||
+			strings.Contains(cmd, "vllm.entrypoints")
+	}
+
+	if svc, ok := services["vllm"]; ok && looksLikeVLLM(svc) {
 		return svc, "vllm", true
 	}
 
-	// Otherwise pick a vllm-looking service, preferring one that actually
-	// carries a command (the flags we need to parse) over an empty one.
 	var fallback composeService
 	var fallbackName string
 	found := false
 	for _, name := range sortedServiceNames(services) {
 		svc := services[name]
-		cmdParts := flattenCommand(svc.Command)
-		cmd := strings.ToLower(strings.Join(cmdParts, " "))
-		img := strings.ToLower(svc.Image)
-		if !strings.Contains(img, "vllm") &&
-			!strings.Contains(cmd, "vllm") &&
-			!strings.Contains(cmd, "vllm.entrypoints") {
+		if !looksLikeVLLM(svc) {
 			continue
 		}
+		cmdParts := flattenCommand(svc.Command)
 		if len(cmdParts) > 0 {
 			return svc, name, true
 		}

ai-studio-cli/internal/provision/node.go

@@ -1,6 +1,7 @@
 package provision
 
 import (
+	"bytes"
 	_ "embed"
 	"fmt"
 	"os"
@@ -127,7 +128,9 @@ func (p *Provisioner) SetupVLLMDeps() error {
 
 	scriptPath := "/tmp/installVllmDeps.sh"
 	fmt.Println("Staging vLLM deps script to /tmp...")
-	if err := os.WriteFile(scriptPath, installVllmDepsScript, 0755); err != nil {
+
+	script := bytes.ReplaceAll(installVllmDepsScript, []byte("\r\n"), []byte("\n"))
+	if err := os.WriteFile(scriptPath, script, 0755); err != nil {
 		return fmt.Errorf("failed to stage script: %v", err)
 	}
 

install.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+set -e
+
+REPO="corespan/aistudio-cli"
+BINARY_NAME="ai-studio-cli"
+
+# 1. Get the latest release tag
+TAG=$(curl -fsSL "https://api.github.com/repos/${REPO}/releases/latest" | grep '"tag_name"' | head -1 | sed -E 's/.*"([^"]+)".*/\1/')
+echo "Latest release: ${TAG}"
+
+# 2. Download the tarball
+TARBALL="${TAG}.tar.gz"
+echo "Downloading ${TARBALL}..."
+curl -fsSL "https://github.com/${REPO}/releases/download/${TAG}/${TARBALL}" -o "/tmp/${TARBALL}"
+
+# 3. Extract and install
+tar -xzf "/tmp/${TARBALL}" -C /tmp
+chmod +x "/tmp/${BINARY_NAME}"
+echo "Installing to /usr/local/bin/..."
+sudo mv "/tmp/${BINARY_NAME}" "/usr/local/bin/${BINARY_NAME}"
+rm -f "/tmp/${TARBALL}"
+
+echo "Done. Run: ${BINARY_NAME} --help"