From aa531b1053a6ce7e00fd42fbf32705071a909c0c Mon Sep 17 00:00:00 2001 From: Ilker Manap Date: Wed, 14 Oct 2026 13:55:34 +0200 Subject: [PATCH] ComfyUI models of one environment share one replica flux and sf3d on one card handed it to each other, each starting ComfyUI and loading its weights from nothing (BACKLOG H3, placement side). An environment of two or more models now gets a hidden catalog entry, comfyui-env-: on its first member's profiles with the most memory any member needs, on demand for the longest idle timeout, carrying every member's weights, repositories and the environment's custom nodes. The members get no replicas of their own; asking for one wants the entry. The agent fetches each member's weights; the gate links all of them, refusing two weights files of one name, links each member's Hugging Face repository, and answers /v1/models for every member. The control plane offers the replica under the members' names with each one's workflow and checkpoint, which the router sends by the model asked for. A member's sibling never evicts the replica serving it, and the members' load counts as the entry's. Co-Authored-By: Claude Opus 5.5 (1M context) --- CHANGELOG.md | 8 ++ docs/admin-api.openapi.json | 49 ++++++++ internal/agent/agent.go | 26 ++++ internal/agent/engines/comfyui.go | 7 +- internal/comfy/gate.go | 137 ++++++++++++++++++++- internal/comfy/gate_test.go | 38 ++++++ internal/comfy/progress_test.go | 2 +- internal/control/catalog/catalog.go | 49 +++++++- internal/control/catalog/environments.go | 2 +- internal/control/catalog/shared.go | 110 +++++++++++++++++ internal/control/catalog/shared_test.go | 51 ++++++++ internal/control/control.go | 29 +++++ internal/control/reconciler/demand.go | 25 +++- internal/control/reconciler/evict.go | 2 + internal/control/reconciler/health.go | 4 +- internal/control/reconciler/shared_test.go | 57 +++++++++ internal/router/canary.go | 2 +- internal/router/endpoint/endpoint.go | 19 +++ internal/router/images.go | 3 +- internal/router/meshes.go | 3 +- internal/router/router.go | 6 +- internal/router/workflows.go | 3 +- 22 files changed, 614 insertions(+), 18 deletions(-) create mode 100644 internal/control/catalog/shared.go create mode 100644 internal/control/catalog/shared_test.go create mode 100644 internal/control/reconciler/shared_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index c9821cb..ab2bab1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -28,6 +28,14 @@ APIs change between minors without a deprecation period. ### Added +- **One ComfyUI for an environment's models**: ComfyUI models in one + `comfyui.environment` (two or more) share one replica — the catalog + makes it a hidden entry, `comfyui-env-`, on the first member's + profiles with the most memory any member needs, carrying every member's + weights and custom nodes. Asking for any member starts it; it answers for + each with that member's workflow and checkpoint, and the members get no + replicas of their own. A picture model and a 3D model on one card stop + handing it to each other with a cold start each way. - **Extended thinking on `/v1/messages`**: `thinking: {type: "enabled"}` tells a model that thinks by its chat template to (`enable_thinking`), and its reasoning comes back as `thinking` blocks before the text, diff --git a/docs/admin-api.openapi.json b/docs/admin-api.openapi.json index 948bdfc..c216c6f 100644 --- a/docs/admin-api.openapi.json +++ b/docs/admin-api.openapi.json @@ -9292,6 +9292,10 @@ "origin": { "$ref": "#/components/schemas/catalog.Origin", "description": "Origin is the registry repository the weights came from, at which commit, and each file's checksum there. Written when a model is served from the registry; for one that was not, worked out from the mirror's record of where it fetched the file. It is what lets the control plane say that the repository has published different weights since, and move the model to them." + }, + "shared": { + "$ref": "#/components/schemas/catalog.SharedComfy", + "description": "Shared is set on an environment's hidden entry only: the members its replica answers for. Made by the catalog, never written." } } }, @@ -9545,6 +9549,51 @@ } } }, + "catalog.SharedComfy": { + "type": "object", + "description": "SharedComfy is what an environment's entry carries of its members.", + "properties": { + "environment": { + "type": "string" + }, + "members": { + "type": "array", + "items": { + "$ref": "#/components/schemas/catalog.SharedMember" + } + } + } + }, + "catalog.SharedMember": { + "type": "object", + "description": "SharedMember is one model a shared replica answers for.", + "properties": { + "name": { + "type": "string" + }, + "aliases": { + "type": "array", + "items": { + "type": "string" + } + }, + "source": { + "type": "string", + "description": "Source is the member's weights, which the replica's node fetches." + }, + "hf_repo": { + "type": "string", + "description": "HFRepo is the repository its node loads the weights by name from." + }, + "workflow": { + "type": "string", + "description": "Workflow and Checkpoint are what a request for it is sent." + }, + "checkpoint": { + "type": "string" + } + } + }, "eval.Live": { "type": "object", "description": "Live is a run as it stands: one running, with how many of its cases are done, or one finished.", diff --git a/internal/agent/agent.go b/internal/agent/agent.go index c21d48e..5ec10f8 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -6,6 +6,7 @@ import ( "crypto/sha256" "crypto/tls" "encoding/hex" + "encoding/json" "errors" "fmt" "io" @@ -1131,6 +1132,31 @@ func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan< stop() spec.Profile.ModelPath = path } + // A shared ComfyUI replica's members' weights, each fetched like the + // model's, and handed to the gate where they landed. + if raw := spec.Profile.Env[engines.SharedModelsEnv]; raw != "" { + var ms []map[string]string + if err := json.Unmarshal([]byte(raw), &ms); err != nil { + report(supervisor.Failed, "the shared replica's models are not readable: "+err.Error(), nil) + return + } + for i, m := range ms { + stop := a.watchFetch(m["source"], func(p weights.Progress) { + send(&llmgwv1.ShardStatus{State: agentsvc.StateToWire(supervisor.Starting), + Fetch: &llmgwv1.FetchProgress{Source: p.Source, BytesDone: p.BytesDone, BytesTotal: p.BytesTotal}}) + }) + path, err := a.weights().Ensure(ctx, m["source"]) + stop() + if err != nil { + report(supervisor.Failed, fmt.Sprintf("model %s: %v", m["name"], err), nil) + return + } + ms[i]["path"] = path + } + b, _ := json.Marshal(ms) + spec.Profile.Env = maps.Clone(spec.Profile.Env) + spec.Profile.Env[engines.SharedModelsEnv] = string(b) + } // Adapters are fetched like the model, and before the engine starts: it // loads them at start, and one missing is a replica that answers for the // base model under the adapter's name. diff --git a/internal/agent/engines/comfyui.go b/internal/agent/engines/comfyui.go index d92e9b7..333b37f 100644 --- a/internal/agent/engines/comfyui.go +++ b/internal/agent/engines/comfyui.go @@ -129,7 +129,7 @@ func (c *ComfyUI) Commands(s ShardSpec) ([]Command, error) { // so taken before baseEnv drops them. Without the repository the gate // links no Hugging Face cache, and sf3d's node went to the network for // its gated repository and was refused with a 401. - for _, k := range []string{comfyNodesEnv, comfyHFRepoEnv} { + for _, k := range []string{comfyNodesEnv, comfyHFRepoEnv, SharedModelsEnv} { if v := p.Env[k]; v != "" { env[k] = v } @@ -144,3 +144,8 @@ const comfyNodesEnv = "LLMGW_COMFYUI_NODES" // comfyHFRepoEnv is the repository a custom node loads the weights by name // from; catalog.HFRepoEnv and comfy.HFRepoEnv by the same name. const comfyHFRepoEnv = "LLMGW_COMFYUI_HF_REPO" + +// SharedModelsEnv is a shared replica's members — name, source, and the +// path the agent fetched it to — as JSON; catalog.SharedModelsEnv and +// comfy.SharedModelsEnv by the same name. +const SharedModelsEnv = "LLMGW_COMFYUI_MODELS" diff --git a/internal/comfy/gate.go b/internal/comfy/gate.go index 63c8159..5bcb547 100644 --- a/internal/comfy/gate.go +++ b/internal/comfy/gate.go @@ -20,6 +20,7 @@ import ( "os/signal" "path/filepath" "regexp" + "slices" "strconv" "strings" "syscall" @@ -112,7 +113,15 @@ func RunGate(ctx context.Context, cfg GateConfig) error { return err } models := filepath.Join(state, "models") - if err := LinkWeights(cfg.Weights, models); err != nil { + shared, err := parseShared(os.Getenv(SharedModelsEnv)) + if err != nil { + return err + } + all := []string{cfg.Weights} + for _, m := range shared { + all = append(all, m.Path) + } + if err := LinkAllWeights(all, models); err != nil { return err } if repo := os.Getenv(HFRepoEnv); repo != "" { @@ -120,6 +129,20 @@ func RunGate(ctx context.Context, cfg GateConfig) error { return fmt.Errorf("showing the weights as %s: %w", repo, err) } } + for _, m := range shared { + if m.HFRepo == "" { + continue + } + if err := LinkHFCache(m.Path, m.HFRepo, hfHubDir()); err != nil { + return fmt.Errorf("showing %s's weights as %s: %w", m.Name, m.HFRepo, err) + } + } + names := []string{cfg.Model} + for _, m := range shared { + if !slices.Contains(names, m.Name) { + names = append(names, m.Name) + } + } extra := filepath.Join(state, "extra_model_paths.yaml") if err := os.WriteFile(extra, []byte(extraModelPaths(models)), 0o644); err != nil { //nolint:gosec // not secret return err @@ -177,7 +200,7 @@ func RunGate(ctx context.Context, cfg GateConfig) error { go book.watch(ctx, upstream, clientID, log) srv := &http.Server{ Addr: cfg.Listen, - Handler: newGate(upstream, cfg.Key, cfg.Model, book, clientID), + Handler: newGate(upstream, cfg.Key, names, book, clientID), ReadHeaderTimeout: 10 * time.Second, } served := make(chan error, 1) @@ -218,12 +241,12 @@ func shutdown(srv *http.Server) { // NewGate is the gate's handler, in front of the ComfyUI at upstream. func NewGate(upstream *url.URL, key, model string) http.Handler { - return newGate(upstream, key, model, newProgressBook(), "") + return newGate(upstream, key, []string{model}, newProgressBook(), "") } // newGate is the gate with the book of its workflows' progress, kept under // clientID; an empty one leaves the workflows' own client ids alone. -func newGate(upstream *url.URL, key, model string, book *progressBook, clientID string) http.Handler { +func newGate(upstream *url.URL, key string, names []string, book *progressBook, clientID string) http.Handler { proxy := httputil.NewSingleHostReverseProxy(upstream) base := proxy.Director proxy.Director = func(r *http.Request) { @@ -269,7 +292,11 @@ func newGate(upstream *url.URL, key, model string, book *progressBook, clientID return } w.Header().Set("Content-Type", "application/json") - fmt.Fprintf(w, `{"object":"list","data":[{"id":%q,"object":"model","owned_by":"comfyui"}]}`, model) + data := make([]map[string]string, 0, len(names)) + for _, n := range names { + data = append(data, map[string]string{"id": n, "object": "model", "owned_by": "comfyui"}) + } + _ = json.NewEncoder(w).Encode(map[string]any{"object": "list", "data": data}) return } if id, ok := strings.CutPrefix(r.URL.Path, "/progress/"); ok && r.Method == http.MethodGet { @@ -306,6 +333,106 @@ func FileName(path string) string { return cachedName.ReplaceAllString(filepath.Base(path), "") } +// SharedModelsEnv is where the agent hands a shared replica's gate its +// members: name, the path their weights were fetched to, and the repository +// a node loads them by name from. +const SharedModelsEnv = "LLMGW_COMFYUI_MODELS" + +// sharedModel is one member of a shared replica. +type sharedModel struct { + Name string `json:"name"` + Path string `json:"path"` + HFRepo string `json:"hf_repo"` +} + +func parseShared(raw string) ([]sharedModel, error) { + if strings.TrimSpace(raw) == "" { + return nil, nil + } + var ms []sharedModel + if err := json.Unmarshal([]byte(raw), &ms); err != nil { + return nil, fmt.Errorf("%s: %w", SharedModelsEnv, err) + } + for _, m := range ms { + if m.Name == "" || m.Path == "" { + return nil, fmt.Errorf("%s: a member with no name or no fetched path", SharedModelsEnv) + } + } + return ms, nil +} + +// LinkAllWeights shows ComfyUI every model's weights at once, for a replica +// that answers for several: links made before are taken away first, then +// each file of each is linked by its name. Two weights files of one name +// are refused — ComfyUI would load one for both; anything else of one name +// (a README, a licence) the first keeps. +func LinkAllWeights(weights []string, dir string) error { + if len(weights) == 1 { + return LinkWeights(weights[0], dir) + } + if err := LinkWeights(weights[0], dir); err != nil { + return err + } + for _, w := range weights[1:] { + files, err := weightFiles(w) + if err != nil { + return err + } + for _, f := range files { + link := filepath.Join(dir, FileName(f)) + if cur, err := os.Readlink(link); err == nil { + if cur == f { + continue + } + if isWeightsFile(f) { + return fmt.Errorf("two models' weights are both named %s (%s and %s): give one comfyui.checkpoint of another name", + FileName(f), cur, f) + } + continue + } + if err := os.Symlink(f, link); err != nil { + return err + } + } + } + return nil +} + +// weightFiles are the files a weights source is: itself, or each of its +// directory's. +func weightFiles(weights string) ([]string, error) { + weights, err := filepath.Abs(weights) + if err != nil { + return nil, err + } + info, err := os.Stat(weights) + if err != nil { + return nil, fmt.Errorf("the weights: %w", err) + } + if !info.IsDir() { + return []string{weights}, nil + } + entries, err := os.ReadDir(weights) + if err != nil { + return nil, err + } + var out []string + for _, e := range entries { + if e.Type().IsRegular() && !strings.HasPrefix(e.Name(), ".") { + out = append(out, filepath.Join(weights, e.Name())) + } + } + return out, nil +} + +func isWeightsFile(path string) bool { + switch strings.ToLower(filepath.Ext(path)) { + case ".safetensors", ".ckpt", ".pt", ".pth", ".bin", ".gguf", ".sft": + return true + } + return false +} + // LinkWeights shows ComfyUI the weights: a link in dir to the file, or to // every file of the directory. Links made before are taken away first, so // a replica started again on other weights sees only those. diff --git a/internal/comfy/gate_test.go b/internal/comfy/gate_test.go index 632be83..90e0cb7 100644 --- a/internal/comfy/gate_test.go +++ b/internal/comfy/gate_test.go @@ -261,3 +261,41 @@ func TestLLMGWsOwnNodesArePutInEveryReplica(t *testing.T) { t.Fatalf("%v %.80s", err, b) } } + +func TestASharedReplicaSeesEveryMembersWeights(t *testing.T) { + root := t.TempDir() + write := func(p string) string { + full := filepath.Join(root, p) + _ = os.MkdirAll(filepath.Dir(full), 0o755) + if err := os.WriteFile(full, []byte("w"), 0o644); err != nil { + t.Fatal(err) + } + return full + } + flux := write("flux1-schnell-fp8.safetensors") + write("sf3d/model.safetensors") + write("sf3d/README.md") + write("tripo/model.ckpt") + write("tripo/README.md") + dir := filepath.Join(root, "models") + if err := LinkAllWeights([]string{flux, filepath.Join(root, "sf3d"), filepath.Join(root, "tripo")}, dir); err != nil { + t.Fatal(err) + } + for _, want := range []string{"flux1-schnell-fp8.safetensors", "model.safetensors", "model.ckpt", "README.md"} { + if _, err := os.Lstat(filepath.Join(dir, want)); err != nil { + t.Errorf("%s not linked: %v", want, err) + } + } + // Two weights of one name: ComfyUI would load one for both. + write("other/model.safetensors") + err := LinkAllWeights([]string{filepath.Join(root, "sf3d"), filepath.Join(root, "other")}, dir) + if err == nil || !strings.Contains(err.Error(), "both named model.safetensors") { + t.Errorf("err %v", err) + } + if ms, err := parseShared(`[{"name":"sf3d","path":"/x","hf_repo":"stabilityai/stable-fast-3d"}]`); err != nil || ms[0].HFRepo == "" { + t.Errorf("parse %+v %v", ms, err) + } + if _, err := parseShared(`[{"name":"sf3d"}]`); err == nil { + t.Error("a member with no path was taken") + } +} diff --git a/internal/comfy/progress_test.go b/internal/comfy/progress_test.go index 49ff825..4161b91 100644 --- a/internal/comfy/progress_test.go +++ b/internal/comfy/progress_test.go @@ -60,7 +60,7 @@ func TestTheGateHearsTheWorkflowsItLetsThrough(t *testing.T) { defer comfy.Close() u, _ := url.Parse(comfy.URL) book := newProgressBook() - gate := httptest.NewServer(newGate(u, "", "m", book, "llmgw-gate-1")) + gate := httptest.NewServer(newGate(u, "", []string{"m"}, book, "llmgw-gate-1")) defer gate.Close() resp, err := http.Post(gate.URL+"/prompt", "application/json", strings.NewReader(`{"prompt":{},"client_id":"router"}`)) if err != nil { diff --git a/internal/control/catalog/catalog.go b/internal/control/catalog/catalog.go index 730d7a5..13ca4ff 100644 --- a/internal/control/catalog/catalog.go +++ b/internal/control/catalog/catalog.go @@ -409,6 +409,10 @@ type Model struct { // say that the repository has published different weights since, and // move the model to them. Origin *Origin `json:"origin,omitempty" koanf:"origin"` + + // Shared is set on an environment's hidden entry only: the members its + // replica answers for. Made by the catalog, never written. + Shared *SharedComfy `json:"shared,omitempty" koanf:"-"` } // RoleGuard is a judge model's role. @@ -650,6 +654,11 @@ const NodesEnv = "LLMGW_COMFYUI_NODES" // HFRepoEnv carries the repository a node loads the weights by name from. const HFRepoEnv = "LLMGW_COMFYUI_HF_REPO" +// SharedModelsEnv carries a shared replica's members — name, weights source, +// repository — to the agent, which fetches each and hands the gate where +// they landed. +const SharedModelsEnv = "LLMGW_COMFYUI_MODELS" + var hfRepo = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9][A-Za-z0-9._-]*$`) var ( @@ -778,6 +787,9 @@ func New(models []Model) (*Catalog, error) { } for i := range models { + if models[i].IsEnvironment() { + continue // the catalog's own, made again below + } m := models[i].normalize() if err := m.validate(); err != nil { return nil, err @@ -851,9 +863,25 @@ func New(models []Model) (*Catalog, error) { } canaryOf[target.Name] = name } - if _, err := c.environments(); err != nil { + envs, err := c.environments() + if err != nil { return nil, err } + for _, e := range envs { + if len(e.Models) < 2 { + continue + } + var members []Model + for _, n := range e.Models { + members = append(members, c.models[n]) + } + em := envModel(e, members) + if _, clash := c.models[em.Name]; clash { + return nil, fmt.Errorf("model %q: the name is the hidden entry of environment %q", em.Name, e.Name) + } + c.models[em.Name] = em + c.order = append(c.order, em.Name) + } return c, nil } @@ -1297,6 +1325,25 @@ func (m Model) EngineProfile(name string) (engines.Profile, bool) { } env[NodesEnv] = string(nodes) } + if p.Engine == EngineComfyUI && m.Shared != nil { + // Every member's weights, for the node to fetch and the gate to + // link: one replica answers for all of them. + type member struct { + Name string `json:"name"` + Source string `json:"source"` + HFRepo string `json:"hf_repo,omitempty"` + } + var ms []member + for _, sm := range m.Shared.Members { + ms = append(ms, member{Name: sm.Name, Source: sm.Source, HFRepo: sm.HFRepo}) + } + b, _ := json.Marshal(ms) + env = maps.Clone(env) + if env == nil { + env = map[string]string{} + } + env[SharedModelsEnv] = string(b) + } if p.Engine == EngineComfyUI && m.ComfyUI != nil && m.ComfyUI.HuggingFaceRepo != "" { env = maps.Clone(env) if env == nil { diff --git a/internal/control/catalog/environments.go b/internal/control/catalog/environments.go index 0686f29..1284fc8 100644 --- a/internal/control/catalog/environments.go +++ b/internal/control/catalog/environments.go @@ -75,7 +75,7 @@ func (c *Catalog) environments() ([]Environment, error) { for _, name := range c.order { m := c.models[name] env := m.ComfyEnvironment() - if env == "" { + if env == "" || m.IsEnvironment() { continue } if !envName.MatchString(env) { diff --git a/internal/control/catalog/shared.go b/internal/control/catalog/shared.go new file mode 100644 index 0000000..51ced8a --- /dev/null +++ b/internal/control/catalog/shared.go @@ -0,0 +1,110 @@ +package catalog + +import ( + "time" +) + +// A shared replica: one ComfyUI serving every model of an environment. +// +// Each ComfyUI model ran in a replica of its own, so on one card a picture +// model and a 3D model took turns, each stopping the other and loading +// ComfyUI and its weights from nothing. An environment of two or more models +// gets a hidden catalog entry of its own, whose replica carries every +// member's weights and custom nodes and answers for each of them; the +// members get no replicas of their own. Placement, idleness, eviction and +// adoption then work as for any model, on that entry. + +// RoleEnvironment is the role of an environment's hidden entry. +const RoleEnvironment = "environment" + +// EnvModelPrefix names an environment's hidden entry: comfyui-env-. +const EnvModelPrefix = "comfyui-env-" + +// SharedComfy is what an environment's entry carries of its members. +type SharedComfy struct { + Environment string `json:"environment"` + Members []SharedMember `json:"members"` +} + +// SharedMember is one model a shared replica answers for. +type SharedMember struct { + Name string `json:"name"` + Aliases []string `json:"aliases,omitempty"` + // Source is the member's weights, which the replica's node fetches. + Source string `json:"source"` + // HFRepo is the repository its node loads the weights by name from. + HFRepo string `json:"hf_repo,omitempty"` + // Workflow and Checkpoint are what a request for it is sent. + Workflow string `json:"workflow"` + Checkpoint string `json:"checkpoint"` +} + +// IsEnvironment reports whether a model is an environment's hidden entry. +func (m Model) IsEnvironment() bool { return m.Role == RoleEnvironment } + +// SharedBy is the hidden entry whose replica serves a model, when it is in +// an environment of two or more. +func (c *Catalog) SharedBy(model string) (string, bool) { + m, ok := c.Resolve(model) + if !ok || m.IsEnvironment() || m.ComfyEnvironment() == "" { + return "", false + } + name := EnvModelPrefix + m.ComfyEnvironment() + if _, ok := c.models[name]; !ok { + return "", false + } + return name, true +} + +// envModel is the hidden entry for an environment of two or more models: on +// the first member's profiles, needing the most memory any of them does, +// on demand for the longest idle timeout of theirs. +func envModel(env Environment, members []Model) Model { + first := members[0] + profs := make([]Profile, len(first.Profiles)) + copy(profs, first.Profiles) + mem := 0 + for _, m := range members { + for _, p := range m.Profiles { + mem = max(mem, p.MinMemoryGB) + } + } + shared := &SharedComfy{Environment: env.Name} + var idle time.Duration + minN, maxN := 0, 1 + class := "" + for _, m := range members { + workflow, ckpt := m.ComfyWorkflow() + hf := "" + if m.ComfyUI != nil { + hf = m.ComfyUI.HuggingFaceRepo + } + shared.Members = append(shared.Members, SharedMember{Name: m.Name, Aliases: m.Aliases, + Source: m.Source, HFRepo: hf, Workflow: workflow, Checkpoint: ckpt}) + idle = max(idle, m.Scaling.IdleTimeout) + minN, maxN = max(minN, m.Scaling.Min), max(maxN, m.Scaling.Max) + if ClassRank(m.EffectiveClass()) > ClassRank(class) || class == "" { + class = m.EffectiveClass() + } + } + for i := range profs { + profs[i].MinMemoryGB = max(profs[i].MinMemoryGB, mem) + profs[i].ComfyUI = nil // the environment's nodes, below + profs[i].Source = "" + } + return Model{ + Name: EnvModelPrefix + env.Name, + Role: RoleEnvironment, + Description: "the shared ComfyUI replica of environment " + env.Name, + Source: first.Source, + ComfyUI: &ComfyUI{Environment: env.Name, Nodes: env.Nodes, Output: OutputEnvironment}, + Profiles: profs, + Scaling: Scaling{Min: minN, Max: maxN, IdleTimeout: idle}, + Class: class, + Shared: shared, + } +} + +// OutputEnvironment is an environment entry's output: it makes nothing of +// its own, its members do. +const OutputEnvironment = "environment" diff --git a/internal/control/catalog/shared_test.go b/internal/control/catalog/shared_test.go new file mode 100644 index 0000000..8f7665d --- /dev/null +++ b/internal/control/catalog/shared_test.go @@ -0,0 +1,51 @@ +package catalog + +import ( + "encoding/json" + "strings" + "testing" + "time" +) + +func TestAnEnvironmentOfTwoGetsOneHiddenEntryThatServesBoth(t *testing.T) { + // flux and a mesh model on one card took turns, each loading ComfyUI + // from nothing; in one environment one replica answers for both. + flux := comfyModel("/w/flux.safetensors", &ComfyUI{Environment: "art", Workflow: "flux-schnell"}) + flux.Scaling = Scaling{Max: 1, IdleTimeout: 15 * time.Minute} + sf3d := inEnv("sf3d", "art", "/w/sf3d/", CustomNode{Repo: "https://example.com/sf3d", Commit: commitA}) + sf3d.ComfyUI.HuggingFaceRepo = "stabilityai/stable-fast-3d" + sf3d.ComfyUI.Checkpoint = "model.safetensors" + sf3d.Scaling = Scaling{Max: 1, IdleTimeout: 10 * time.Minute} + cat, err := New([]Model{flux, sf3d, inEnv("alone", "solo", "/w/alone.ckpt")}) + if err != nil { + t.Fatal(err) + } + host, ok := cat.SharedBy("sf3d") + if !ok || host != "comfyui-env-art" { + t.Fatalf("sf3d shared by %q %v", host, ok) + } + if _, ok := cat.SharedBy("alone"); ok { + t.Error("an environment of one got a hidden entry") + } + em, _ := cat.Resolve(host) + if !em.IsEnvironment() || em.MakesImages() || em.MakesMeshes() || em.Scaling.IdleTimeout != 15*time.Minute || + len(em.Shared.Members) != 2 || em.Shared.Members[1].HFRepo != "stabilityai/stable-fast-3d" { + t.Errorf("entry %+v", em) + } + p, ok := em.EngineProfile(em.Profiles[0].Name) + if !ok { + t.Fatal("no engine profile") + } + var ms []map[string]string + if err := json.Unmarshal([]byte(p.Env[SharedModelsEnv]), &ms); err != nil || len(ms) != 2 || ms[1]["source"] != "/w/sf3d/" { + t.Errorf("members for the agent %q %v", p.Env[SharedModelsEnv], err) + } + if !strings.Contains(p.Env[NodesEnv], "sf3d") { + t.Errorf("nodes %q", p.Env[NodesEnv]) + } + // A catalog built again from its own models is the same catalog. + again, err := New(cat.Models()) + if err != nil || len(again.Models()) != len(cat.Models()) { + t.Errorf("rebuilt: %v, %d models vs %d", err, len(again.Models()), len(cat.Models())) + } +} diff --git a/internal/control/control.go b/internal/control/control.go index 81b8eff..c03a7d2 100644 --- a/internal/control/control.go +++ b/internal/control/control.go @@ -371,6 +371,35 @@ func (c *Controller) endpoints() []endpoint.Info { Weight: rep.Weight, Priority: prof.Priority, } + if m.Shared != nil { + // One replica answering for every member of its environment: + // offered under their names, each sent its own workflow and + // checkpoint. Its own hidden name nobody asks for. + info.Models = nil + info.Comfy = map[string]endpoint.ComfyRoute{} + ok := true + for _, sm := range m.Shared.Members { + w, err := c.workflow(sm.Workflow) + if err != nil { + c.log.Warn("a shared ComfyUI member's workflow could not be read; its replica is not offered", + "model", sm.Name, "replica", rep.ID, "err", err) + ok = false + break + } + route := endpoint.ComfyRoute{Workflow: w, Checkpoint: sm.Checkpoint} + for _, n := range append([]string{sm.Name}, sm.Aliases...) { + info.Models = append(info.Models, n) + info.Comfy[n] = route + } + if info.Workflow == "" { + info.Workflow, info.UpstreamModel = w, sm.Checkpoint + } + } + if ok { + out = append(out, info) + } + continue + } if prof.Engine == catalog.EngineComfyUI { workflow, ckpt := m.ComfyWorkflowFor(prof) w, err := c.workflow(workflow) diff --git a/internal/control/reconciler/demand.go b/internal/control/reconciler/demand.go index a656283..ad7a404 100644 --- a/internal/control/reconciler/demand.go +++ b/internal/control/reconciler/demand.go @@ -55,7 +55,7 @@ func (r *Reconciler) IdleFor(id string) (time.Duration, bool) { // from "there is genuinely nothing serving this", which are different answers // to give a client. func (r *Reconciler) Want(model string) bool { - m, ok := r.cat.Load().Resolve(model) + m, ok := r.cat.Load().Resolve(r.hostOf(model)) if !ok || !m.Scaling.OnDemand() { return false } @@ -127,6 +127,9 @@ func (r *Reconciler) baseFor(m catalog.Model) int { if m.IsExternal() { return 0 // served elsewhere: nothing to place, whatever its scaling says } + if _, shared := r.cat.Load().SharedBy(m.Name); shared { + return 0 // its environment's replica serves it + } if lo := r.minFor(m); !m.Scaling.OnDemand() || lo > 0 { return lo } @@ -199,6 +202,17 @@ type Load struct { // were measured to hold, so scaling down does not land on the edge. func (r *Reconciler) RecordLoad(loads map[string]Load) { now := r.now() + // A shared replica's load is its members' together. + for name, l := range loads { + if host := r.hostOf(name); host != name { + h := loads[host] + h.Replicas = max(h.Replicas, l.Replicas) + h.Running += l.Running + h.Waiting += l.Waiting + loads[host] = h + delete(loads, name) + } + } // The profile each model's replicas run on, for its measurement; read // before the loop, which holds the lock. running := map[string]string{} @@ -280,3 +294,12 @@ func (r *Reconciler) RecordLoad(loads map[string]Load) { // after a scale-down may be. Below one so that the load that was there a // minute ago does not put them straight back over. const scaleDownHeadroom = 0.7 + +// hostOf is the model whose replicas serve name: its environment's hidden +// entry when it shares one, else itself. +func (r *Reconciler) hostOf(name string) string { + if host, ok := r.cat.Load().SharedBy(name); ok { + return host + } + return name +} diff --git a/internal/control/reconciler/evict.go b/internal/control/reconciler/evict.go index 5214c19..139d1a7 100644 --- a/internal/control/reconciler/evict.go +++ b/internal/control/reconciler/evict.go @@ -88,6 +88,8 @@ func (r *Reconciler) forgetWantIfGone(model, stopped string) { // longest idle first. func (r *Reconciler) evictable(wantedBy string) []evictionCandidate { cat := r.cat.Load() + // A member never evicts the replica that serves it. + wantedBy = r.hostOf(wantedBy) now := r.now() wantRank := -1 diff --git a/internal/control/reconciler/health.go b/internal/control/reconciler/health.go index de9f43c..006b3eb 100644 --- a/internal/control/reconciler/health.go +++ b/internal/control/reconciler/health.go @@ -130,7 +130,7 @@ func (r *Reconciler) ModelHealth() []ModelHealth { // Unservable says why a model cannot be served at all, so a request for it is // told now rather than asked to retry something that will never come up. func (r *Reconciler) Unservable(model string) (string, bool) { - m, ok := r.cat.Load().Resolve(model) + m, ok := r.cat.Load().Resolve(r.hostOf(model)) if !ok { return "", false } @@ -176,7 +176,7 @@ func (r *Reconciler) Unservable(model string) (string, bool) { // on another model; eta is zero when when the holder lets go is not known. func (r *Reconciler) WaitFor(model string) (heldBy string, eta time.Duration, ok bool) { cat := r.cat.Load() - m, found := cat.Resolve(model) + m, found := cat.Resolve(r.hostOf(model)) if !found || !m.Scaling.OnDemand() { return "", 0, false } diff --git a/internal/control/reconciler/shared_test.go b/internal/control/reconciler/shared_test.go new file mode 100644 index 0000000..2262eb6 --- /dev/null +++ b/internal/control/reconciler/shared_test.go @@ -0,0 +1,57 @@ +package reconciler + +import ( + "testing" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog" +) + +func sharedCatalog(t *testing.T) *catalog.Catalog { + t.Helper() + mk := func(name, out string) catalog.Model { + return catalog.Model{Name: name, Source: "/w/" + name + ".safetensors", + ComfyUI: &catalog.ComfyUI{Environment: "art", Workflow: "flux-schnell", Output: out}, + Scaling: catalog.Scaling{Max: 1, IdleTimeout: time.Hour}, + Profiles: []catalog.Profile{{Name: "p", Engine: catalog.EngineComfyUI, GPUClass: "h100-80g", TP: 1}}} + } + cat, err := catalog.New([]catalog.Model{mk("flux", catalog.OutputPictures), mk("sf3d", catalog.OutputMeshes)}) + if err != nil { + t.Fatal(err) + } + return cat +} + +func TestAMemberAskedForStartsItsEnvironmentsOneReplica(t *testing.T) { + // A picture model and a 3D model on one card took turns, each from + // nothing: in one environment, asking for either starts one replica + // that serves both, and neither gets one of its own. + h := newHarness(t, testModel(1), gpuNode("1", 1)) + h.r.cat.Store(sharedCatalog(t)) + if !h.r.Want("sf3d") { + t.Fatal("a member was not wanted") + } + h.r.Reconcile(t.Context()) + h.markReady() + h.r.Want("flux") + h.r.Reconcile(t.Context()) + var models []string + for _, rep := range h.r.Replicas() { + if rep.State.Counts() { + models = append(models, rep.Model) + } + } + if len(models) != 1 || models[0] != "comfyui-env-art" { + t.Fatalf("replicas %v", models) + } + // The member's sibling never evicts the replica serving it. + reps := h.r.Replicas() + h.r.RecordUsage(map[string]time.Time{reps[0].ID: h.clock.now()}) + h.advance(DefaultConfig().MinLifetime + DefaultConfig().EvictionGrace + time.Minute) + if got := h.r.evictable("flux"); len(got) != 0 { + t.Errorf("a member was offered its own environment's replica: %+v", got) + } + if _, _, ok := h.r.WaitFor("flux"); ok { + t.Error("a member served by a replica that is up was said to be waiting") + } +} diff --git a/internal/router/canary.go b/internal/router/canary.go index 99f1a6f..60661e2 100644 --- a/internal/router/canary.go +++ b/internal/router/canary.go @@ -111,7 +111,7 @@ func (r *Router) isCanaryTarget(name string) bool { // hiddenModel is a model clients cannot name: a judge, or a canary. func (r *Router) hiddenModel(name string) bool { - return r.isGuardModel(name) || r.isCanaryTarget(name) + return r.isGuardModel(name) || r.isCanaryTarget(name) || r.modelMap()[name].IsEnvironment() } // setVariantHeader tells the client which side served it. diff --git a/internal/router/endpoint/endpoint.go b/internal/router/endpoint/endpoint.go index 049322d..000d4f2 100644 --- a/internal/router/endpoint/endpoint.go +++ b/internal/router/endpoint/endpoint.go @@ -51,6 +51,25 @@ type Info struct { // Workflow is the ComfyUI workflow sent for an endpoint whose engine is // EngineComfyUI: a built-in's name or the workflow's JSON. Workflow string + // Comfy is, for a shared ComfyUI replica that answers for several + // models, each one's workflow and checkpoint, by every name it goes by. + // Nil for a replica of one model, which Workflow and UpstreamModel say. + Comfy map[string]ComfyRoute +} + +// ComfyRoute is what a request for one model of a shared ComfyUI replica is +// sent. +type ComfyRoute struct { + Workflow string + Checkpoint string +} + +// ComfyFor is the workflow and checkpoint a request for model is sent here. +func (i Info) ComfyFor(model string) (workflow, checkpoint string) { + if r, ok := i.Comfy[model]; ok { + return r.Workflow, r.Checkpoint + } + return i.Workflow, i.UpstreamModel } // EngineExternal marks an endpoint that is a provider outside the cluster. diff --git a/internal/router/images.go b/internal/router/images.go index 35bce2b..3c975b2 100644 --- a/internal/router/images.go +++ b/internal/router/images.go @@ -182,7 +182,8 @@ func (r *Router) generateImages(ctx context.Context, w http.ResponseWriter, req Model string `json:"model"` } _ = json.Unmarshal(body, &upstream) - tpl, err := comfy.Load(ep.Workflow) + wf, _ := ep.ComfyFor(req.Model) + tpl, err := comfy.Load(wf) if err != nil { return proxy.Outcome{}, fmt.Errorf("%s: %w", ep.ID, err) } diff --git a/internal/router/meshes.go b/internal/router/meshes.go index e13f847..3342d78 100644 --- a/internal/router/meshes.go +++ b/internal/router/meshes.go @@ -101,7 +101,8 @@ func (r *Router) generateMeshes(ctx context.Context, w http.ResponseWriter, req if err != nil { return proxy.Outcome{}, err } - tpl, err := comfy.Load(ep.Workflow) + wf, _ := ep.ComfyFor(req.Model) + tpl, err := comfy.Load(wf) if err != nil { return proxy.Outcome{}, fmt.Errorf("%s: %w", ep.ID, err) } diff --git a/internal/router/router.go b/internal/router/router.go index cfea90c..8717a93 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -1134,8 +1134,10 @@ func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *pa trace.WithAttributes(attribute.String("llmgw.endpoint", dec.Endpoint.ID), attribute.String("url.full", dec.Endpoint.URL+req.Path))) body := req.UpstreamBody - if dec.Endpoint.UpstreamModel != "" && req.ContentType == "" { - body = withModel(body, dec.Endpoint.UpstreamModel) + // A shared ComfyUI replica is told which of its models' checkpoint; + // any other endpoint its own upstream name. + if _, up := dec.Endpoint.ComfyFor(model); up != "" && req.ContentType == "" { + body = withModel(body, up) } // llama-server is asked for the model with the adapter applied to // this request; vLLM serves each adapter as a model of its own name, diff --git a/internal/router/workflows.go b/internal/router/workflows.go index 433e39e..94362c6 100644 --- a/internal/router/workflows.go +++ b/internal/router/workflows.go @@ -297,7 +297,8 @@ func (r *Router) generateRun(ctx context.Context, w http.ResponseWriter, req *pa if err != nil { return proxy.Outcome{}, err } - tpl, err := comfy.Load(ep.Workflow) + wf, _ := ep.ComfyFor(req.Model) + tpl, err := comfy.Load(wf) if err != nil { return proxy.Outcome{}, fmt.Errorf("%s: %w", ep.ID, err) }