Files
llmgw/internal/sourcecheck/sourcecheck.go
T

85 lines
2.8 KiB
Go

package sourcecheck
import (
"context"
"fmt"
"net/http"
"strings"
"time"
"gitea.mynodes.xyz/ilker/llmgw/internal/safedial"
)
// Check asks whether a model's weights are actually there.
//
// A source that does not resolve is caught today by a replica failing to start,
// minutes later, on a node — and the error names the node rather than the typo.
// Asking now turns that into a sentence while the operator still has the file
// open.
//
// Only URLs. A path is checked on the node that will use it, by the node, which
// is the only machine whose answer means anything; a bare model id is resolved
// by the engine and is not this function's business.
//
// A package of its own rather than a helper in the CLI, because the admin page
// adds models too and was not checking anything. Two ways in that disagree
// about whether the weights exist is the kind of difference nobody finds until
// a replica fails on a node, hours later, naming the node.
func Check(ctx context.Context, source string) (size int64, err error) {
url, ok := HTTPSource(source)
if !ok {
return 0, nil
}
ctx, cancel := context.WithTimeout(ctx, 15*time.Second)
defer cancel()
req, err := http.NewRequestWithContext(ctx, http.MethodHead, url, nil)
if err != nil {
return 0, err
}
// Not anything link-local: instance metadata answering a HEAD would be
// read through this, by whoever can add a model.
resp, err := safedial.Client(safedial.ControlPlane, 0).Do(req)
if err != nil {
return 0, fmt.Errorf("cannot reach %s: %w", url, err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return 0, fmt.Errorf("%s answered %s", url, resp.Status)
}
// Absent when the server does not send it, which is not an error: a
// resumable fetch does not need the length in advance, and refusing here
// would reject a mirror that is merely terse. Go reports that as -1;
// normalised to 0 so callers have one spelling of "unknown".
if resp.ContentLength < 0 {
return 0, nil
}
return resp.ContentLength, nil
}
// HTTPSource returns the URL to check, and whether there is one.
//
// hf: is expanded the way the fetcher expands it, so the check asks about the
// same file the download will ask for rather than about a string that looks
// similar.
func HTTPSource(source string) (string, bool) {
switch {
case strings.HasPrefix(source, "http://"), strings.HasPrefix(source, "https://"):
return source, true
case strings.HasPrefix(source, "hf:"):
rest := strings.TrimPrefix(source, "hf:")
parts := strings.SplitN(rest, "/", 3)
if len(parts) != 3 {
// A repository rather than a file. vLLM and SGLang download those
// themselves, so there is nothing here to check.
return "", false
}
return fmt.Sprintf("https://huggingface.co/%s/%s/resolve/main/%s",
parts[0], parts[1], parts[2]), true
default:
return "", false
}
}