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-<name>: 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) <noreply@anthropic.com>
This commit is contained in:
Ilker ManapandClaude Opus 5.5 committed 2026-10-14 13:55:34 +02:00
1 parent 7c53380e68
commit aa531b1053
22 files changed
+614 -18

No files matched your search

+8
View File
@@ -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-<name>`, 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,
+49
View File
@@ -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.",
+26
View File
@@ -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.
+6 -1
View File
@@ -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"
+132 -5
View File
@@ -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.
+38
View File
@@ -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")
}
}
+1 -1
View File
@@ -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 {
+48 -1
View File
@@ -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 {
+1 -1
View File
@@ -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) {
+110
View File
@@ -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-<name>.
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"
+51
View File
@@ -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()))
}
}
+29
View File
@@ -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)
+24 -1
View File
@@ -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
}
+2
View File
@@ -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
+2 -2
View File
@@ -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
}
@@ -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")
}
}
+1 -1
View File
@@ -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.
+19
View File
@@ -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.
+2 -1
View File
@@ -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)
}
+2 -1
View File
@@ -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)
}
+4 -2
View File
@@ -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,
+2 -1
View File
@@ -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)
}