547 lines
17 KiB
Go
547 lines
17 KiB
Go
package worker
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"log"
|
|
"log/slog"
|
|
"math"
|
|
"net/http"
|
|
"os"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"runtime"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/jfraeys/fetch_ml/internal/container"
|
|
"github.com/jfraeys/fetch_ml/internal/envpool"
|
|
"github.com/jfraeys/fetch_ml/internal/jupyter"
|
|
"github.com/jfraeys/fetch_ml/internal/logging"
|
|
"github.com/jfraeys/fetch_ml/internal/metrics"
|
|
"github.com/jfraeys/fetch_ml/internal/network"
|
|
"github.com/jfraeys/fetch_ml/internal/queue"
|
|
"github.com/jfraeys/fetch_ml/internal/resources"
|
|
"github.com/jfraeys/fetch_ml/internal/tracking"
|
|
"github.com/jfraeys/fetch_ml/internal/tracking/factory"
|
|
trackingplugins "github.com/jfraeys/fetch_ml/internal/tracking/plugins"
|
|
"github.com/prometheus/client_golang/prometheus"
|
|
"github.com/prometheus/client_golang/prometheus/collectors"
|
|
"github.com/prometheus/client_golang/prometheus/promhttp"
|
|
)
|
|
|
|
// MLServer wraps network.SSHClient for backward compatibility.
|
|
type MLServer struct {
|
|
*network.SSHClient
|
|
}
|
|
|
|
// JupyterManager is the subset of the Jupyter service manager used by the worker.
|
|
// It exists to keep task execution testable.
|
|
type JupyterManager interface {
|
|
StartService(ctx context.Context, req *jupyter.StartRequest) (*jupyter.JupyterService, error)
|
|
StopService(ctx context.Context, serviceID string) error
|
|
RemoveService(ctx context.Context, serviceID string, purge bool) error
|
|
RestoreWorkspace(ctx context.Context, name string) (string, error)
|
|
ListServices() []*jupyter.JupyterService
|
|
}
|
|
|
|
// isValidName validates that input strings contain only safe characters.
|
|
// isValidName checks if the input string is a valid name.
|
|
func isValidName(input string) bool {
|
|
return len(input) > 0 && len(input) < 256
|
|
}
|
|
|
|
// NewMLServer creates a new ML server connection.
|
|
// NewMLServer returns a new MLServer instance.
|
|
func NewMLServer(cfg *Config) (*MLServer, error) {
|
|
if cfg.LocalMode {
|
|
return &MLServer{SSHClient: network.NewLocalClient(cfg.BasePath)}, nil
|
|
}
|
|
|
|
client, err := network.NewSSHClient(cfg.Host, cfg.User, cfg.SSHKey, cfg.Port, cfg.KnownHosts)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &MLServer{SSHClient: client}, nil
|
|
}
|
|
|
|
// Worker represents an ML task worker.
|
|
type Worker struct {
|
|
id string
|
|
config *Config
|
|
server *MLServer
|
|
queue queue.Backend
|
|
resources *resources.Manager
|
|
running map[string]context.CancelFunc // Store cancellation functions for graceful shutdown
|
|
runningMu sync.RWMutex
|
|
ctx context.Context
|
|
cancel context.CancelFunc
|
|
logger *logging.Logger
|
|
metrics *metrics.Metrics
|
|
metricsSrv *http.Server
|
|
|
|
datasetCache map[string]time.Time
|
|
datasetCacheMu sync.RWMutex
|
|
datasetCacheTTL time.Duration
|
|
|
|
// Graceful shutdown fields
|
|
shutdownCh chan struct{}
|
|
activeTasks sync.Map // map[string]*queue.Task - track active tasks
|
|
gracefulWait sync.WaitGroup
|
|
|
|
podman *container.PodmanManager
|
|
jupyter JupyterManager
|
|
trackingRegistry *tracking.Registry
|
|
envPool *envpool.Pool
|
|
|
|
prewarmMu sync.Mutex
|
|
prewarmTargetID string
|
|
prewarmCancel context.CancelFunc
|
|
prewarmStartedAt time.Time
|
|
}
|
|
|
|
func envInt(name string) (int, bool) {
|
|
v := strings.TrimSpace(os.Getenv(name))
|
|
if v == "" {
|
|
return 0, false
|
|
}
|
|
n, err := strconv.Atoi(v)
|
|
if err != nil {
|
|
return 0, false
|
|
}
|
|
return n, true
|
|
}
|
|
|
|
func parseCPUFromConfig(cfg *Config) int {
|
|
if n, ok := envInt("FETCH_ML_TOTAL_CPU"); ok && n >= 0 {
|
|
return n
|
|
}
|
|
if cfg != nil {
|
|
if cfg.Resources.PodmanCPUs != "" {
|
|
if f, err := strconv.ParseFloat(strings.TrimSpace(cfg.Resources.PodmanCPUs), 64); err == nil {
|
|
if f < 0 {
|
|
return 0
|
|
}
|
|
return int(math.Floor(f))
|
|
}
|
|
}
|
|
}
|
|
return runtime.NumCPU()
|
|
}
|
|
|
|
func parseGPUCountFromConfig(cfg *Config) int {
|
|
factory := &GPUDetectorFactory{}
|
|
detector := factory.CreateDetector(cfg)
|
|
return detector.DetectGPUCount()
|
|
}
|
|
|
|
func (w *Worker) getGPUDetector() GPUDetector {
|
|
factory := &GPUDetectorFactory{}
|
|
return factory.CreateDetector(w.config)
|
|
}
|
|
|
|
func parseGPUSlotsPerGPUFromConfig() int {
|
|
if n, ok := envInt("FETCH_ML_GPU_SLOTS_PER_GPU"); ok && n > 0 {
|
|
return n
|
|
}
|
|
return 1
|
|
}
|
|
|
|
func (w *Worker) setupMetricsExporter() {
|
|
if !w.config.Metrics.Enabled {
|
|
return
|
|
}
|
|
|
|
reg := prometheus.NewRegistry()
|
|
reg.MustRegister(
|
|
collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}),
|
|
collectors.NewGoCollector(),
|
|
)
|
|
|
|
labels := prometheus.Labels{"worker_id": w.id}
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_tasks_processed_total",
|
|
Help: "Total tasks processed successfully by this worker.",
|
|
ConstLabels: labels,
|
|
}, func() float64 {
|
|
return float64(w.metrics.TasksProcessed.Load())
|
|
}))
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_tasks_failed_total",
|
|
Help: "Total tasks failed by this worker.",
|
|
ConstLabels: labels,
|
|
}, func() float64 {
|
|
return float64(w.metrics.TasksFailed.Load())
|
|
}))
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_tasks_active",
|
|
Help: "Number of tasks currently running on this worker.",
|
|
ConstLabels: labels,
|
|
}, func() float64 {
|
|
return float64(w.runningCount())
|
|
}))
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_tasks_queued",
|
|
Help: "Latest observed queue depth from Redis.",
|
|
ConstLabels: labels,
|
|
}, func() float64 {
|
|
return float64(w.metrics.QueuedTasks.Load())
|
|
}))
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_data_transferred_bytes_total",
|
|
Help: "Total bytes transferred while fetching datasets.",
|
|
ConstLabels: labels,
|
|
}, func() float64 {
|
|
return float64(w.metrics.DataTransferred.Load())
|
|
}))
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_data_fetch_time_seconds_total",
|
|
Help: "Total time spent fetching datasets (seconds).",
|
|
ConstLabels: labels,
|
|
}, func() float64 {
|
|
return float64(w.metrics.DataFetchTime.Load()) / float64(time.Second)
|
|
}))
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_execution_time_seconds_total",
|
|
Help: "Total execution time for completed tasks (seconds).",
|
|
ConstLabels: labels,
|
|
}, func() float64 {
|
|
return float64(w.metrics.ExecutionTime.Load()) / float64(time.Second)
|
|
}))
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_prewarm_env_hit_total",
|
|
Help: "Total environment prewarm hits (warmed image already existed).",
|
|
ConstLabels: labels,
|
|
}, func() float64 {
|
|
return float64(w.metrics.PrewarmEnvHit.Load())
|
|
}))
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_prewarm_env_miss_total",
|
|
Help: "Total environment prewarm misses (warmed image did not exist yet).",
|
|
ConstLabels: labels,
|
|
}, func() float64 {
|
|
return float64(w.metrics.PrewarmEnvMiss.Load())
|
|
}))
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_prewarm_env_built_total",
|
|
Help: "Total environment prewarm images built.",
|
|
ConstLabels: labels,
|
|
}, func() float64 {
|
|
return float64(w.metrics.PrewarmEnvBuilt.Load())
|
|
}))
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_prewarm_env_time_seconds_total",
|
|
Help: "Total time spent building prewarm images (seconds).",
|
|
ConstLabels: labels,
|
|
}, func() float64 {
|
|
return float64(w.metrics.PrewarmEnvTime.Load()) / float64(time.Second)
|
|
}))
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_prewarm_snapshot_hit_total",
|
|
Help: "Total prewarmed snapshot hits (snapshots found in .prewarm/).",
|
|
ConstLabels: labels,
|
|
}, func() float64 {
|
|
return float64(w.metrics.PrewarmSnapshotHit.Load())
|
|
}))
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_prewarm_snapshot_miss_total",
|
|
Help: "Total prewarmed snapshot misses (snapshots not found in .prewarm/).",
|
|
ConstLabels: labels,
|
|
}, func() float64 {
|
|
return float64(w.metrics.PrewarmSnapshotMiss.Load())
|
|
}))
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_prewarm_snapshot_built_total",
|
|
Help: "Total snapshots prewarmed into .prewarm/.",
|
|
ConstLabels: labels,
|
|
}, func() float64 {
|
|
return float64(w.metrics.PrewarmSnapshotBuilt.Load())
|
|
}))
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_prewarm_snapshot_time_seconds_total",
|
|
Help: "Total time spent prewarming snapshots (seconds).",
|
|
ConstLabels: labels,
|
|
}, func() float64 {
|
|
return float64(w.metrics.PrewarmSnapshotTime.Load()) / float64(time.Second)
|
|
}))
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_worker_max_concurrency",
|
|
Help: "Configured maximum concurrent tasks for this worker.",
|
|
ConstLabels: labels,
|
|
}, func() float64 {
|
|
return float64(w.config.MaxWorkers)
|
|
}))
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_resources_cpu_total",
|
|
Help: "Total CPU tokens managed by the worker resource manager.",
|
|
ConstLabels: labels,
|
|
}, func() float64 {
|
|
return float64(w.resources.Snapshot().TotalCPU)
|
|
}))
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_resources_cpu_free",
|
|
Help: "Free CPU tokens currently available in the worker resource manager.",
|
|
ConstLabels: labels,
|
|
}, func() float64 {
|
|
return float64(w.resources.Snapshot().FreeCPU)
|
|
}))
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_resources_acquire_total",
|
|
Help: "Total resource acquisition attempts.",
|
|
ConstLabels: labels,
|
|
}, func() float64 {
|
|
return float64(w.resources.Snapshot().AcquireTotal)
|
|
}))
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_resources_acquire_wait_total",
|
|
Help: "Total resource acquisitions that had to wait for resources.",
|
|
ConstLabels: labels,
|
|
}, func() float64 {
|
|
return float64(w.resources.Snapshot().AcquireWaitTotal)
|
|
}))
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_resources_acquire_timeout_total",
|
|
Help: "Total resource acquisition attempts that timed out.",
|
|
ConstLabels: labels,
|
|
}, func() float64 {
|
|
return float64(w.resources.Snapshot().AcquireTimeoutTotal)
|
|
}))
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_resources_acquire_wait_seconds_total",
|
|
Help: "Total seconds spent waiting for resources across all acquisitions.",
|
|
ConstLabels: labels,
|
|
}, func() float64 {
|
|
return w.resources.Snapshot().AcquireWaitSeconds
|
|
}))
|
|
|
|
snap := w.resources.Snapshot()
|
|
for i := range snap.GPUFree {
|
|
gpuLabels := prometheus.Labels{"worker_id": w.id, "gpu_index": strconv.Itoa(i)}
|
|
idx := i
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_resources_gpu_slots_total",
|
|
Help: "Total GPU slots per GPU index.",
|
|
ConstLabels: gpuLabels,
|
|
}, func() float64 {
|
|
return float64(w.resources.Snapshot().SlotsPerGPU)
|
|
}))
|
|
reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{
|
|
Name: "fetchml_resources_gpu_slots_free",
|
|
Help: "Free GPU slots per GPU index.",
|
|
ConstLabels: gpuLabels,
|
|
}, func() float64 {
|
|
s := w.resources.Snapshot()
|
|
if idx < 0 || idx >= len(s.GPUFree) {
|
|
return 0
|
|
}
|
|
return float64(s.GPUFree[idx])
|
|
}))
|
|
}
|
|
|
|
mux := http.NewServeMux()
|
|
mux.Handle("/metrics", promhttp.HandlerFor(reg, promhttp.HandlerOpts{}))
|
|
|
|
srv := &http.Server{
|
|
Addr: w.config.Metrics.ListenAddr,
|
|
Handler: mux,
|
|
ReadHeaderTimeout: 5 * time.Second,
|
|
}
|
|
|
|
w.metricsSrv = srv
|
|
go func() {
|
|
w.logger.Info("metrics exporter listening",
|
|
"addr", w.config.Metrics.ListenAddr,
|
|
"enabled", w.config.Metrics.Enabled)
|
|
if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
|
w.logger.Warn("metrics exporter stopped",
|
|
"error", err)
|
|
}
|
|
}()
|
|
}
|
|
|
|
// NewWorker creates a new worker instance.
|
|
func NewWorker(cfg *Config, _ string) (*Worker, error) {
|
|
srv, err := NewMLServer(cfg)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer func() {
|
|
if err != nil {
|
|
if closeErr := srv.Close(); closeErr != nil {
|
|
log.Printf("Warning: failed to close server connection during error cleanup: %v", closeErr)
|
|
}
|
|
}
|
|
}()
|
|
|
|
backendCfg := queue.BackendConfig{
|
|
Backend: queue.QueueBackend(strings.ToLower(strings.TrimSpace(cfg.Queue.Backend))),
|
|
RedisAddr: cfg.RedisAddr,
|
|
RedisPassword: cfg.RedisPassword,
|
|
RedisDB: cfg.RedisDB,
|
|
SQLitePath: cfg.Queue.SQLitePath,
|
|
MetricsFlushInterval: cfg.MetricsFlushInterval,
|
|
}
|
|
queueClient, err := queue.NewBackend(backendCfg)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer func() {
|
|
if err != nil {
|
|
if closeErr := queueClient.Close(); closeErr != nil {
|
|
log.Printf("Warning: failed to close task queue during error cleanup: %v", closeErr)
|
|
}
|
|
}
|
|
}()
|
|
|
|
// Create data_dir if it doesn't exist (for production without NAS)
|
|
if cfg.DataDir != "" {
|
|
if _, err := srv.Exec(fmt.Sprintf("mkdir -p %s", cfg.DataDir)); err != nil {
|
|
log.Printf("Warning: failed to create data_dir %s: %v", cfg.DataDir, err)
|
|
}
|
|
}
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer func() {
|
|
if err != nil {
|
|
cancel()
|
|
}
|
|
}()
|
|
ctx = logging.EnsureTrace(ctx)
|
|
ctx = logging.CtxWithWorker(ctx, cfg.WorkerID)
|
|
|
|
baseLogger := logging.NewLogger(slog.LevelInfo, false)
|
|
logger := baseLogger.Component(ctx, "worker")
|
|
metricsObj := &metrics.Metrics{}
|
|
|
|
podmanMgr, err := container.NewPodmanManager(logger)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to create podman manager: %w", err)
|
|
}
|
|
|
|
jupyterMgr, err := jupyter.NewServiceManager(logger, jupyter.GetDefaultServiceConfig())
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to create jupyter service manager: %w", err)
|
|
}
|
|
|
|
trackingRegistry := tracking.NewRegistry(logger)
|
|
pluginLoader := factory.NewPluginLoader(logger, podmanMgr)
|
|
|
|
if len(cfg.Plugins) == 0 {
|
|
logger.Warn("no plugins configured, defining defaults")
|
|
// Register defaults manually for backward compatibility/local dev
|
|
mlflowPlugin, err := trackingplugins.NewMLflowPlugin(
|
|
logger,
|
|
podmanMgr,
|
|
trackingplugins.MLflowOptions{
|
|
ArtifactBasePath: filepath.Join(cfg.BasePath, "tracking", "mlflow"),
|
|
},
|
|
)
|
|
if err == nil {
|
|
trackingRegistry.Register(mlflowPlugin)
|
|
}
|
|
|
|
tensorboardPlugin, err := trackingplugins.NewTensorBoardPlugin(
|
|
logger,
|
|
podmanMgr,
|
|
trackingplugins.TensorBoardOptions{
|
|
LogBasePath: filepath.Join(cfg.BasePath, "tracking", "tensorboard"),
|
|
},
|
|
)
|
|
if err == nil {
|
|
trackingRegistry.Register(tensorboardPlugin)
|
|
}
|
|
|
|
trackingRegistry.Register(trackingplugins.NewWandbPlugin())
|
|
} else {
|
|
if err := pluginLoader.LoadPlugins(cfg.Plugins, trackingRegistry); err != nil {
|
|
return nil, fmt.Errorf("failed to load plugins: %w", err)
|
|
}
|
|
}
|
|
|
|
worker := &Worker{
|
|
id: cfg.WorkerID,
|
|
config: cfg,
|
|
server: srv,
|
|
queue: queueClient,
|
|
running: make(map[string]context.CancelFunc),
|
|
datasetCache: make(map[string]time.Time),
|
|
datasetCacheTTL: cfg.DatasetCacheTTL,
|
|
ctx: ctx,
|
|
cancel: cancel,
|
|
logger: logger,
|
|
metrics: metricsObj,
|
|
shutdownCh: make(chan struct{}),
|
|
podman: podmanMgr,
|
|
jupyter: jupyterMgr,
|
|
trackingRegistry: trackingRegistry,
|
|
envPool: envpool.New(""),
|
|
}
|
|
|
|
rm, rmErr := resources.NewManager(resources.Options{
|
|
TotalCPU: parseCPUFromConfig(cfg),
|
|
GPUCount: parseGPUCountFromConfig(cfg),
|
|
SlotsPerGPU: parseGPUSlotsPerGPUFromConfig(),
|
|
})
|
|
if rmErr != nil {
|
|
return nil, fmt.Errorf("failed to init resource manager: %w", rmErr)
|
|
}
|
|
worker.resources = rm
|
|
|
|
if !cfg.LocalMode {
|
|
gpuType := strings.ToLower(strings.TrimSpace(os.Getenv("FETCH_ML_GPU_TYPE")))
|
|
if cfg.AppleGPU.Enabled {
|
|
logger.Warn("apple MPS GPU mode is intended for development; do not use in production",
|
|
"gpu_type", "apple",
|
|
)
|
|
}
|
|
if gpuType == "amd" {
|
|
logger.Warn("amd GPU mode is intended for development; do not use in production",
|
|
"gpu_type", "amd",
|
|
)
|
|
}
|
|
}
|
|
|
|
worker.setupMetricsExporter()
|
|
|
|
// Pre-pull tracking images in background
|
|
go worker.prePullImages()
|
|
|
|
return worker, nil
|
|
}
|
|
|
|
func (w *Worker) prePullImages() {
|
|
if w.config.LocalMode {
|
|
return
|
|
}
|
|
|
|
w.logger.Info("starting image pre-pulling")
|
|
|
|
// Pull worker image
|
|
if w.config.PodmanImage != "" {
|
|
w.pullImage(w.config.PodmanImage)
|
|
}
|
|
|
|
// Pull plugin images
|
|
for name, cfg := range w.config.Plugins {
|
|
if !cfg.Enabled || cfg.Image == "" {
|
|
continue
|
|
}
|
|
w.logger.Info("pre-pulling plugin image", "plugin", name, "image", cfg.Image)
|
|
w.pullImage(cfg.Image)
|
|
}
|
|
}
|
|
|
|
func (w *Worker) pullImage(image string) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
|
|
defer cancel()
|
|
|
|
cmd := exec.CommandContext(ctx, "podman", "pull", image)
|
|
if output, err := cmd.CombinedOutput(); err != nil {
|
|
w.logger.Warn("failed to pull image", "image", image, "error", err, "output", string(output))
|
|
} else {
|
|
w.logger.Info("image pulled successfully", "image", image)
|
|
}
|
|
}
|