4cdb68907e
Update utility modules: - File utilities with secure file operations - Environment pool with resource tracking - Error types with scheduler error categories - Logging with audit context support - Network/SSH with connection pooling - Privacy/PII handling with tenant boundaries - Resource manager with scheduler allocation - Security monitor with audit integration - Tracking plugins (MLflow, TensorBoard) with auth - Crypto signing with tenant keys - Database init with multi-user support
116 lines
3.1 KiB
Go
116 lines
3.1 KiB
Go
package factory
|
|
|
|
import (
|
|
"fmt"
|
|
|
|
"github.com/jfraeys/fetch_ml/internal/container"
|
|
"github.com/jfraeys/fetch_ml/internal/logging"
|
|
"github.com/jfraeys/fetch_ml/internal/tracking"
|
|
"github.com/jfraeys/fetch_ml/internal/tracking/plugins"
|
|
)
|
|
|
|
// PluginConfig represents the configuration for a single plugin.
|
|
type PluginConfig struct {
|
|
Settings map[string]any `toml:"settings" yaml:"settings"`
|
|
Image string `toml:"image" yaml:"image"`
|
|
Mode string `toml:"mode" yaml:"mode"`
|
|
LogBasePath string `toml:"log_base_path" yaml:"log_base_path"`
|
|
ArtifactPath string `toml:"artifact_path" yaml:"artifact_path"`
|
|
Enabled bool `toml:"enabled" yaml:"enabled"`
|
|
}
|
|
|
|
// PluginFactory is a function that creates a Plugin instance.
|
|
type PluginFactory func(
|
|
logger *logging.Logger,
|
|
podman *container.PodmanManager,
|
|
cfg PluginConfig,
|
|
) (tracking.Plugin, error)
|
|
|
|
// PluginLoader uses dependency injection to load plugins.
|
|
type PluginLoader struct {
|
|
logger *logging.Logger
|
|
podman *container.PodmanManager
|
|
factories map[string]PluginFactory
|
|
}
|
|
|
|
// NewPluginLoader creates a new PluginLoader.
|
|
func NewPluginLoader(logger *logging.Logger, podman *container.PodmanManager) *PluginLoader {
|
|
loader := &PluginLoader{
|
|
logger: logger,
|
|
podman: podman,
|
|
factories: make(map[string]PluginFactory),
|
|
}
|
|
|
|
// Register default factories
|
|
loader.RegisterFactory("mlflow", createMLflowPlugin)
|
|
loader.RegisterFactory("tensorboard", createTensorBoardPlugin)
|
|
loader.RegisterFactory("wandb", createWandbPlugin)
|
|
|
|
return loader
|
|
}
|
|
|
|
// RegisterFactory allows external packages to register new plugin types (Marketplace ready).
|
|
func (l *PluginLoader) RegisterFactory(name string, factory PluginFactory) {
|
|
l.factories[name] = factory
|
|
}
|
|
|
|
// LoadPlugins loads plugins from the provided configuration map and registers them.
|
|
func (l *PluginLoader) LoadPlugins(
|
|
plugins map[string]PluginConfig,
|
|
registry *tracking.Registry,
|
|
) error {
|
|
for name, pluginCfg := range plugins {
|
|
if !pluginCfg.Enabled {
|
|
continue
|
|
}
|
|
|
|
factory, ok := l.factories[name]
|
|
if !ok {
|
|
l.logger.Warn("unknown plugin type", "name", name)
|
|
continue
|
|
}
|
|
|
|
plugin, err := factory(l.logger, l.podman, pluginCfg)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to create plugin %s: %w", name, err)
|
|
}
|
|
|
|
registry.Register(plugin)
|
|
l.logger.Info("plugin loaded", "name", name, "mode", pluginCfg.Mode)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// Factory Implementations
|
|
func createMLflowPlugin(
|
|
logger *logging.Logger,
|
|
podman *container.PodmanManager,
|
|
cfg PluginConfig,
|
|
) (tracking.Plugin, error) {
|
|
opts := plugins.MLflowOptions{
|
|
Image: cfg.Image,
|
|
ArtifactBasePath: cfg.ArtifactPath,
|
|
}
|
|
return plugins.NewMLflowPlugin(logger, podman, opts)
|
|
}
|
|
|
|
func createTensorBoardPlugin(
|
|
logger *logging.Logger,
|
|
podman *container.PodmanManager,
|
|
cfg PluginConfig,
|
|
) (tracking.Plugin, error) {
|
|
opts := plugins.TensorBoardOptions{
|
|
Image: cfg.Image,
|
|
LogBasePath: cfg.LogBasePath,
|
|
}
|
|
return plugins.NewTensorBoardPlugin(logger, podman, opts)
|
|
}
|
|
|
|
func createWandbPlugin(
|
|
logger *logging.Logger,
|
|
_ *container.PodmanManager,
|
|
_ PluginConfig,
|
|
) (tracking.Plugin, error) {
|
|
return plugins.NewWandbPlugin(), nil
|
|
}
|