toki
c46874055a
- Add edge node unit tests and transport package - Add node test client and related artifacts - Update bootstrap, node, and config modules - Add proto generated files - Update Makefile and configuration files
96 lines
2.4 KiB
Go
96 lines
2.4 KiB
Go
// Package bootstrap wires the IOP Node application using go.uber.org/fx.
|
|
package bootstrap
|
|
|
|
import (
|
|
"context"
|
|
|
|
"go.uber.org/fx"
|
|
"go.uber.org/zap"
|
|
|
|
"iop/apps/node/internal/adapters"
|
|
"iop/apps/node/internal/adapters/cli"
|
|
"iop/apps/node/internal/adapters/mock"
|
|
"iop/apps/node/internal/adapters/ollama"
|
|
"iop/apps/node/internal/adapters/vllm"
|
|
"iop/apps/node/internal/node"
|
|
"iop/apps/node/internal/router"
|
|
"iop/apps/node/internal/runtime"
|
|
"iop/apps/node/internal/store"
|
|
"iop/apps/node/internal/transport"
|
|
"iop/packages/config"
|
|
"iop/packages/observability"
|
|
)
|
|
|
|
// Module returns the fx options that wire the node application.
|
|
func Module(cfg *config.NodeConfig) fx.Option {
|
|
return fx.Options(
|
|
fx.Provide(
|
|
func() *config.NodeConfig { return cfg },
|
|
|
|
func(cfg *config.NodeConfig) (*zap.Logger, error) {
|
|
return observability.NewLogger(cfg.Logging.Level)
|
|
},
|
|
|
|
func(cfg *config.NodeConfig, logger *zap.Logger) (*store.Store, error) {
|
|
return store.New(cfg.SQLite.DSN, logger)
|
|
},
|
|
|
|
func(cfg *config.NodeConfig, logger *zap.Logger) *adapters.Registry {
|
|
reg := adapters.NewRegistry()
|
|
|
|
// mock adapter is always registered
|
|
reg.Register(mock.New(logger))
|
|
|
|
if cfg.Adapters.Ollama.Enabled {
|
|
reg.Register(ollama.New(cfg.Adapters.Ollama, logger))
|
|
}
|
|
if cfg.Adapters.Vllm.Enabled {
|
|
reg.Register(vllm.New(cfg.Adapters.Vllm, logger))
|
|
}
|
|
if cfg.Adapters.CLI.Enabled {
|
|
reg.Register(cli.New(cfg.Adapters.CLI, logger))
|
|
}
|
|
return reg
|
|
},
|
|
|
|
func(reg *adapters.Registry, logger *zap.Logger) runtime.Router {
|
|
return router.New(reg, logger)
|
|
},
|
|
|
|
func(
|
|
cfg *config.NodeConfig,
|
|
r runtime.Router,
|
|
reg *adapters.Registry,
|
|
st *store.Store,
|
|
logger *zap.Logger,
|
|
) *node.Node {
|
|
return node.New(cfg, r, reg, st, logger)
|
|
},
|
|
),
|
|
|
|
fx.Invoke(func(lc fx.Lifecycle, n *node.Node, cfg *config.NodeConfig, logger *zap.Logger) {
|
|
var sess *transport.Session
|
|
lc.Append(fx.Hook{
|
|
OnStart: func(ctx context.Context) error {
|
|
s, err := transport.DialEdge(ctx, cfg.Transport.EdgeAddr, n, cfg.Node.ID, logger)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
sess = s
|
|
go func() {
|
|
if err := observability.ServeMetrics(cfg.Metrics.Port); err != nil {
|
|
logger.Warn("metrics server exited", zap.Error(err))
|
|
}
|
|
}()
|
|
return nil
|
|
},
|
|
OnStop: func(_ context.Context) error {
|
|
if sess != nil {
|
|
return sess.Close()
|
|
}
|
|
return nil
|
|
},
|
|
})
|
|
}),
|
|
)
|
|
}
|