| | |
| | | package model |
| | | |
| | | import ( |
| | | "context" |
| | | "errors" |
| | | "fmt" |
| | | "io" |
| | | "net/http" |
| | | "os" |
| | | "os/exec" |
| | | "path/filepath" |
| | | "strconv" |
| | | "strings" |
| | | "time" |
| | | |
| | | "voicesnap/internal/logger" |
| | |
| | | } |
| | | |
| | | func DownloadProfile(profile ModelProfile, urls []string, modelsDir string, progress ProgressCallback) error { |
| | | return DownloadProfileWithContext(context.Background(), profile, urls, modelsDir, progress) |
| | | } |
| | | |
| | | func DownloadProfileWithContext(ctx context.Context, profile ModelProfile, urls []string, modelsDir string, progress ProgressCallback) error { |
| | | if err := os.MkdirAll(modelsDir, 0755); err != nil { |
| | | return fmt.Errorf("failed to create models dir: %w", err) |
| | | } |
| | | cleanupStaleDownloadDirs(profile, modelsDir) |
| | | |
| | | var lastErr error |
| | | for i, url := range urls { |
| | | if err := ctx.Err(); err != nil { |
| | | return err |
| | | } |
| | | if url == "" { |
| | | continue |
| | | } |
| | | if err := downloadAndInstallFromURL(profile, url, modelsDir, i+1, progress); err != nil { |
| | | if err := downloadAndInstallFromURL(ctx, profile, url, modelsDir, i+1, progress); err != nil { |
| | | if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { |
| | | return err |
| | | } |
| | | lastErr = err |
| | | logger.Info("Model install from URL %d failed: %v", i+1, err) |
| | | continue |
| | |
| | | return fmt.Errorf("download failed: %w", lastErr) |
| | | } |
| | | |
| | | func downloadAndInstallFromURL(profile ModelProfile, url, modelsDir string, attempt int, progress ProgressCallback) error { |
| | | func cleanupStaleDownloadDirs(profile ModelProfile, modelsDir string) { |
| | | downloadRoot := filepath.Join(modelsDir, ".downloads", profile.ID) |
| | | if err := os.RemoveAll(downloadRoot); err != nil { |
| | | logger.Info("Failed to remove stale download dir %s: %v", downloadRoot, err) |
| | | } |
| | | } |
| | | |
| | | func downloadAndInstallFromURL(ctx context.Context, profile ModelProfile, url, modelsDir string, attempt int, progress ProgressCallback) error { |
| | | runID := fmt.Sprintf("%d-%d", time.Now().UnixNano(), attempt) |
| | | downloadDir := filepath.Join(modelsDir, ".downloads", profile.ID, runID) |
| | | extractDir := filepath.Join(downloadDir, "extract") |
| | |
| | | defer os.RemoveAll(stagingDir) |
| | | |
| | | logger.Info("Downloading model from URL %d: %s", attempt, url) |
| | | if err := downloadFile(url, archivePath, progress); err != nil { |
| | | if err := downloadFileWithContext(ctx, url, archivePath, progress); err != nil { |
| | | return err |
| | | } |
| | | if err := ctx.Err(); err != nil { |
| | | return err |
| | | } |
| | | |
| | | logger.Info("Extracting model archive...") |
| | | if err := extractArchive(archivePath, extractDir); err != nil { |
| | | if err := extractArchiveWithContext(ctx, archivePath, extractDir); err != nil { |
| | | return fmt.Errorf("extraction failed: %w", err) |
| | | } |
| | | if err := ctx.Err(); err != nil { |
| | | return err |
| | | } |
| | | |
| | | sourceDir, err := findInstallSource(profile, extractDir) |
| | | if err != nil { |
| | | return err |
| | | } |
| | | if err := ctx.Err(); err != nil { |
| | | return err |
| | | } |
| | | |
| | | if err := moveOrCopyDir(sourceDir, stagingDir); err != nil { |
| | | return fmt.Errorf("failed to stage model: %w", err) |
| | | } |
| | | if err := ctx.Err(); err != nil { |
| | | return err |
| | | } |
| | | |
| | | validation := ValidateModelDir(profile, stagingDir) |
| | | if !validation.Valid { |
| | | return fmt.Errorf("downloaded model is incomplete: missing=%v problems=%v", validation.Missing, validation.Problems) |
| | | } |
| | | if err := ctx.Err(); err != nil { |
| | | return err |
| | | } |
| | | |
| | | finalDir := filepath.Join(modelsDir, profile.InstallDirName) |
| | |
| | | } |
| | | |
| | | func downloadFile(url, destPath string, progress ProgressCallback) error { |
| | | resp, err := http.Get(url) |
| | | return downloadFileWithContext(context.Background(), url, destPath, progress) |
| | | } |
| | | |
| | | func downloadFileWithContext(ctx context.Context, url, destPath string, progress ProgressCallback) error { |
| | | const maxAttempts = 6 |
| | | var lastErr error |
| | | var lastPercent = -1.0 |
| | | |
| | | for attempt := 1; attempt <= maxAttempts; attempt++ { |
| | | if err := ctx.Err(); err != nil { |
| | | return err |
| | | } |
| | | downloaded := existingFileSize(destPath) |
| | | err := downloadFileAttempt(ctx, url, destPath, downloaded, &lastPercent, progress) |
| | | if err == nil { |
| | | return nil |
| | | } |
| | | if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { |
| | | return err |
| | | } |
| | | if isNonRetryableDownloadError(err) { |
| | | return err |
| | | } |
| | | lastErr = err |
| | | logger.Info("Download attempt %d/%d failed: %v", attempt, maxAttempts, err) |
| | | if attempt < maxAttempts { |
| | | select { |
| | | case <-ctx.Done(): |
| | | return ctx.Err() |
| | | case <-time.After(time.Duration(attempt) * time.Second): |
| | | } |
| | | } |
| | | } |
| | | |
| | | return fmt.Errorf("download failed after %d attempts: %w", maxAttempts, lastErr) |
| | | } |
| | | |
| | | func downloadFileAttempt(ctx context.Context, url, destPath string, resumeFrom int64, lastPercent *float64, progress ProgressCallback) error { |
| | | req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) |
| | | if err != nil { |
| | | return nonRetryableDownloadError{err: err} |
| | | } |
| | | if resumeFrom > 0 { |
| | | req.Header.Set("Range", fmt.Sprintf("bytes=%d-", resumeFrom)) |
| | | } |
| | | |
| | | resp, err := http.DefaultClient.Do(req) |
| | | if err != nil { |
| | | return err |
| | | } |
| | | defer resp.Body.Close() |
| | | |
| | | if resp.StatusCode != http.StatusOK { |
| | | return fmt.Errorf("HTTP %d: %s", resp.StatusCode, resp.Status) |
| | | appendMode := resumeFrom > 0 && resp.StatusCode == http.StatusPartialContent |
| | | if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusPartialContent { |
| | | err := fmt.Errorf("HTTP %d: %s", resp.StatusCode, resp.Status) |
| | | if resp.StatusCode < http.StatusInternalServerError { |
| | | return nonRetryableDownloadError{err: err} |
| | | } |
| | | return err |
| | | } |
| | | if resumeFrom > 0 && resp.StatusCode == http.StatusOK { |
| | | logger.Info("Download server ignored Range request; restarting download") |
| | | resumeFrom = 0 |
| | | appendMode = false |
| | | } |
| | | |
| | | totalBytes := resp.ContentLength |
| | | out, err := os.Create(destPath) |
| | | if appendMode { |
| | | if total := parseContentRangeTotal(resp.Header.Get("Content-Range")); total > 0 { |
| | | totalBytes = total |
| | | } else if resp.ContentLength > 0 { |
| | | totalBytes = resumeFrom + resp.ContentLength |
| | | } |
| | | } |
| | | |
| | | flag := os.O_CREATE | os.O_WRONLY |
| | | if appendMode { |
| | | flag |= os.O_APPEND |
| | | } else { |
| | | flag |= os.O_TRUNC |
| | | } |
| | | out, err := os.OpenFile(destPath, flag, 0644) |
| | | if err != nil { |
| | | return err |
| | | } |
| | | defer out.Close() |
| | | |
| | | downloaded := resumeFrom |
| | | if !appendMode { |
| | | downloaded = 0 |
| | | } |
| | | buf := make([]byte, 32*1024) |
| | | var downloaded int64 |
| | | lastPercent := -1.0 |
| | | |
| | | for { |
| | | if err := ctx.Err(); err != nil { |
| | | return err |
| | | } |
| | | n, readErr := resp.Body.Read(buf) |
| | | if n > 0 { |
| | | if _, writeErr := out.Write(buf[:n]); writeErr != nil { |
| | |
| | | |
| | | if totalBytes > 0 && progress != nil { |
| | | pct := float64(downloaded) / float64(totalBytes) * 100 |
| | | if pct-lastPercent >= 0.5 { |
| | | lastPercent = pct |
| | | if pct-*lastPercent >= 0.5 { |
| | | *lastPercent = pct |
| | | progress(pct, downloaded, totalBytes) |
| | | } |
| | | } |
| | | } |
| | | if readErr != nil { |
| | | if readErr == io.EOF { |
| | | break |
| | | if totalBytes > 0 && downloaded < totalBytes { |
| | | return fmt.Errorf("short download: got %d of %d bytes", downloaded, totalBytes) |
| | | } |
| | | return nil |
| | | } |
| | | return readErr |
| | | } |
| | | } |
| | | } |
| | | |
| | | return nil |
| | | func existingFileSize(path string) int64 { |
| | | info, err := os.Stat(path) |
| | | if err != nil || info.IsDir() { |
| | | return 0 |
| | | } |
| | | return info.Size() |
| | | } |
| | | |
| | | func parseContentRangeTotal(value string) int64 { |
| | | if value == "" { |
| | | return 0 |
| | | } |
| | | slash := strings.LastIndex(value, "/") |
| | | if slash < 0 || slash == len(value)-1 { |
| | | return 0 |
| | | } |
| | | total, err := strconv.ParseInt(value[slash+1:], 10, 64) |
| | | if err != nil { |
| | | return 0 |
| | | } |
| | | return total |
| | | } |
| | | |
| | | type nonRetryableDownloadError struct { |
| | | err error |
| | | } |
| | | |
| | | func (e nonRetryableDownloadError) Error() string { |
| | | return e.err.Error() |
| | | } |
| | | |
| | | func (e nonRetryableDownloadError) Unwrap() error { |
| | | return e.err |
| | | } |
| | | |
| | | func isNonRetryableDownloadError(err error) bool { |
| | | var target nonRetryableDownloadError |
| | | return errors.As(err, &target) |
| | | } |
| | | |
| | | func extractArchive(archivePath, destDir string) error { |
| | | cmd := exec.Command("tar", "-xf", archivePath, "-C", destDir) |
| | | return extractArchiveWithContext(context.Background(), archivePath, destDir) |
| | | } |
| | | |
| | | func extractArchiveWithContext(ctx context.Context, archivePath, destDir string) error { |
| | | cmd := exec.CommandContext(ctx, "tar", "-xf", archivePath, "-C", destDir) |
| | | output, err := cmd.CombinedOutput() |
| | | if err != nil { |
| | | if ctxErr := ctx.Err(); ctxErr != nil { |
| | | return ctxErr |
| | | } |
| | | return fmt.Errorf("tar extraction failed: %v, output: %s", err, string(output)) |
| | | } |
| | | return nil |