diff --git a/provider/llamaswap/embed.go b/provider/llamaswap/embed.go index 2705955..74bfca3 100644 --- a/provider/llamaswap/embed.go +++ b/provider/llamaswap/embed.go @@ -75,6 +75,10 @@ func (m *embedModel) Embed(ctx context.Context, req embeddings.EmbedRequest, opt return nil, &llm.APIError{Provider: m.p.name, Model: m.id, Message: fmt.Sprintf("embeddings response entry index %d invalid or empty", d.Index)} } + if vectors[d.Index] != nil { + return nil, &llm.APIError{Provider: m.p.name, Model: m.id, + Message: fmt.Sprintf("embeddings response repeats index %d", d.Index)} + } vectors[d.Index] = d.Embedding } for i, v := range vectors { diff --git a/provider/llamaswap/music.go b/provider/llamaswap/music.go index a3631cc..cd383b9 100644 --- a/provider/llamaswap/music.go +++ b/provider/llamaswap/music.go @@ -25,8 +25,14 @@ import ( // musicPollInterval is the delay between query_result polls. Long enough to // be polite to the queue, short enough that a ~10s xl-turbo song isn't -// dominated by poll latency. -const musicPollInterval = 2 * time.Second +// dominated by poll latency. A var so tests can shrink it. +var musicPollInterval = 2 * time.Second + +// musicPollMaxConsecutiveFailures bounds how many consecutive BAD polls +// (transport error, unparseable payload, task momentarily absent) are +// tolerated before aborting. A multi-minute GPU job must not die to one +// blip; a genuinely broken upstream still fails within ~5 intervals. +const musicPollMaxConsecutiveFailures = 5 // MusicModel implements musicgen.Provider. The id selects which upstream // llama-swap loads. @@ -44,7 +50,8 @@ type musicModel struct { } // releaseTaskRequest is the ACE-Step POST /release_task shape. audio_duration -// is the v1 param name; verify against docs/en/API.md at smoke time — an +// is the v1 param name; verify against the upstream ACE-Step-1.5 repo's +// docs/en/API.md at smoke time — an // unknown field is ignored upstream, degrading to the default clip length, // never an error. type releaseTaskRequest struct { @@ -65,20 +72,16 @@ type queryItem struct { Result string `json:"result"` } -// musicFormatMIME maps ACE-Step's audio_format values to MIME types. -func musicFormatMIME(format string) string { - switch strings.ToLower(strings.TrimSpace(format)) { - case "wav", "wav32": - return "audio/wav" - case "flac": - return "audio/flac" - case "opus": - return "audio/ogg" - case "aac": - return "audio/aac" - default: // "", "mp3" - return "audio/mpeg" +// musicFormatMIME resolves the clip MIME from the response Content-Type +// or the requested format, reusing speechMIME's table (one format->MIME +// map for the whole provider). wav32 is ACE-Step-specific: normalize it +// to wav before the shared lookup. +func musicFormatMIME(contentType, format string) string { + format = strings.ToLower(strings.TrimSpace(format)) + if format == "wav32" { + format = "wav" } + return speechMIME(contentType, format) } // Generate implements musicgen.Model. @@ -90,6 +93,9 @@ func (m *musicModel) Generate(ctx context.Context, req musicgen.Request, opts .. if req.DurationSeconds < 0 { return nil, fmt.Errorf("%w: duration must be >= 0, got %d", llm.ErrUnsupported, req.DurationSeconds) } + if req.Steps != nil && *req.Steps <= 0 { + return nil, fmt.Errorf("%w: inference steps must be > 0, got %d", llm.ErrUnsupported, *req.Steps) + } taskID, err := m.releaseTask(ctx, req) if err != nil { @@ -139,7 +145,12 @@ func (m *musicModel) releaseTask(ctx context.Context, req musicgen.Request) (str } // pollResult polls query_result until the task succeeds, fails, or ctx -// expires. +// expires. Transient trouble — a transport blip, a momentarily +// unparseable payload, the task briefly absent from the response — is +// tolerated up to musicPollMaxConsecutiveFailures in a row: a +// multi-minute exclusive-GPU job must not die to one flaky poll. Only an +// explicit status=2, a run of consecutive failures, or the ctx deadline +// aborts. func (m *musicModel) pollResult(ctx context.Context, taskID string) (*queryItem, error) { path, err := upstreamPath(m.id, "/query_result") if err != nil { @@ -148,22 +159,29 @@ func (m *musicModel) pollResult(ctx context.Context, taskID string) (*queryItem, body := map[string]any{"task_id_list": []string{taskID}} ticker := time.NewTicker(musicPollInterval) defer ticker.Stop() + failures := 0 + var lastErr error for { - // Tolerant envelope: items under "data" or a bare array. - var raw json.RawMessage - if err := m.p.doJSON(ctx, http.MethodPost, path, m.id, body, &raw); err != nil { - return nil, err - } - item, err := findQueryItem(raw, taskID) - if err != nil { - return nil, &llm.APIError{Provider: m.p.name, Model: m.id, Message: err.Error()} - } - switch item.Status { - case 1: + item, pollErr := m.pollOnce(ctx, path, body, taskID) + switch { + case pollErr != nil: + // Ctx expiry is never transient — bail with the deadline error. + if ctx.Err() != nil { + return nil, fmt.Errorf("llama-swap: music generation: %w", ctx.Err()) + } + failures++ + lastErr = pollErr + if failures >= musicPollMaxConsecutiveFailures { + return nil, &llm.APIError{Provider: m.p.name, Model: m.id, + Message: fmt.Sprintf("music poll failed %d times in a row: %v", failures, lastErr)} + } + case item.Status == 1: return item, nil - case 2: + case item.Status == 2: return nil, &llm.APIError{Provider: m.p.name, Model: m.id, Message: "music generation failed upstream: " + truncateForError([]byte(item.Result))} + default: + failures = 0 // healthy queued/running poll } select { case <-ctx.Done(): @@ -173,6 +191,16 @@ func (m *musicModel) pollResult(ctx context.Context, taskID string) (*queryItem, } } +// pollOnce performs one query_result round trip and locates the task. +func (m *musicModel) pollOnce(ctx context.Context, path string, body any, taskID string) (*queryItem, error) { + // Tolerant envelope: items under "data" or a bare array. + var raw json.RawMessage + if err := m.p.doJSON(ctx, http.MethodPost, path, m.id, body, &raw); err != nil { + return nil, err + } + return findQueryItem(raw, taskID) +} + // findQueryItem digs the task's entry out of the query_result payload, // tolerating {"data": [...]}, {"data": {...}}, and bare-array envelopes. func findQueryItem(raw json.RawMessage, taskID string) (*queryItem, error) { @@ -192,7 +220,10 @@ func findQueryItem(raw json.RawMessage, taskID string) (*queryItem, error) { items = []queryItem{one} } for i := range items { - if items[i].TaskID == taskID || items[i].TaskID == "" && len(items) == 1 { + // Single-item responses without a task_id echo are assumed to be + // ours — we only ever poll one task; a mismatch surfaces as a + // transient miss and is retried by the caller. + if items[i].TaskID == taskID || (items[i].TaskID == "" && len(items) == 1) { return &items[i], nil } } @@ -209,7 +240,12 @@ func (m *musicModel) fetchResult(ctx context.Context, format string, item *query Message: "music result blob missing file URL: " + truncateForError([]byte(item.Result))} } // The file URL is server-relative (e.g. "/v1/audio?path=..."); route it - // back through the same upstream. + // back through the same upstream. upstreamPath additionally refuses + // dot-dot/scheme smuggling in this SERVER-SUPPLIED value. + if !strings.HasPrefix(result.File, "/") { + return nil, &llm.APIError{Provider: m.p.name, Model: m.id, + Message: "music result file URL is not server-relative: " + truncateForError([]byte(result.File))} + } path, err := upstreamPath(m.id, result.File) if err != nil { return nil, err @@ -221,10 +257,7 @@ func (m *musicModel) fetchResult(ctx context.Context, format string, item *query if len(raw) == 0 { return nil, &llm.APIError{Provider: m.p.name, Model: m.id, Message: "music response contained no audio"} } - mimeType := mimeFromContentType(contentType, "audio/") - if mimeType == "" { - mimeType = musicFormatMIME(format) - } + mimeType := musicFormatMIME(contentType, format) return &musicgen.Result{ Audio: musicgen.Audio{Data: raw, MIME: mimeType}, Raw: json.RawMessage(item.Result), diff --git a/provider/llamaswap/music_embed_test.go b/provider/llamaswap/music_embed_test.go index ea3f7a5..5d15019 100644 --- a/provider/llamaswap/music_embed_test.go +++ b/provider/llamaswap/music_embed_test.go @@ -51,6 +51,10 @@ func aceStepStub(t *testing.T, mp3 []byte) (*httptest.Server, *atomic.Int32) { } func TestMusicGenerate(t *testing.T) { + old := musicPollInterval + musicPollInterval = 5 * time.Millisecond + defer func() { musicPollInterval = old }() + mp3 := []byte("ID3fakeaudio") srv, polls := aceStepStub(t, mp3) defer srv.Close() @@ -79,6 +83,10 @@ func TestMusicGenerate(t *testing.T) { } func TestMusicGenerateUpstreamFailure(t *testing.T) { + old := musicPollInterval + musicPollInterval = 5 * time.Millisecond + defer func() { musicPollInterval = old }() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/upstream/musicgen-acestep/release_task": @@ -222,3 +230,66 @@ func TestInstructedQuery(t *testing.T) { t.Errorf("InstructedQuery = %q", got) } } + +func TestMusicGenerateSurvivesTransientPollFailures(t *testing.T) { + old := musicPollInterval + musicPollInterval = 5 * time.Millisecond + defer func() { musicPollInterval = old }() + + mp3 := []byte("ID3fakeaudio") + var polls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/upstream/musicgen-acestep/release_task": + _, _ = w.Write([]byte(`{"data": {"task_id": "t-3"}}`)) + case "/upstream/musicgen-acestep/query_result": + switch polls.Add(1) { + case 1: + w.WriteHeader(http.StatusBadGateway) // transient transport blip + case 2: + _, _ = w.Write([]byte(`{"data": []}`)) // task momentarily absent + default: + _, _ = w.Write([]byte(`{"data": [{"task_id": "t-3", "status": 1, + "result": "{\"file\": \"/v1/audio?path=out.mp3\"}"}]}`)) + } + case "/upstream/musicgen-acestep/v1/audio": + w.Header().Set("Content-Type", "audio/mpeg") + _, _ = w.Write(mp3) + } + })) + defer srv.Close() + + p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client())) + mm, _ := p.MusicModel("musicgen-acestep") + res, err := mm.Generate(context.Background(), musicgen.Request{Prompt: "p"}) + if err != nil { + t.Fatalf("Generate should survive 2 transient failures: %v", err) + } + if len(res.Audio.Data) == 0 { + t.Fatal("no audio") + } +} + +func TestMusicGenerateRejectsHostileFileURL(t *testing.T) { + old := musicPollInterval + musicPollInterval = 5 * time.Millisecond + defer func() { musicPollInterval = old }() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/upstream/musicgen-acestep/release_task": + _, _ = w.Write([]byte(`{"data": {"task_id": "t-4"}}`)) + case "/upstream/musicgen-acestep/query_result": + _, _ = w.Write([]byte(`{"data": [{"task_id": "t-4", "status": 1, + "result": "{\"file\": \"/v1/audio?path=../../api/models/unload\"}"}]}`)) + } + })) + defer srv.Close() + + p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client())) + mm, _ := p.MusicModel("musicgen-acestep") + _, err := mm.Generate(context.Background(), musicgen.Request{Prompt: "p"}) + if err == nil { + t.Fatal("dot-dot result file URL accepted") + } +}