Download cli/main.go from p80x/store: direct link, hf CLI and curl.
- Browser
- Download file 13 kB
-
https://huggingface.co/spaces/p80x/store/resolve/main/cli/main.go
- Command line
-
hf download hf://spaces/p80x/store/cli/main.go
-
curl -L -o main.go https://huggingface.co/spaces/p80x/store/resolve/main/cli/main.go
13 kB
| package main | |
| import ( | |
| "bytes" | |
| "encoding/json" | |
| "fmt" | |
| "io" | |
| "mime/multipart" | |
| "net/http" | |
| "os" | |
| "path/filepath" | |
| "strings" | |
| "time" | |
| ) | |
| type Config struct { | |
| ServerURL string `json:"server_url"` | |
| APIKey string `json:"api_key"` | |
| } | |
| func main() { | |
| if len(os.Args) < 2 { | |
| printUsageAndExit() | |
| } | |
| command := os.Args[1] | |
| // Help check | |
| if command == "-h" || command == "--help" || command == "help" { | |
| printUsageAndExit() | |
| } | |
| switch command { | |
| case "save": | |
| // User can run either: hf-save <dirs>... OR hf-save save <dirs>... | |
| // We'll support both, but the primary expectation is: hf-save outputs logs | |
| if len(os.Args) < 3 { | |
| fmt.Println("Error: No directories specified to save.") | |
| os.Exit(1) | |
| } | |
| runSave(os.Args[2:]) | |
| case "list": | |
| // user runs hf-list or hf-save list | |
| date := "" | |
| if len(os.Args) > 2 { | |
| date = os.Args[2] | |
| } | |
| runList(date) | |
| case "mount": | |
| // hf-mount outputs OR hf-mount 2026-06-15 16-35-20_outputs | |
| if len(os.Args) < 3 { | |
| fmt.Println("Error: Please specify what to restore.") | |
| os.Exit(1) | |
| } | |
| runMount(os.Args[2:]) | |
| default: | |
| // If command is not recognized, treat all arguments as directories to save | |
| // This enables direct usage like: hf-save outputs logs | |
| runSave(os.Args[1:]) | |
| } | |
| } | |
| func printUsageAndExit() { | |
| fmt.Println(`hf-save: Extreme fast backup/restore tool for ephemeral GPU instances | |
| Usage: | |
| hf-save <dir1> [dir2]... Save directories to the backup server | |
| hf-list List available backup dates | |
| hf-list <date> List all backups on a specific date (e.g. hf-list 2026-06-15) | |
| hf-mount <dir_name> Restore latest backup containing directory name | |
| hf-mount <date> <run> Restore specific backup run (e.g. hf-mount 2026-06-15 16-35-20_outputs) | |
| Aliases / Subcommands: | |
| hf-save save <dir>... | |
| hf-save list [date] | |
| hf-save mount <args>... | |
| `) | |
| os.Exit(0) | |
| } | |
| func loadConfig() Config { | |
| home, err := os.UserHomeDir() | |
| if err != nil { | |
| fmt.Printf("Error determining user home directory: %v\n", err) | |
| os.Exit(1) | |
| } | |
| configPath := filepath.Join(home, ".config", "hf-save", "config.json") | |
| file, err := os.Open(configPath) | |
| if err != nil { | |
| fmt.Printf("Error: Configuration file not found at %s. Please run the init script first.\n", configPath) | |
| os.Exit(1) | |
| } | |
| defer file.Close() | |
| var cfg Config | |
| if err := json.NewDecoder(file).Decode(&cfg); err != nil { | |
| fmt.Printf("Error parsing configuration: %v\n", err) | |
| os.Exit(1) | |
| } | |
| cfg.ServerURL = strings.TrimSuffix(cfg.ServerURL, "/") | |
| return cfg | |
| } | |
| type uploadFile struct { | |
| relPath string | |
| absPath string | |
| } | |
| func ensureServerAwake(serverURL string) { | |
| client := &http.Client{Timeout: 5 * time.Second} | |
| fmt.Print("Checking connection to backup server... ") | |
| for i := 0; i < 30; i++ { | |
| resp, err := client.Get(serverURL + "/health") | |
| if err == nil { | |
| if resp.StatusCode == http.StatusOK { | |
| resp.Body.Close() | |
| fmt.Println("Connected.") | |
| return | |
| } | |
| resp.Body.Close() | |
| } | |
| if i == 0 { | |
| fmt.Println("\nServer is asleep or starting up. Waking it up (this may take 1-2 minutes)...") | |
| } | |
| fmt.Printf("Waiting for server to wake up... (Attempt %d/30)\r", i+1) | |
| time.Sleep(5 * time.Second) | |
| } | |
| fmt.Println("\nWarning: Could not verify server is awake. Attempting operation anyway...") | |
| } | |
| type progressReader struct { | |
| r io.Reader | |
| total int64 | |
| current int64 | |
| } | |
| func (pr *progressReader) Read(p []byte) (n int, err error) { | |
| n, err = pr.r.Read(p) | |
| if n > 0 { | |
| pr.current += int64(n) | |
| percent := float64(pr.current) / float64(pr.total) * 100 | |
| mbCurrent := float64(pr.current) / (1024 * 1024) | |
| mbTotal := float64(pr.total) / (1024 * 1024) | |
| fmt.Printf("Uploading: %.1f%% (%.1f/%.1f MB)\r", percent, mbCurrent, mbTotal) | |
| } | |
| return | |
| } | |
| type writeCounter struct { | |
| total int64 | |
| current int64 | |
| name string | |
| } | |
| func (wc *writeCounter) Write(p []byte) (int, error) { | |
| n := len(p) | |
| wc.current += int64(n) | |
| percent := float64(wc.current) / float64(wc.total) * 100 | |
| mbCurrent := float64(wc.current) / (1024 * 1024) | |
| mbTotal := float64(wc.total) / (1024 * 1024) | |
| fmt.Printf("Downloading %s: %.1f%% (%.1f/%.1f MB)\r", wc.name, percent, mbCurrent, mbTotal) | |
| return n, nil | |
| } | |
| var quoteEscaper = strings.NewReplacer("\\", "\\\\", "\"", "\\\"") | |
| func escapeQuotes(s string) string { | |
| return quoteEscaper.Replace(s) | |
| } | |
| func calculateUploadSize(prefix string, files []uploadFile, boundary string) (int64, error) { | |
| var totalSize int64 | |
| // 1. Prefix field size | |
| var prefixBuf bytes.Buffer | |
| w := multipart.NewWriter(&prefixBuf) | |
| w.SetBoundary(boundary) | |
| _ = w.WriteField("prefix", prefix) | |
| totalSize += int64(prefixBuf.Len()) | |
| // 2. Files size | |
| for _, file := range files { | |
| var fileHeaderBuf bytes.Buffer | |
| fw := multipart.NewWriter(&fileHeaderBuf) | |
| fw.SetBoundary(boundary) | |
| _, err := fw.CreateFormFile("files", file.relPath) | |
| if err != nil { | |
| return 0, err | |
| } | |
| totalSize += int64(fileHeaderBuf.Len()) | |
| info, err := os.Stat(file.absPath) | |
| if err != nil { | |
| return 0, err | |
| } | |
| totalSize += info.Size() | |
| totalSize += 2 // trailing \r\n after file part data | |
| } | |
| // 3. Closing boundary size (\r\n--boundary--\r\n) | |
| totalSize += int64(len("\r\n--" + boundary + "--\r\n")) | |
| return totalSize, nil | |
| } | |
| // bytes helper import for size calculation | |
| type bytesBufferWriter struct { | |
| bytes.Buffer | |
| } | |
| func runSave(targets []string) { | |
| cfg := loadConfig() | |
| ensureServerAwake(cfg.ServerURL) | |
| for _, target := range targets { | |
| saveSingleTarget(cfg, target) | |
| } | |
| } | |
| func saveSingleTarget(cfg Config, target string) { | |
| cleanDir := filepath.Clean(target) | |
| info, err := os.Stat(cleanDir) | |
| if err != nil { | |
| fmt.Printf("Error: Target '%s' not found, skipping.\n", target) | |
| return | |
| } | |
| var filesToUpload []uploadFile | |
| if !info.IsDir() { | |
| filesToUpload = append(filesToUpload, uploadFile{ | |
| relPath: filepath.Base(cleanDir), | |
| absPath: cleanDir, | |
| }) | |
| } else { | |
| err = filepath.Walk(cleanDir, func(path string, fileInfo os.FileInfo, walkErr error) error { | |
| if walkErr != nil { | |
| return walkErr | |
| } | |
| if fileInfo.IsDir() { | |
| return nil | |
| } | |
| parts := strings.Split(path, string(filepath.Separator)) | |
| for _, part := range parts { | |
| if part == ".git" || part == "node_modules" || part == "__pycache__" || part == ".venv" || part == "venv" { | |
| return nil | |
| } | |
| } | |
| rel, err := filepath.Rel(filepath.Dir(cleanDir), path) | |
| if err == nil { | |
| filesToUpload = append(filesToUpload, uploadFile{ | |
| relPath: filepath.ToSlash(rel), | |
| absPath: path, | |
| }) | |
| } | |
| return nil | |
| }) | |
| if err != nil { | |
| fmt.Printf("Error scanning target %s: %v\n", target, err) | |
| return | |
| } | |
| } | |
| if len(filesToUpload) == 0 { | |
| fmt.Printf("No files found to save in target %s.\n", target) | |
| return | |
| } | |
| dateStr := time.Now().Format("2006-01-02") | |
| timeStr := time.Now().Format("15-04-05") | |
| folderSuffix := filepath.Base(cleanDir) | |
| prefix := fmt.Sprintf("%s/%s_%s", dateStr, timeStr, folderSuffix) | |
| fmt.Printf("\n--- Saving target: %s (%d files) ---\n", target, len(filesToUpload)) | |
| fmt.Printf("Uploading snapshot: %s\n", prefix) | |
| // Create a dummy multipart writer to generate a consistent boundary string | |
| dummyWriter := multipart.NewWriter(nil) | |
| boundary := dummyWriter.Boundary() | |
| totalSize, err := calculateUploadSize(prefix, filesToUpload, boundary) | |
| if err != nil { | |
| fmt.Printf("Failed to calculate upload size: %v\n", err) | |
| return | |
| } | |
| pipeReader, pipeWriter := io.Pipe() | |
| writer := multipart.NewWriter(pipeWriter) | |
| writer.SetBoundary(boundary) | |
| go func() { | |
| defer pipeWriter.Close() | |
| defer writer.Close() | |
| err := writer.WriteField("prefix", prefix) | |
| if err != nil { | |
| return | |
| } | |
| for _, file := range filesToUpload { | |
| fileWriter, err := writer.CreateFormFile("files", file.relPath) | |
| if err != nil { | |
| return | |
| } | |
| f, err := os.Open(file.absPath) | |
| if err != nil { | |
| return | |
| } | |
| _, err = io.Copy(fileWriter, f) | |
| f.Close() | |
| if err != nil { | |
| return | |
| } | |
| } | |
| }() | |
| reqURL := fmt.Sprintf("%s/upload", cfg.ServerURL) | |
| pr := &progressReader{ | |
| r: pipeReader, | |
| total: totalSize, | |
| } | |
| req, err := http.NewRequest(http.MethodPost, reqURL, pr) | |
| if err != nil { | |
| fmt.Printf("Error creating upload request: %v\n", err) | |
| return | |
| } | |
| req.Header.Set("Content-Type", writer.FormDataContentType()) | |
| req.ContentLength = totalSize | |
| if cfg.APIKey != "" { | |
| req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", cfg.APIKey)) | |
| } | |
| client := &http.Client{} | |
| resp, err := client.Do(req) | |
| if err != nil { | |
| fmt.Printf("\nUpload failed: %v\n", err) | |
| return | |
| } | |
| defer resp.Body.Close() | |
| fmt.Println() // print newline after progress percentage | |
| if resp.StatusCode != http.StatusOK { | |
| body, _ := io.ReadAll(resp.Body) | |
| fmt.Printf("Upload failed with status %s: %s\n", resp.Status, string(body)) | |
| return | |
| } | |
| var result map[string]interface{} | |
| if err := json.NewDecoder(resp.Body).Decode(&result); err == nil { | |
| fmt.Printf("Successfully saved %s! (%v files)\n", target, result["file_count"]) | |
| } else { | |
| fmt.Printf("Successfully saved %s!\n", target) | |
| } | |
| } | |
| func runList(date string) { | |
| cfg := loadConfig() | |
| ensureServerAwake(cfg.ServerURL) | |
| reqURL := fmt.Sprintf("%s/list", cfg.ServerURL) | |
| if date != "" { | |
| reqURL += "?date=" + date | |
| } | |
| resp, err := http.Get(reqURL) | |
| if err != nil { | |
| fmt.Printf("Request failed: %v\n", err) | |
| os.Exit(1) | |
| } | |
| defer resp.Body.Close() | |
| if resp.StatusCode != http.StatusOK { | |
| fmt.Printf("Request failed: %s\n", resp.Status) | |
| os.Exit(1) | |
| } | |
| var results []string | |
| if err := json.NewDecoder(resp.Body).Decode(&results); err != nil { | |
| fmt.Printf("Failed to decode response: %v\n", err) | |
| os.Exit(1) | |
| } | |
| if len(results) == 0 { | |
| fmt.Println("No backups found.") | |
| return | |
| } | |
| for _, item := range results { | |
| fmt.Println(item) | |
| } | |
| } | |
| func runMount(args []string) { | |
| cfg := loadConfig() | |
| ensureServerAwake(cfg.ServerURL) | |
| var snapshot string | |
| if len(args) == 1 { | |
| folderName := args[0] | |
| dates, err := fetchList(cfg.ServerURL, "") | |
| if err != nil { | |
| fmt.Printf("Error: Failed to fetch date indexes: %v\n", err) | |
| os.Exit(1) | |
| } | |
| found := false | |
| for i := len(dates) - 1; i >= 0; i-- { | |
| runs, err := fetchList(cfg.ServerURL, dates[i]) | |
| if err != nil { | |
| continue | |
| } | |
| for j := len(runs) - 1; j >= 0; j-- { | |
| if strings.HasSuffix(runs[j], "_"+folderName) { | |
| snapshot = dates[i] + "/" + runs[j] | |
| found = true | |
| break | |
| } | |
| } | |
| if found { | |
| break | |
| } | |
| } | |
| if !found { | |
| fmt.Printf("Error: No backups found matching folder name: %s\n", folderName) | |
| os.Exit(1) | |
| } | |
| } else if len(args) == 2 { | |
| snapshot = args[0] + "/" + args[1] | |
| } else { | |
| fmt.Println("Error: Invalid arguments for mount command.") | |
| os.Exit(1) | |
| } | |
| fmt.Printf("Mounting snapshot: %s\n", snapshot) | |
| manifestURL := fmt.Sprintf("%s/download?snapshot=%s", cfg.ServerURL, snapshot) | |
| resp, err := http.Get(manifestURL) | |
| if err != nil { | |
| fmt.Printf("Failed to download snapshot manifest: %v\n", err) | |
| os.Exit(1) | |
| } | |
| defer resp.Body.Close() | |
| if resp.StatusCode != http.StatusOK { | |
| body, _ := io.ReadAll(resp.Body) | |
| fmt.Printf("Failed to fetch manifest: %s - %s\n", resp.Status, string(body)) | |
| os.Exit(1) | |
| } | |
| var files []string | |
| if err := json.NewDecoder(resp.Body).Decode(&files); err != nil { | |
| fmt.Printf("Failed to decode snapshot list: %v\n", err) | |
| os.Exit(1) | |
| } | |
| if len(files) == 0 { | |
| fmt.Println("Snapshot is empty.") | |
| return | |
| } | |
| fmt.Printf("Restoring %d files...\n", len(files)) | |
| for _, file := range files { | |
| targetLocalPath := filepath.FromSlash(file) | |
| localDir := filepath.Dir(targetLocalPath) | |
| if localDir != "." { | |
| if err := os.MkdirAll(localDir, 0755); err != nil { | |
| fmt.Printf("Failed to create folder %s: %v\n", localDir, err) | |
| os.Exit(1) | |
| } | |
| } | |
| fileURL := fmt.Sprintf("%s/download?file=%s/%s", cfg.ServerURL, snapshot, file) | |
| err = downloadFile(fileURL, targetLocalPath, filepath.Base(targetLocalPath)) | |
| if err != nil { | |
| fmt.Printf("Error downloading file %s: %v\n", file, err) | |
| os.Exit(1) | |
| } | |
| } | |
| fmt.Println("Restore completed successfully!") | |
| } | |
| func fetchList(serverURL, date string) ([]string, error) { | |
| reqURL := fmt.Sprintf("%s/list", serverURL) | |
| if date != "" { | |
| reqURL += "?date=" + date | |
| } | |
| resp, err := http.Get(reqURL) | |
| if err != nil { | |
| return nil, err | |
| } | |
| defer resp.Body.Close() | |
| if resp.StatusCode != http.StatusOK { | |
| return nil, fmt.Errorf("unexpected status %s", resp.Status) | |
| } | |
| var res []string | |
| err = json.NewDecoder(resp.Body).Decode(&res) | |
| return res, err | |
| } | |
| func downloadFile(url, dest string, fileName string) error { | |
| resp, err := http.Get(url) | |
| if err != nil { | |
| return err | |
| } | |
| defer resp.Body.Close() | |
| if resp.StatusCode != http.StatusOK { | |
| body, _ := io.ReadAll(resp.Body) | |
| return fmt.Errorf("server returned %s - %s", resp.Status, string(body)) | |
| } | |
| out, err := os.Create(dest) | |
| if err != nil { | |
| return err | |
| } | |
| defer out.Close() | |
| if resp.ContentLength > 5*1024*1024 { | |
| counter := &writeCounter{total: resp.ContentLength, name: fileName} | |
| _, err = io.Copy(out, io.TeeReader(resp.Body, counter)) | |
| fmt.Println() | |
| } else { | |
| fmt.Printf("Downloading %s...\n", fileName) | |
| _, err = io.Copy(out, resp.Body) | |
| } | |
| return err | |
| } | |