Add a worker pool for health check go-routines

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)
 			}
 		}
 	}

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"
+	}
+}

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())
+}

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())
+	}
+}

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:

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