diff --git a/core/application/studio_connectors.go b/core/application/studio_connectors.go index 6c59454aca..36d642c95b 100644 --- a/core/application/studio_connectors.go +++ b/core/application/studio_connectors.go @@ -2,6 +2,7 @@ package application import ( "context" + "encoding/json" "errors" "github.com/mudler/xlog" @@ -59,19 +60,33 @@ func (t connectorTransport) Call(ctx context.Context, scope studio.ConnectorCall // RegisterAddonConnector registers the host operation with which a connector // add-on pushes a message. It registers nothing when this process has no host // operation channel or no connector manager. The message goes through the same -// function as a message the manager pulls itself. +// function as a message the manager pulls itself. A dedicated add-on owner +// refuses it instead. func (a *Application) RegisterAddonConnector() error { if a.addonHostOps == nil || a.studioConnectors == nil { return nil } + if a.ownsAddonsOnly() { + // A pushed message would start its chat turn here, and the owner has no + // models. Nothing forwards it to a frontend, so refuse it and let the + // add-on see the failure. + xlog.Warn("Connectors: this add-on owner refuses connector.inbound; push connector add-ons need a combined or distributed process") + return a.addonHostOps.Broker.Register(hostops.Spec{Op: contract.ConnectorInboundOp, Scope: hostops.ScopeInbound, Handler: refuseConnectorInbound}) + } return hostops.RegisterConnector(a.addonHostOps.Broker, a.studioConnectors) } +// refuseConnectorInbound is connector.inbound on a dedicated add-on owner. +func refuseConnectorInbound(context.Context, hostops.Call) (json.RawMessage, error) { + return nil, hostops.ErrForbidden +} + // StudioConnectors returns the connector manager of the chats, or nil. func (a *Application) StudioConnectors() *studio.Connectors { return a.studioConnectors } // startStudioConnectors builds the connector manager and starts its loop. In a -// cluster the polling runs on the one frontend that holds the advisory lock. +// cluster the polling runs on the one frontend that holds the advisory lock. A +// dedicated add-on owner builds the manager but never starts the loop. func (a *Application) startStudioConnectors(svc *studio.Service) { if a.studioConnStore == nil { return @@ -96,9 +111,16 @@ func (a *Application) startStudioConnectors(svc *studio.Service) { return } svc.SetConnectors(c) + a.studioConnectorsHook.Store(c) + if a.ownsAddonsOnly() { + // The manager stays: frontends get the stored bot credentials from it + // and add-ons push messages to it. Only the polling is left to the + // frontends, which have the models. + a.studioConnectors = c + return + } ctx := a.backgroundContext() ctx, cancel := context.WithCancel(ctx) a.studioConnectors, a.studioConnectorsStop = c, cancel - a.studioConnectorsHook.Store(c) go c.Run(ctx) } diff --git a/core/application/studio_owner_role_test.go b/core/application/studio_owner_role_test.go new file mode 100644 index 0000000000..b08539c540 --- /dev/null +++ b/core/application/studio_owner_role_test.go @@ -0,0 +1,84 @@ +// SPDX-License-Identifier: MIT +package application + +import ( + "context" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/services/addons/hostops" + addonruntime "github.com/mudler/LocalAI/core/services/addons/runtime" + "github.com/mudler/LocalAI/core/services/studio" + "github.com/mudler/LocalAI/pkg/addons/contract" + "github.com/mudler/LocalAI/pkg/system" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("studio background work on the add-on owner", func() { + newApp := func(opts ...config.AppOption) *Application { + ctx, cancel := context.WithCancel(context.Background()) + DeferCleanup(cancel) + state, err := system.GetSystemState(system.WithModelPath(GinkgoT().TempDir()), system.WithBackendPath(GinkgoT().TempDir())) + Expect(err).NotTo(HaveOccurred()) + base := []config.AppOption{config.WithContext(ctx), config.WithSystemState(state), config.DisableMetricsEndpoint, + config.WithDisableLocalAIAssistant(true), config.WithDisableStats(true), config.WithDataPath(GinkgoT().TempDir())} + app, err := New(append(base, opts...)...) + if app != nil { + DeferCleanup(func() { Expect(app.Shutdown()).To(Succeed()) }) + } + Expect(err).NotTo(HaveOccurred()) + return app + } + + It("runs connector polling, triggers, loops and agent jobs on a combined or frontend process", func() { + for _, c := range []config.AddonConfig{{}, {Frontend: true}} { + app := newApp(config.WithAddons(c)) + Expect(app.studioConnectorsStop).NotTo(BeNil(), "the connector loop did not start") + Expect(app.studioTriggersStop).NotTo(BeNil(), "the trigger loop did not start") + Expect(app.studioLoopsStop).NotTo(BeNil(), "the loop host did not start") + Expect(app.studioJobsStop).NotTo(BeNil(), "the agent job host did not start") + } + }) + + It("runs none of them on a dedicated add-on owner, which keeps the connector authority and inbox", func() { + app := newApp(config.WithAddons(config.AddonConfig{DedicatedOwner: true})) + Expect(app.StudioService()).NotTo(BeNil()) + Expect(app.studioConnectorsStop).To(BeNil(), "the owner polls connectors") + Expect(app.studioTriggersStop).To(BeNil(), "the owner runs triggers") + Expect(app.studioLoopsStop).To(BeNil(), "the owner runs loops") + Expect(app.studioJobsStop).To(BeNil(), "the owner runs agent jobs") + // Frontends reach the connector credentials and add-ons push messages + // through the manager, so it stays. + Expect(app.StudioConnectors()).NotTo(BeNil()) + }) +}) + +var _ = Describe("connector.inbound on the add-on owner", func() { + newBroker := func() *hostops.Broker { + return hostops.New(hostops.NewInFlight(), nil, contract.DefaultLimits()) + } + + It("is registered and refuses a pushed message on a dedicated owner", func() { + b := newBroker() + a := &Application{ + applicationConfig: &config.ApplicationConfig{Addons: config.AddonConfig{DedicatedOwner: true}}, + addonHostOps: &addonruntime.HostOps{Broker: b}, + studioConnectors: &studio.Connectors{}, + } + Expect(a.RegisterAddonConnector()).To(Succeed()) + Expect(b.Ops()).To(ContainElement(contract.ConnectorInboundOp)) + _, err := refuseConnectorInbound(context.Background(), hostops.Call{}) + Expect(err).To(MatchError(hostops.ErrForbidden)) + }) + + It("takes the message into the chats on a combined process", func() { + b := newBroker() + a := &Application{ + applicationConfig: &config.ApplicationConfig{}, + addonHostOps: &addonruntime.HostOps{Broker: b}, + studioConnectors: &studio.Connectors{}, + } + Expect(a.RegisterAddonConnector()).To(Succeed()) + Expect(b.Ops()).To(ContainElement(contract.ConnectorInboundOp)) + }) +}) diff --git a/core/application/studio_service.go b/core/application/studio_service.go index ff41700b56..27f74fd3af 100644 --- a/core/application/studio_service.go +++ b/core/application/studio_service.go @@ -199,12 +199,23 @@ func (a *Application) initStudioService() { return } a.startStudioConnectors(svc) - a.startStudioTriggers(svc) - a.startStudioLoops(svc) - a.startStudioJobs(svc) + if a.ownsAddonsOnly() { + xlog.Info("Chats: this is the add-on owner; connectors, triggers, loops and agent jobs run on the frontends") + } else { + a.startStudioTriggers(svc) + a.startStudioLoops(svc) + a.startStudioJobs(svc) + } a.prepareStudioConversion(svc) } +// ownsAddonsOnly says whether this process is a dedicated add-on owner +// (--addons-role=owner). It hosts the add-ons and has no models, so the +// background work that starts chat turns belongs to the frontends. +func (a *Application) ownsAddonsOnly() bool { + return a.applicationConfig.Addons.DedicatedOwner +} + // errStudioStarting is the error of a command hook called before the chat // service is up. No chat turn runs then, so it is not expected. var errStudioStarting = errors.New("chats: the service is starting")