steve/majordomo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: music poll resilience + hostile-URL guard + embed dup-index (gadfly round 1)

Browse Source 
- pollResult tolerates up to 5 CONSECUTIVE bad polls (transport blip,
  unparseable payload, task momentarily absent) instead of killing a
  multi-minute exclusive-GPU job on the first hiccup; only status=2, a
  failure run, or ctx deadline aborts
- server-supplied result.File must be server-relative; combined with the
  upstreamPath dot-dot/scheme rejection this stops a hostile upstream
  from steering the follow-up GET at other proxy endpoints (test:
  ../../api/models/unload refused)
- WithSteps(<=0) rejected; embed responses repeating an index rejected;
  musicFormatMIME now wraps speechMIME (one format table, wav32
  normalized); poll interval is a test-shrinkable var (CI no longer
  burns 2s+ per music test); parens + comments per review

Co-Authored-By: Claude Fable 5 <[email protected]>
This commit is contained in:
steve
2026-07-13 00:52:06 -04:00
co-authored by Claude Fable 5
parent cacce61ddc
commit cf2d83f157
3 changed files with 143 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
provider/llamaswap/music.go
+68 -35
View File
@@ -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),