SCHED-1898: Speed up Exporter's Slurm metrics collection

docs/slurm-exporter.md

@@ -134,7 +134,10 @@ The exporter provides self-monitoring metrics to track its own health and perfor
 | **slurm_exporter_collection_duration_seconds**<br>*Gauge* | Duration of the most recent metrics collection from SLURM APIs<br><br>**Labels:** None |
 | **slurm_exporter_collection_attempts_total**<br>*Counter* | Total number of metrics collection attempts<br><br>**Labels:** None |
 | **slurm_exporter_collection_failures_total**<br>*Counter* | Total number of failed metrics collection attempts<br><br>**Labels:** None |
+| **slurm_exporter_collector_duration_seconds**<br>*Gauge* | Duration of the most recent sub-collector run during metrics collection<br><br>**Labels:** `collector` (one of `nodes`, `jobs`, `diag`) |
 | **slurm_exporter_collector_errors_total**<br>*Counter* | Total number of errors per sub-collector during metrics collection. Sub-collectors are isolated, so a failure in one does not drop the others' metrics for that cycle; the failed collector's metrics simply gap until it recovers.<br><br>**Labels:** `collector` (one of `nodes`, `jobs`, `diag`) |
+| **slurm_exporter_collector_inflight**<br>*Gauge* | Whether a sub-collector run is currently in progress<br><br>**Labels:** `collector` (one of `nodes`, `jobs`, `diag`) |
+| **slurm_exporter_collector_skipped_total**<br>*Counter* | Total number of skipped sub-collector runs because the previous run for the same collector was still in progress<br><br>**Labels:** `collector` (one of `nodes`, `jobs`, `diag`) |
 | **slurm_exporter_metrics_requests_total**<br>*Counter* | Total number of requests to the `/metrics` endpoint<br><br>**Labels:** None |
 | **slurm_exporter_metrics_exported**<br>*Gauge* | Number of metrics exported in the last scrape<br><br>**Labels:** None |
 

internal/exporter/collector.go

@@ -7,6 +7,7 @@ import (
 	"iter"
 	"maps"
 	"strconv"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -51,7 +52,8 @@ type MetricsCollector struct {
 	controllerServerThreadCount *prometheus.Desc
 
 	// Atomic pointer to the current state for lock-free reads
-	state atomic.Pointer[metricsCollectorState]
+	stateMu sync.Mutex
+	state   atomic.Pointer[metricsCollectorState]
 
 	// Monitoring contains self-monitoring metrics
 	Monitoring *MonitoringMetrics
@@ -333,6 +335,8 @@ func (c *MetricsCollector) Describe(ch chan<- *prometheus.Desc) {
 func (c *MetricsCollector) updateState(ctx context.Context) (err error) {
 	logger := log.FromContext(ctx).WithName(ControllerName)
 	startTime := time.Now()
+	c.stateMu.Lock()
+	defer c.stateMu.Unlock()
 
 	defer func() {
 		duration := time.Since(startTime).Seconds()
@@ -369,11 +373,13 @@ func (c *MetricsCollector) updateState(ctx context.Context) (err error) {
 
 	var errs []error
 	for _, col := range collectors {
+		colStart := time.Now()
 		if cerr := col.run(); cerr != nil {
 			logger.Error(cerr, "Sub-collector failed", "collector", col.name)
 			c.Monitoring.RecordCollectorError(col.name)
 			errs = append(errs, fmt.Errorf("%s: %w", col.name, cerr))
 		}
+		c.Monitoring.RecordCollectorDuration(col.name, time.Since(colStart).Seconds())
 	}
 
 	logger.Info("Collected metrics", "elapsed_seconds", time.Since(startTime).Seconds())
@@ -383,54 +389,144 @@ func (c *MetricsCollector) updateState(ctx context.Context) (err error) {
 
 // collectNodes fetches nodes and updates node-derived metrics on the new state.
 func (c *MetricsCollector) collectNodes(ctx context.Context, previousState, newState *metricsCollectorState) error {
+	nodes, err := c.listNodes(ctx)
+	if err != nil {
+		return err
+	}
+	c.applyNodes(ctx, nodes, previousState, newState)
+	return nil
+}
+
+func (c *MetricsCollector) listNodes(ctx context.Context) ([]slurmapi.Node, error) {
 	logger := log.FromContext(ctx).WithName(ControllerName)
 
 	nodesStart := time.Now()
 	nodes, err := c.slurmAPIClient.ListNodes(ctx)
 	if err != nil {
-		return fmt.Errorf("get nodes from SLURM API: %w", err)
+		return nil, fmt.Errorf("get nodes from SLURM API: %w", err)
 	}
 	logger.Info("Fetched nodes", "count", len(nodes), "elapsed_seconds", time.Since(nodesStart).Seconds())
-	newState.nodes = nodes
 
+	return nodes, nil
+}
+
+func (c *MetricsCollector) applyNodes(ctx context.Context, nodes []slurmapi.Node, previousState, newState *metricsCollectorState) {
+	newState.nodes = nodes
 	c.updateNodeFailureMetrics(nodes, previousState.nodes)
 	c.updateNodeStateMetrics(nodes, previousState, newState, time.Now())
 	newState.lastGPUSecondsUpdate = c.updateGPUSecondsMetrics(ctx, nodes, previousState.lastGPUSecondsUpdate, time.Now())
+}
+
+func cloneMetricsCollectorState(previousState *metricsCollectorState) *metricsCollectorState {
+	if previousState == nil {
+		previousState = newMetricsCollectorState()
+	}
+	newState := newMetricsCollectorState()
+	newState.lastGPUSecondsUpdate = previousState.lastGPUSecondsUpdate
+	newState.nodes = previousState.nodes
+	newState.jobs = previousState.jobs
+	newState.diag = previousState.diag
+	maps.Copy(newState.nodeUnavailabilityStartTimes, previousState.nodeUnavailabilityStartTimes)
+	maps.Copy(newState.nodeDrainingStartTimes, previousState.nodeDrainingStartTimes)
+	return newState
+}
+
+func (c *MetricsCollector) refreshNodes(ctx context.Context) error {
+	nodes, err := c.listNodes(ctx)
+	if err != nil {
+		return err
+	}
+
+	c.stateMu.Lock()
+	defer c.stateMu.Unlock()
 
+	previousState := c.state.Load()
+	if previousState == nil {
+		previousState = newMetricsCollectorState()
+	}
+	newState := cloneMetricsCollectorState(previousState)
+	c.applyNodes(ctx, nodes, previousState, newState)
+	c.state.Store(newState)
 	return nil
 }
 
 // collectJobs fetches jobs and stores them on the new state.
 func (c *MetricsCollector) collectJobs(ctx context.Context, newState *metricsCollectorState) error {
+	jobs, err := c.listJobs(ctx)
+	if err != nil {
+		return err
+	}
+	newState.jobs = jobs
+	return nil
+}
+
+func (c *MetricsCollector) listJobs(ctx context.Context) ([]slurmapi.Job, error) {
 	logger := log.FromContext(ctx).WithName(ControllerName)
 
 	jobsStart := time.Now()
 	jobs, err := c.slurmAPIClient.ListJobsWithParams(ctx, c.jobListParams)
 	if err != nil {
-		return fmt.Errorf("get jobs from SLURM API: %w", err)
+		return nil, fmt.Errorf("get jobs from SLURM API: %w", err)
 	}
 	logger.Info("Fetched jobs",
 		"source", string(c.jobListParams.Source),
 		"count", len(jobs),
 		"elapsed_seconds", time.Since(jobsStart).Seconds(),
 	)
-	newState.jobs = jobs
 
+	return jobs, nil
+}
+
+func (c *MetricsCollector) refreshJobs(ctx context.Context) error {
+	jobs, err := c.listJobs(ctx)
+	if err != nil {
+		return err
+	}
+
+	c.stateMu.Lock()
+	defer c.stateMu.Unlock()
+
+	newState := cloneMetricsCollectorState(c.state.Load())
+	newState.jobs = jobs
+	c.state.Store(newState)
 	return nil
 }
 
 // collectDiag fetches controller diagnostics and stores them on the new state.
 func (c *MetricsCollector) collectDiag(ctx context.Context, newState *metricsCollectorState) error {
+	diag, err := c.getDiag(ctx)
+	if err != nil {
+		return err
+	}
+	newState.diag = diag
+	return nil
+}
+
+func (c *MetricsCollector) getDiag(ctx context.Context) (*api.V0041OpenapiDiagResp, error) {
 	logger := log.FromContext(ctx).WithName(ControllerName)
 
 	diagStart := time.Now()
 	diag, err := c.slurmAPIClient.GetDiag(ctx)
 	if err != nil {
-		return fmt.Errorf("get diag from SLURM API: %w", err)
+		return nil, fmt.Errorf("get diag from SLURM API: %w", err)
 	}
 	logger.Info("Fetched controller diag", "elapsed_seconds", time.Since(diagStart).Seconds())
-	newState.diag = diag
 
+	return diag, nil
+}
+
+func (c *MetricsCollector) refreshDiag(ctx context.Context) error {
+	diag, err := c.getDiag(ctx)
+	if err != nil {
+		return err
+	}
+
+	c.stateMu.Lock()
+	defer c.stateMu.Unlock()
+
+	newState := cloneMetricsCollectorState(c.state.Load())
+	newState.diag = diag
+	c.state.Store(newState)
 	return nil
 }
 

internal/exporter/exporter.go

@@ -5,6 +5,8 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
@@ -51,6 +53,12 @@ type Exporter struct {
 	monitoringServer *http.Server
 }
 
+type asyncSubCollector struct {
+	name    string
+	running atomic.Bool
+	run     func(context.Context) error
+}
+
 // NewClusterExporter creates a new SLURM cluster exporter
 func NewClusterExporter(slurmAPIClient slurmapi.Client, params Params) *Exporter {
 	registry := prometheus.NewRegistry()
@@ -102,13 +110,67 @@ func (e *Exporter) collectionLoop(ctx context.Context) {
 	logger := log.FromContext(ctx).WithName(ControllerName)
 	logger.Info("Starting metrics collection loop", "interval", e.params.CollectionInterval)
 
+	loopCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
 	ticker := time.NewTicker(e.params.CollectionInterval)
 	defer ticker.Stop()
 
-	if err := e.collector.updateState(ctx); err != nil {
-		logger.Error(err, "Initial metrics collection failed")
+	collectors := []*asyncSubCollector{
+		{name: "nodes", run: e.collector.refreshNodes},
+		{name: "jobs", run: e.collector.refreshJobs},
+		{name: "diag", run: e.collector.refreshDiag},
 	}
 
+	startCollectors := func() {
+		start := time.Now()
+		var wg sync.WaitGroup
+		var failed atomic.Bool
+		started := false
+
+		for _, col := range collectors {
+			if !col.running.CompareAndSwap(false, true) {
+				e.monitoringMetrics.RecordCollectorSkipped(col.name)
+				logger.Info("Skipping sub-collector; previous run still in progress", "collector", col.name)
+				continue
+			}
+
+			started = true
+			wg.Add(1)
+			e.monitoringMetrics.SetCollectorInflight(col.name, true)
+			go func(col *asyncSubCollector) {
+				defer wg.Done()
+				defer col.running.Store(false)
+				defer e.monitoringMetrics.SetCollectorInflight(col.name, false)
+
+				start := time.Now()
+				if err := col.run(loopCtx); err != nil {
+					if loopCtx.Err() == nil {
+						logger.Error(err, "Sub-collector failed", "collector", col.name)
+						e.monitoringMetrics.RecordCollectorError(col.name)
+						failed.Store(true)
+					}
+				}
+				e.monitoringMetrics.RecordCollectorDuration(col.name, time.Since(start).Seconds())
+			}(col)
+		}
+
+		if !started {
+			e.monitoringMetrics.RecordCollection(0, nil)
+			return
+		}
+		go func() {
+			wg.Wait()
+			var err error
+			if failed.Load() {
+				err = errors.New("sub-collector failed")
+			}
+			e.monitoringMetrics.RecordCollection(time.Since(start).Seconds(), err)
+		}()
+	}
+
+	startCollectors()
+
 	for {
 		select {
 		case <-ctx.Done():
@@ -118,9 +180,7 @@ func (e *Exporter) collectionLoop(ctx context.Context) {
 			logger.Info("Stopping metrics collection loop via stop channel")
 			return
 		case <-ticker.C:
-			if err := e.collector.updateState(ctx); err != nil {
-				logger.Error(err, "Metrics collection failed")
-			}
+			startCollectors()
 		}
 	}
 }

internal/exporter/exporter_test.go

@@ -0,0 +1,86 @@
+package exporter
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	api "github.com/SlinkyProject/slurm-client/api/v0041"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+	"sigs.k8s.io/controller-runtime/pkg/log"
+	"sigs.k8s.io/controller-runtime/pkg/log/zap"
+
+	"nebius.ai/slurm-operator/internal/slurmapi"
+	"nebius.ai/slurm-operator/internal/slurmapi/fake"
+)
+
+func TestExporterCollectionLoopSkipsRunningSubCollector(t *testing.T) {
+	log.SetLogger(zap.New(zap.UseDevMode(true)))
+
+	mockClient := &fake.MockClient{}
+	var jobsCalls atomic.Int32
+	jobsStarted := make(chan struct{})
+	jobsDone := make(chan struct{})
+	releaseJobs := make(chan struct{})
+	var closeStarted, closeDone sync.Once
+
+	mockClient.EXPECT().ListNodes(mock.Anything).Return([]slurmapi.Node{}, nil).Maybe()
+	mockClient.EXPECT().GetDiag(mock.Anything).Return(&api.V0041OpenapiDiagResp{}, nil).Maybe()
+	mockClient.EXPECT().ListJobsWithParams(mock.Anything, mock.Anything).RunAndReturn(
+		func(ctx context.Context, _ slurmapi.ListJobsParams) ([]slurmapi.Job, error) {
+			jobsCalls.Add(1)
+			closeStarted.Do(func() { close(jobsStarted) })
+			defer closeDone.Do(func() { close(jobsDone) })
+
+			select {
+			case <-ctx.Done():
+				return nil, ctx.Err()
+			case <-releaseJobs:
+				return nil, nil
+			}
+		},
+	)
+
+	exporter := NewClusterExporter(mockClient, Params{CollectionInterval: 5 * time.Millisecond})
+	registry := prometheus.NewRegistry()
+	require.NoError(t, exporter.monitoringMetrics.Register(registry))
+
+	ctx, cancel := context.WithCancel(context.Background())
+	loopDone := make(chan struct{})
+	go func() {
+		exporter.collectionLoop(ctx)
+		close(loopDone)
+	}()
+
+	select {
+	case <-jobsStarted:
+	case <-time.After(time.Second):
+		t.Fatal("jobs collector did not start")
+	}
+
+	require.Eventually(t, func() bool {
+		families, err := registry.Gather()
+		require.NoError(t, err)
+		return collectorCounterValue(families, "slurm_exporter_collector_skipped_total", "jobs") > 0
+	}, time.Second, 10*time.Millisecond)
+	assert.Equal(t, int32(1), jobsCalls.Load())
+
+	cancel()
+	close(releaseJobs)
+
+	select {
+	case <-jobsDone:
+	case <-time.After(time.Second):
+		t.Fatal("jobs collector did not stop")
+	}
+	select {
+	case <-loopDone:
+	case <-time.After(time.Second):
+		t.Fatal("collection loop did not stop")
+	}
+}