4756348c48
Add security hardening features for worker execution: - Worker config with sandboxing options (network_mode, read_only, secrets) - Execution setup with security context propagation - Podman container runtime security enhancements - Security configuration management in config package - Add homelab-sandbox.yaml example configuration Supports running jobs in isolated, restricted environments.
490 lines
15 KiB
Go
490 lines
15 KiB
Go
package worker
|
|
|
|
import (
|
|
"fmt"
|
|
"math"
|
|
"net/url"
|
|
"os"
|
|
"path/filepath"
|
|
"runtime"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/jfraeys/fetch_ml/internal/auth"
|
|
"github.com/jfraeys/fetch_ml/internal/config"
|
|
"github.com/jfraeys/fetch_ml/internal/fileutil"
|
|
"github.com/jfraeys/fetch_ml/internal/queue"
|
|
"github.com/jfraeys/fetch_ml/internal/storage"
|
|
"github.com/jfraeys/fetch_ml/internal/tracking/factory"
|
|
"gopkg.in/yaml.v3"
|
|
)
|
|
|
|
const (
|
|
defaultMetricsFlushInterval = 500 * time.Millisecond
|
|
datasetCacheDefaultTTL = 30 * time.Minute
|
|
)
|
|
|
|
type QueueConfig struct {
|
|
Backend string `yaml:"backend"`
|
|
SQLitePath string `yaml:"sqlite_path"`
|
|
FilesystemPath string `yaml:"filesystem_path"`
|
|
FallbackToFilesystem bool `yaml:"fallback_to_filesystem"`
|
|
}
|
|
|
|
// Config holds worker configuration.
|
|
type Config struct {
|
|
Host string `yaml:"host"`
|
|
User string `yaml:"user"`
|
|
SSHKey string `yaml:"ssh_key"`
|
|
Port int `yaml:"port"`
|
|
BasePath string `yaml:"base_path"`
|
|
TrainScript string `yaml:"train_script"`
|
|
RedisURL string `yaml:"redis_url"`
|
|
RedisAddr string `yaml:"redis_addr"`
|
|
RedisPassword string `yaml:"redis_password"`
|
|
RedisDB int `yaml:"redis_db"`
|
|
Queue QueueConfig `yaml:"queue"`
|
|
KnownHosts string `yaml:"known_hosts"`
|
|
WorkerID string `yaml:"worker_id"`
|
|
MaxWorkers int `yaml:"max_workers"`
|
|
PollInterval int `yaml:"poll_interval_seconds"`
|
|
Resources config.ResourceConfig `yaml:"resources"`
|
|
LocalMode bool `yaml:"local_mode"`
|
|
|
|
// Authentication
|
|
Auth auth.Config `yaml:"auth"`
|
|
|
|
// Metrics exporter
|
|
Metrics MetricsConfig `yaml:"metrics"`
|
|
// Metrics buffering
|
|
MetricsFlushInterval time.Duration `yaml:"metrics_flush_interval"`
|
|
|
|
// Data management
|
|
DataManagerPath string `yaml:"data_manager_path"`
|
|
AutoFetchData bool `yaml:"auto_fetch_data"`
|
|
DataDir string `yaml:"data_dir"`
|
|
DatasetCacheTTL time.Duration `yaml:"dataset_cache_ttl"`
|
|
|
|
SnapshotStore SnapshotStoreConfig `yaml:"snapshot_store"`
|
|
|
|
// Provenance enforcement
|
|
// Default: fail-closed (trustworthiness-by-default). Set true to opt into best-effort.
|
|
ProvenanceBestEffort bool `yaml:"provenance_best_effort"`
|
|
|
|
// Phase 1: opt-in prewarming of next task artifacts (snapshot/datasets/env).
|
|
PrewarmEnabled bool `yaml:"prewarm_enabled"`
|
|
|
|
// Podman execution
|
|
PodmanImage string `yaml:"podman_image"`
|
|
ContainerWorkspace string `yaml:"container_workspace"`
|
|
ContainerResults string `yaml:"container_results"`
|
|
GPUDevices []string `yaml:"gpu_devices"`
|
|
GPUVendor string `yaml:"gpu_vendor"`
|
|
GPUVisibleDevices []int `yaml:"gpu_visible_devices"`
|
|
GPUVisibleDeviceIDs []string `yaml:"gpu_visible_device_ids"`
|
|
|
|
// Apple M-series GPU configuration
|
|
AppleGPU AppleGPUConfig `yaml:"apple_gpu"`
|
|
|
|
// Task lease and retry settings
|
|
TaskLeaseDuration time.Duration `yaml:"task_lease_duration"` // Worker lease (default: 30min)
|
|
HeartbeatInterval time.Duration `yaml:"heartbeat_interval"` // Renew lease (default: 1min)
|
|
MaxRetries int `yaml:"max_retries"` // Maximum retry attempts (default: 3)
|
|
GracefulTimeout time.Duration `yaml:"graceful_timeout"` // Shutdown timeout (default: 5min)
|
|
|
|
// Plugins configuration
|
|
Plugins map[string]factory.PluginConfig `yaml:"plugins"`
|
|
|
|
// Sandboxing configuration
|
|
Sandbox SandboxConfig `yaml:"sandbox"`
|
|
}
|
|
|
|
// MetricsConfig controls the Prometheus exporter.
|
|
type MetricsConfig struct {
|
|
Enabled bool `yaml:"enabled"`
|
|
ListenAddr string `yaml:"listen_addr"`
|
|
}
|
|
|
|
type SnapshotStoreConfig struct {
|
|
Enabled bool `yaml:"enabled"`
|
|
Endpoint string `yaml:"endpoint"`
|
|
Secure bool `yaml:"secure"`
|
|
Region string `yaml:"region"`
|
|
Bucket string `yaml:"bucket"`
|
|
Prefix string `yaml:"prefix"`
|
|
AccessKey string `yaml:"access_key"`
|
|
SecretKey string `yaml:"secret_key"`
|
|
SessionToken string `yaml:"session_token"`
|
|
Timeout time.Duration `yaml:"timeout"`
|
|
MaxRetries int `yaml:"max_retries"`
|
|
}
|
|
|
|
// AppleGPUConfig holds configuration for Apple M-series GPU support
|
|
type AppleGPUConfig struct {
|
|
Enabled bool `yaml:"enabled"`
|
|
MetalDevice string `yaml:"metal_device"`
|
|
MPSRuntime string `yaml:"mps_runtime"`
|
|
}
|
|
|
|
// SandboxConfig holds container sandbox settings
|
|
type SandboxConfig struct {
|
|
NetworkMode string `yaml:"network_mode"` // "none", "slirp4netns", "bridge"
|
|
ReadOnlyRoot bool `yaml:"read_only_root"`
|
|
AllowSecrets bool `yaml:"allow_secrets"`
|
|
AllowedSecrets []string `yaml:"allowed_secrets"` // e.g., ["HF_TOKEN", "WANDB_API_KEY"]
|
|
SeccompProfile string `yaml:"seccomp_profile"`
|
|
MaxRuntimeHours int `yaml:"max_runtime_hours"`
|
|
}
|
|
|
|
// Validate checks sandbox configuration
|
|
func (s *SandboxConfig) Validate() error {
|
|
validNetworks := map[string]bool{"none": true, "slirp4netns": true, "bridge": true, "": true}
|
|
if !validNetworks[s.NetworkMode] {
|
|
return fmt.Errorf("invalid network_mode: %s", s.NetworkMode)
|
|
}
|
|
if s.MaxRuntimeHours < 0 {
|
|
return fmt.Errorf("max_runtime_hours must be positive")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// LoadConfig loads worker configuration from a YAML file.
|
|
func LoadConfig(path string) (*Config, error) {
|
|
data, err := fileutil.SecureFileRead(path)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var cfg Config
|
|
if err := yaml.Unmarshal(data, &cfg); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if strings.TrimSpace(cfg.RedisURL) != "" {
|
|
cfg.RedisURL = os.ExpandEnv(strings.TrimSpace(cfg.RedisURL))
|
|
cfg.RedisAddr = cfg.RedisURL
|
|
cfg.RedisPassword = ""
|
|
cfg.RedisDB = 0
|
|
}
|
|
|
|
// Get smart defaults for current environment
|
|
smart := config.GetSmartDefaults()
|
|
|
|
// Use PathRegistry for consistent path management
|
|
paths := config.FromEnv()
|
|
|
|
if cfg.Port == 0 {
|
|
cfg.Port = config.DefaultSSHPort
|
|
}
|
|
if cfg.Host == "" {
|
|
host, err := smart.Host()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to get default host: %w", err)
|
|
}
|
|
cfg.Host = host
|
|
}
|
|
if cfg.BasePath == "" {
|
|
// Prefer PathRegistry over smart defaults for consistency
|
|
cfg.BasePath = paths.ExperimentsDir()
|
|
}
|
|
if cfg.RedisAddr == "" {
|
|
redisAddr, err := smart.RedisAddr()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to get default redis address: %w", err)
|
|
}
|
|
cfg.RedisAddr = redisAddr
|
|
}
|
|
if cfg.KnownHosts == "" {
|
|
knownHosts, err := smart.KnownHostsPath()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to get default known hosts path: %w", err)
|
|
}
|
|
cfg.KnownHosts = knownHosts
|
|
}
|
|
if cfg.WorkerID == "" {
|
|
cfg.WorkerID = fmt.Sprintf("worker-%s", uuid.New().String()[:8])
|
|
}
|
|
cfg.Resources.ApplyDefaults()
|
|
if cfg.MaxWorkers > 0 {
|
|
cfg.Resources.MaxWorkers = cfg.MaxWorkers
|
|
} else {
|
|
maxWorkers, err := smart.MaxWorkers()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to get default max workers: %w", err)
|
|
}
|
|
cfg.MaxWorkers = maxWorkers
|
|
cfg.Resources.MaxWorkers = maxWorkers
|
|
}
|
|
if cfg.PollInterval == 0 {
|
|
pollInterval, err := smart.PollInterval()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to get default poll interval: %w", err)
|
|
}
|
|
cfg.PollInterval = pollInterval
|
|
}
|
|
if cfg.DataManagerPath == "" {
|
|
cfg.DataManagerPath = "./data_manager"
|
|
}
|
|
if cfg.DataDir == "" {
|
|
// Use PathRegistry for consistent data directory
|
|
cfg.DataDir = paths.DataDir()
|
|
}
|
|
if cfg.SnapshotStore.Timeout == 0 {
|
|
cfg.SnapshotStore.Timeout = 10 * time.Minute
|
|
}
|
|
if cfg.SnapshotStore.MaxRetries == 0 {
|
|
cfg.SnapshotStore.MaxRetries = 3
|
|
}
|
|
if cfg.Metrics.ListenAddr == "" {
|
|
cfg.Metrics.ListenAddr = ":9100"
|
|
}
|
|
if cfg.MetricsFlushInterval == 0 {
|
|
cfg.MetricsFlushInterval = defaultMetricsFlushInterval
|
|
}
|
|
if cfg.DatasetCacheTTL == 0 {
|
|
cfg.DatasetCacheTTL = datasetCacheDefaultTTL
|
|
}
|
|
|
|
if strings.TrimSpace(cfg.Queue.Backend) == "" {
|
|
cfg.Queue.Backend = string(queue.QueueBackendRedis)
|
|
}
|
|
if strings.EqualFold(strings.TrimSpace(cfg.Queue.Backend), string(queue.QueueBackendSQLite)) {
|
|
if strings.TrimSpace(cfg.Queue.SQLitePath) == "" {
|
|
cfg.Queue.SQLitePath = filepath.Join(cfg.DataDir, "queue.db")
|
|
}
|
|
cfg.Queue.SQLitePath = storage.ExpandPath(cfg.Queue.SQLitePath)
|
|
}
|
|
if strings.EqualFold(strings.TrimSpace(cfg.Queue.Backend), string(queue.QueueBackendFS)) || cfg.Queue.FallbackToFilesystem {
|
|
if strings.TrimSpace(cfg.Queue.FilesystemPath) == "" {
|
|
cfg.Queue.FilesystemPath = filepath.Join(cfg.DataDir, "queue-fs")
|
|
}
|
|
cfg.Queue.FilesystemPath = storage.ExpandPath(cfg.Queue.FilesystemPath)
|
|
}
|
|
|
|
if strings.TrimSpace(cfg.GPUVendor) == "" {
|
|
if cfg.AppleGPU.Enabled {
|
|
cfg.GPUVendor = string(GPUTypeApple)
|
|
} else if len(cfg.GPUDevices) > 0 ||
|
|
len(cfg.GPUVisibleDevices) > 0 ||
|
|
len(cfg.GPUVisibleDeviceIDs) > 0 {
|
|
cfg.GPUVendor = string(GPUTypeNVIDIA)
|
|
} else {
|
|
cfg.GPUVendor = string(GPUTypeNone)
|
|
}
|
|
}
|
|
|
|
// Set lease and retry defaults
|
|
if cfg.TaskLeaseDuration == 0 {
|
|
cfg.TaskLeaseDuration = 30 * time.Minute
|
|
}
|
|
if cfg.HeartbeatInterval == 0 {
|
|
cfg.HeartbeatInterval = 1 * time.Minute
|
|
}
|
|
if cfg.MaxRetries == 0 {
|
|
cfg.MaxRetries = 3
|
|
}
|
|
if cfg.GracefulTimeout == 0 {
|
|
cfg.GracefulTimeout = 5 * time.Minute
|
|
}
|
|
|
|
return &cfg, nil
|
|
}
|
|
|
|
// Validate implements config.Validator interface.
|
|
func (c *Config) Validate() error {
|
|
if c.Port != 0 {
|
|
if err := config.ValidatePort(c.Port); err != nil {
|
|
return fmt.Errorf("invalid SSH port: %w", err)
|
|
}
|
|
}
|
|
|
|
if c.BasePath != "" {
|
|
// Convert relative paths to absolute
|
|
c.BasePath = storage.ExpandPath(c.BasePath)
|
|
if !filepath.IsAbs(c.BasePath) {
|
|
// Resolve relative to current working directory, not DefaultBasePath
|
|
cwd, err := os.Getwd()
|
|
if err != nil {
|
|
return fmt.Errorf("failed to get current directory: %w", err)
|
|
}
|
|
c.BasePath = filepath.Join(cwd, c.BasePath)
|
|
}
|
|
}
|
|
|
|
backend := strings.ToLower(strings.TrimSpace(c.Queue.Backend))
|
|
if backend == "" {
|
|
backend = string(queue.QueueBackendRedis)
|
|
c.Queue.Backend = backend
|
|
}
|
|
if backend != string(queue.QueueBackendRedis) && backend != string(queue.QueueBackendSQLite) && backend != string(queue.QueueBackendFS) {
|
|
return fmt.Errorf("queue.backend must be one of %q, %q, or %q", queue.QueueBackendRedis, queue.QueueBackendSQLite, queue.QueueBackendFS)
|
|
}
|
|
|
|
if backend == string(queue.QueueBackendSQLite) {
|
|
if strings.TrimSpace(c.Queue.SQLitePath) == "" {
|
|
return fmt.Errorf("queue.sqlite_path is required when queue.backend is %q", queue.QueueBackendSQLite)
|
|
}
|
|
c.Queue.SQLitePath = storage.ExpandPath(c.Queue.SQLitePath)
|
|
if !filepath.IsAbs(c.Queue.SQLitePath) {
|
|
c.Queue.SQLitePath = filepath.Join(config.DefaultLocalDataDir, c.Queue.SQLitePath)
|
|
}
|
|
}
|
|
if backend == string(queue.QueueBackendFS) || c.Queue.FallbackToFilesystem {
|
|
if strings.TrimSpace(c.Queue.FilesystemPath) == "" {
|
|
return fmt.Errorf("queue.filesystem_path is required when filesystem queue is enabled")
|
|
}
|
|
c.Queue.FilesystemPath = storage.ExpandPath(c.Queue.FilesystemPath)
|
|
if !filepath.IsAbs(c.Queue.FilesystemPath) {
|
|
c.Queue.FilesystemPath = filepath.Join(config.DefaultLocalDataDir, c.Queue.FilesystemPath)
|
|
}
|
|
}
|
|
|
|
if c.RedisAddr != "" {
|
|
addr := strings.TrimSpace(c.RedisAddr)
|
|
if strings.HasPrefix(addr, "redis://") {
|
|
u, err := url.Parse(addr)
|
|
if err != nil {
|
|
return fmt.Errorf("invalid Redis configuration: invalid redis url: %w", err)
|
|
}
|
|
if u.Scheme != "redis" || strings.TrimSpace(u.Host) == "" {
|
|
return fmt.Errorf("invalid Redis configuration: invalid redis url")
|
|
}
|
|
} else {
|
|
if err := config.ValidateRedisAddr(addr); err != nil {
|
|
return fmt.Errorf("invalid Redis configuration: %w", err)
|
|
}
|
|
}
|
|
}
|
|
|
|
if c.MaxWorkers < 1 {
|
|
return fmt.Errorf("max_workers must be at least 1, got %d", c.MaxWorkers)
|
|
}
|
|
|
|
switch strings.ToLower(strings.TrimSpace(c.GPUVendor)) {
|
|
case string(GPUTypeNVIDIA), string(GPUTypeApple), string(GPUTypeNone), "amd":
|
|
// ok
|
|
default:
|
|
return fmt.Errorf(
|
|
"gpu_vendor must be one of %q, %q, %q, %q",
|
|
string(GPUTypeNVIDIA),
|
|
"amd",
|
|
string(GPUTypeApple),
|
|
string(GPUTypeNone),
|
|
)
|
|
}
|
|
|
|
// Strict GPU visibility configuration:
|
|
// - gpu_visible_devices and gpu_visible_device_ids are mutually exclusive.
|
|
// - UUID-style gpu_visible_device_ids is NVIDIA-only.
|
|
vendor := strings.ToLower(strings.TrimSpace(c.GPUVendor))
|
|
if len(c.GPUVisibleDevices) > 0 && len(c.GPUVisibleDeviceIDs) > 0 {
|
|
return fmt.Errorf("gpu_visible_devices and gpu_visible_device_ids are mutually exclusive")
|
|
}
|
|
if len(c.GPUVisibleDeviceIDs) > 0 {
|
|
if vendor != string(GPUTypeNVIDIA) {
|
|
return fmt.Errorf(
|
|
"gpu_visible_device_ids is only supported when gpu_vendor is %q",
|
|
string(GPUTypeNVIDIA),
|
|
)
|
|
}
|
|
for _, id := range c.GPUVisibleDeviceIDs {
|
|
id = strings.TrimSpace(id)
|
|
if id == "" {
|
|
return fmt.Errorf("gpu_visible_device_ids contains an empty value")
|
|
}
|
|
if !strings.HasPrefix(id, "GPU-") {
|
|
return fmt.Errorf("gpu_visible_device_ids values must start with %q, got %q", "GPU-", id)
|
|
}
|
|
}
|
|
}
|
|
if vendor == string(GPUTypeApple) || vendor == string(GPUTypeNone) {
|
|
if len(c.GPUVisibleDevices) > 0 || len(c.GPUVisibleDeviceIDs) > 0 {
|
|
return fmt.Errorf(
|
|
"gpu_visible_devices and gpu_visible_device_ids are not supported when gpu_vendor is %q",
|
|
vendor,
|
|
)
|
|
}
|
|
}
|
|
if vendor == "amd" {
|
|
if len(c.GPUVisibleDeviceIDs) > 0 {
|
|
return fmt.Errorf("gpu_visible_device_ids is not supported when gpu_vendor is %q", vendor)
|
|
}
|
|
for _, idx := range c.GPUVisibleDevices {
|
|
if idx < 0 {
|
|
return fmt.Errorf("gpu_visible_devices contains negative index %d", idx)
|
|
}
|
|
}
|
|
}
|
|
|
|
if c.SnapshotStore.Enabled {
|
|
if strings.TrimSpace(c.SnapshotStore.Endpoint) == "" {
|
|
return fmt.Errorf("snapshot_store.endpoint is required when snapshot_store.enabled is true")
|
|
}
|
|
if strings.TrimSpace(c.SnapshotStore.Bucket) == "" {
|
|
return fmt.Errorf("snapshot_store.bucket is required when snapshot_store.enabled is true")
|
|
}
|
|
ak := strings.TrimSpace(c.SnapshotStore.AccessKey)
|
|
sk := strings.TrimSpace(c.SnapshotStore.SecretKey)
|
|
if (ak == "") != (sk == "") {
|
|
return fmt.Errorf(
|
|
"snapshot_store.access_key and snapshot_store.secret_key must both be set or both be empty",
|
|
)
|
|
}
|
|
if c.SnapshotStore.Timeout < 0 {
|
|
return fmt.Errorf("snapshot_store.timeout must be >= 0")
|
|
}
|
|
if c.SnapshotStore.MaxRetries < 0 {
|
|
return fmt.Errorf("snapshot_store.max_retries must be >= 0")
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// envInt reads an integer from environment variable
|
|
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
|
|
}
|
|
|
|
// parseCPUFromConfig determines total CPU from environment or config
|
|
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()
|
|
}
|
|
|
|
// parseGPUCountFromConfig detects GPU count from config
|
|
func parseGPUCountFromConfig(cfg *Config) int {
|
|
factory := &GPUDetectorFactory{}
|
|
detector := factory.CreateDetector(cfg)
|
|
return detector.DetectGPUCount()
|
|
}
|
|
|
|
// parseGPUSlotsPerGPUFromConfig reads GPU slots per GPU from environment
|
|
func parseGPUSlotsPerGPUFromConfig() int {
|
|
if n, ok := envInt("FETCH_ML_GPU_SLOTS_PER_GPU"); ok && n > 0 {
|
|
return n
|
|
}
|
|
return 1
|
|
}
|