diff --git a/README.md b/README.md index fdfee61..6f1ad97 100644 --- a/README.md +++ b/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/ ``` diff --git a/ai-studio-cli/cmd/bench.go b/ai-studio-cli/cmd/bench.go index ae578e0..5ad491b 100644 --- a/ai-studio-cli/cmd/bench.go +++ b/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,6 +149,8 @@ func (r *benchRunner) runBenchmark(cmd *cobra.Command, composePath string) error return err } + r.ensureTokenizer(composePath, modelName) + gpu := r.gpuTag if gpu == "" { gpu = detectGPUTag() @@ -132,8 +158,6 @@ func (r *benchRunner) runBenchmark(cmd *cobra.Command, composePath string) error 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() -} \ No newline at end of file +} diff --git a/ai-studio-cli/cmd/bench_ui.go b/ai-studio-cli/cmd/bench_ui.go index 5c75a95..3c6bbb2 100644 --- a/ai-studio-cli/cmd/bench_ui.go +++ b/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,7 +111,7 @@ 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) @@ -99,6 +119,18 @@ func runBenchUI() error { 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 { diff --git a/ai-studio-cli/cmd/compose.go b/ai-studio-cli/cmd/compose.go index 7d2e771..a9cb355 100644 --- a/ai-studio-cli/cmd/compose.go +++ b/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 } diff --git a/ai-studio-cli/internal/provision/node.go b/ai-studio-cli/internal/provision/node.go index a16ec8a..f806fea 100644 --- a/ai-studio-cli/internal/provision/node.go +++ b/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) } diff --git a/install.sh b/install.sh new file mode 100644 index 0000000..07cf2b9 --- /dev/null +++ b/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"