diff --git a/autopilot-daemon/pkg/cmd/main.go b/autopilot-daemon/pkg/cmd/main.go index aed09c2..864c482 100644 --- a/autopilot-daemon/pkg/cmd/main.go +++ b/autopilot-daemon/pkg/cmd/main.go @@ -6,11 +6,13 @@ import ( "fmt" "net/http" "os" + "runtime" "time" "github.com/IBM/autopilot/pkg/handler" "github.com/IBM/autopilot/pkg/healthcheck" "github.com/IBM/autopilot/pkg/utils" + "github.com/IBM/autopilot/pkg/worker" "github.com/prometheus/client_golang/prometheus" "github.com/prometheus/client_golang/prometheus/promhttp" "k8s.io/klog/v2" @@ -23,6 +25,7 @@ func main() { v := flag.String("loglevel", "2", "Log level") repeat := flag.String("w", "24h", "Run all tests periodically on each node. Time set in interval format. Defaults to 24h") invasive := flag.String("invasive-check-timer", "4h", "Run invasive checks (e.g., dcgmi level 3) on each node when GPUs are free. Time set in interval format. Defaults to 4h. Set to 0 to avoid invasive checks") + poolSizeInput := flag.Int("workers", 0, "Number of workers to use for concurrent health checks. Defaults to 0 which uses 2*number_of_logical_CPU_cores") flag.Parse() @@ -122,8 +125,19 @@ func main() { // Create a Watcher over nodes. Needed to export metrics from data created by external jobs (i.e., dcgm Jobs) go utils.WatchNode() + // Set the pool size based on the number of CPU cores + poolSize := runtime.NumCPU() * 2 // use 2 workers per CPU core + if *poolSizeInput > 0 { + // if user has set a limit, use it + poolSize = *poolSizeInput + } + + // Create a WorkerPool to handle tasks concurrently + workerPool := worker.CreateWorkerPool(poolSize) + klog.Infof("Starting WorkerPool with %d workers", poolSize) + // Run the health checks at startup, then start the timer - healthcheck.PeriodicCheck() + workerPool.Submit(worker.TaskPeriodicCheck) // Parse the repeat and invasive intervals to durations repeatDuration, err := utils.ParseInterval(*repeat) @@ -144,10 +158,10 @@ func main() { for { select { case <-periodicChecksTicker.C: - healthcheck.PeriodicCheck() + workerPool.Submit(worker.TaskPeriodicCheck) case <-invasiveChecksTicker.C: if invasiveDuration > 0 { - healthcheck.InvasiveCheck() + workerPool.Submit(worker.TaskInvasiveCheck) } } } diff --git a/autopilot-daemon/pkg/worker/enum.go b/autopilot-daemon/pkg/worker/enum.go new file mode 100644 index 0000000..c01ec13 --- /dev/null +++ b/autopilot-daemon/pkg/worker/enum.go @@ -0,0 +1,22 @@ +package worker + +// TaskType represents the type of task to be processed by the worker pool. +// it is either a periodic check or an invasive check. +type TaskType int + +const ( + TaskPeriodicCheck TaskType = iota + 1 + TaskInvasiveCheck +) + +// String returns a string representation of the TaskType. +func (t TaskType) String() string { + switch t { + case TaskPeriodicCheck: + return "Periodic Check" + case TaskInvasiveCheck: + return "Invasive Check" + default: + return "Unknown Task Type" + } +} diff --git a/autopilot-daemon/pkg/worker/pool.go b/autopilot-daemon/pkg/worker/pool.go new file mode 100644 index 0000000..5c129e8 --- /dev/null +++ b/autopilot-daemon/pkg/worker/pool.go @@ -0,0 +1,47 @@ +package worker + +import ( + "sync" + + "k8s.io/klog/v2" +) + +// WorkerPool manages a pool of go-routines that process tasks concurrently. +type WorkerPool struct { + // runningTasks keeps track of tasks currently being processed + runningTasks *sync.Map + // taskQueue is a channel where tasks are submitted for processing + taskQueue chan TaskType +} + +// CreateWorkerPool initializes a new WorkerPool with a specified number of workers. +func CreateWorkerPool(numberOfWorkers int) *WorkerPool { + syncMap := &sync.Map{} + taskQueue := make(chan TaskType) + + // start the specified number of workers + for i := 0; i < numberOfWorkers; i++ { + go worker(taskQueue, syncMap) // start a worker goroutine, see worker.go for implementation + } + + return &WorkerPool{ + runningTasks: syncMap, + taskQueue: taskQueue, + } +} + +// Submit adds a task to the worker pool for processing. +func (wp *WorkerPool) Submit(task TaskType) { + // check if the task is running + if _, exists := wp.runningTasks.Load(task); exists { + klog.InfoS("Task already running, skipping submission", "task", task.String()) + return // task is already running, do not submit again + } + + // mark the task as running, so it won't be submitted again + // the worker will remove it from runningTasks when done + wp.runningTasks.Store(task, struct{}{}) + wp.taskQueue <- task + + klog.InfoS("Task submitted to worker pool", "task", task.String()) +} diff --git a/autopilot-daemon/pkg/worker/worker.go b/autopilot-daemon/pkg/worker/worker.go new file mode 100644 index 0000000..9017511 --- /dev/null +++ b/autopilot-daemon/pkg/worker/worker.go @@ -0,0 +1,31 @@ +package worker + +import ( + "sync" + + "github.com/IBM/autopilot/pkg/healthcheck" + + "k8s.io/klog/v2" +) + +// worker is a function that processes tasks from the task queue. +// it runs in a separate goroutine and listens for tasks to process. +// it removes the task from the runningTasks map once completed. +func worker(c chan TaskType, sm *sync.Map) { + for task := range c { + klog.InfoS("Processing task", "task", task.String()) + + switch task { + case TaskPeriodicCheck: + healthcheck.PeriodicCheck() + case TaskInvasiveCheck: + healthcheck.InvasiveCheck() + default: + klog.Errorf("Unknown task type: %v", task) + } + + // mark the task as completed + sm.Delete(task) + klog.InfoS("Task completed", "task", task.String()) + } +} diff --git a/helm-charts/autopilot/templates/autopilot.yaml b/helm-charts/autopilot/templates/autopilot.yaml index 9b042a4..6f3de7d 100644 --- a/helm-charts/autopilot/templates/autopilot.yaml +++ b/helm-charts/autopilot/templates/autopilot.yaml @@ -51,7 +51,7 @@ spec: - sh - -c - | - /usr/local/bin/autopilot --port {{ .Values.service.port }} --loglevel={{ .Values.loglevel }} --bw {{ .Values.PCIeBW }} --w {{ .Values.repeat }} --invasive-check-timer {{ .Values.invasive }} + /usr/local/bin/autopilot --port {{ .Values.service.port }} --loglevel={{ .Values.loglevel }} --bw {{ .Values.PCIeBW }} --w {{ .Values.repeat }} --invasive-check-timer {{ .Values.invasive }} --workers {{ .Values.workers }} imagePullPolicy: {{ .Values.image.pullPolicy }} name: autopilot securityContext: diff --git a/helm-charts/autopilot/values.yaml b/helm-charts/autopilot/values.yaml index 29e36f3..689aae9 100644 --- a/helm-charts/autopilot/values.yaml +++ b/helm-charts/autopilot/values.yaml @@ -5,6 +5,9 @@ image: repository: quay.io/autopilot/autopilot pullPolicy: Always +# Workers for concurrent tasks. Defaults to 0 which uses 2*number_of_logical_CPU_cores (depends on the resource limits). +workers: 2 + # Bandwidth threshold below which PCIe links are considered defective (Gb/s) # It is recommended to set a threshold that is 25% or lower of the expected peak PCIe bandwidth capability, which maps to maximum peak from 16 lanes to 4 lanes. For example, for a PCIe Gen4x16, reported peak bandwidth is 63GB/s. A degradation at 25% is 15.75GB/s, which corresponds to PCIe Gen4x4. The measured bandwidth is expected to be at least 80% of the expected peak PCIe generation bandwidth. PCIeBW: 4