Documentation
¶
Overview ¶
Package scheduler implements the sqi-server assignment loop and worker registry — the authoritative component for deciding which task runs on which worker.
The Scheduler runs three concurrent loops:
Lease subscriber: a core-NATS request/reply subscriber on work.lease.> ([handleLeaseRequest]). Idle workers ask for work; the scheduler selects a priority-ordered batch of ready tasks the worker is eligible for that fits its free CPU cores ([selectLeaseBatch]), atomically leases each (store.TaskStore.LeaseReadyTask), and replies with the assignment payloads. When no work is available the request parks in the waiter registry until new work appears or the hold elapses, then replies.
Worker registry: a NATS push-consumer for worker.register messages that persists capability data via store.WorkerStore.RegisterWorker and keeps the WorkersTotal Prometheus gauge current.
Heartbeat sweep: a NATS push-consumer that updates each worker's LastHeartbeatAt on worker.heartbeat messages, paired with a periodic timer that marks workers offline once their heartbeat goes stale (store.WorkerStore.ListStaleWorkers), terminates their open attempts (store.TaskAttemptStore.TerminateWorkerAttempts), and returns their in-flight tasks to the ready queue (store.TaskStore.ReclaimWorkerTasks). The same tick refreshes the queue-depth, idle-worker, and usage-claim Prometheus gauges.
Worker selection. A task is matched to a worker by capability tags, compute-location affinity, and queue/farm filtering (WorkerEligible), subject to per-queue and per-farm maximum-concurrent-task limits ([policyGate]). Once a worker is chosen, a provisional store.TaskAttempt is created and any required usage pool slots are claimed atomically (store.UsageClaimStore.TryClaimSlots); if the pool is saturated the assignment is rolled back and the task stays ready for the next tick. Attempt numbers come from store.TaskAttemptStore.LatestTaskAttempt — 1 for a fresh task, N+1 on retry.
Assignment payload. [buildAssignPayload] re-parses the job's raw OpenJD template to extract the matching step's OnRun action, embedded files, and ordered environments (job environments first, then step, per the OpenJD spec). The path-map field is reserved but empty until named storage location CRUD and resolved-mode path translation are implemented.
Status and log ingestion. A push-consumer on the SQI_TASK stream ([handleTaskStatusMessage]) decodes protocol.TaskStatusMsg from task.status.<job>, updates the task/attempt, releases held usage pool slots, and drives step/job completion including openjd.ResolveDependencies for multi-step jobs. A push-consumer on SQI_LOGS ([handleLogChunk]) persists each task.logs.<task> chunk as a store.TaskLog row, recording both the worker-assigned sequence number and the NATS stream sequence that serves as the log-tail pagination cursor.
Cancellation. [CancelJob] and [CancelTask] are the server-side entry points called by the REST layer: they close running attempts, transition tasks to store.TaskStatusCanceled, publish task.cancel.<taskID> signals to assigned workers (bus.Client.PublishTaskCancel), and release held usage pool slots. The logic lives in cancellation.go; the SQI_CANCEL stream and publish helper live in the bus package.
Wire protocol and metrics. All worker messages use the versioned JSON types in worker/protocol (protocol.RegisterMsg, protocol.HeartbeatMsg, protocol.AssignMsg, protocol.TaskStatusMsg, protocol.LogChunkMsg). Each heartbeat-sweep tick refreshes the scheduler's Prometheus metrics (metrics.Metrics): queue depth by queue, idle workers by farm, and active usage-pool claims per pool.
Index ¶
- func WorkerEligible(worker store.Worker, job store.Job, step store.Step, ...) bool
- func WorkerEligibleWithReason(worker store.Worker, job store.Job, step store.Step, ...) (reason string, eligible bool)
- type Config
- type HeartbeatMsg
- type RegisterMsg
- type Scheduler
- func (s *Scheduler) CancelJob(ctx context.Context, jobID string) error
- func (s *Scheduler) CancelTask(ctx context.Context, taskID string) error
- func (s *Scheduler) ReleaseTaskUsage(ctx context.Context, attemptID string) error
- func (s *Scheduler) RetryJob(ctx context.Context, jobID string) (int, error)
- func (s *Scheduler) RetryTask(ctx context.Context, taskID string) error
- func (s *Scheduler) Run(ctx context.Context) error
- func (s *Scheduler) Stop()
- func (s *Scheduler) WakeQueue(queueID string)
- func (s *Scheduler) WorkerTimeout() time.Duration
Constants ¶
This section is empty.
Variables ¶
This section is empty.
Functions ¶
func WorkerEligible ¶
func WorkerEligible( worker store.Worker, job store.Job, step store.Step, pools map[string]store.UsagePool, activeCounts map[string]int, ) bool
WorkerEligible reports whether worker w can run a task that belongs to step s and job j, given the current usage pool state.
pools maps pool name → store.UsagePool. activeCounts maps pool name → current active claim count.
Both maps may be nil or empty when no usage pools exist.
func WorkerEligibleWithReason ¶
func WorkerEligibleWithReason( worker store.Worker, job store.Job, step store.Step, pools map[string]store.UsagePool, activeCounts map[string]int, ) (reason string, eligible bool)
WorkerEligibleWithReason is the same as WorkerEligible but also returns a human-readable rejection reason for diagnostic logging.
Types ¶
type Config ¶
type Config struct {
// FarmID is the farm this scheduler instance manages. The assignment loop
// only considers tasks that belong to this farm.
FarmID string
// AssignInterval is how often the assignment loop polls the store for ready
// tasks. Lower values reduce latency; higher values reduce DB pressure.
// Default: 1 s.
AssignInterval time.Duration
// AssignBatchSize is the maximum number of ready tasks fetched per
// assignment loop tick. Limits burst DB load without capping throughput
// over time. Default: 50.
AssignBatchSize int
// AssignWorkers is the size of the goroutine pool that processes assignment
// decisions concurrently. Default: 4.
AssignWorkers int
// WorkerTimeout is the maximum time since a worker's last heartbeat before
// the heartbeat sweep considers it offline. Default: 30 s.
WorkerTimeout time.Duration
// HeartbeatSweepInterval is how often the heartbeat sweep runs. Should be
// well below WorkerTimeout so stale workers are caught promptly. Default: 15 s.
HeartbeatSweepInterval time.Duration
// AssignedTaskTimeout is the maximum time a task may sit in 'assigned'
// without transitioning to 'running' before the reaper returns it to the
// ready queue. It guards the brief leased→running window: a task a worker
// leased but never reported running (e.g. the worker crashed between lease
// and process start). The heartbeat sweep cannot recover such tasks while the
// worker is still heartbeating, so this reaper runs independently on the
// heartbeat sweep tick. Default: 30 s.
AssignedTaskTimeout time.Duration
// OfflineWorkerRetention is how long a worker may remain in
// [store.WorkerStatusOffline] before the retention sweep hard-deletes it,
// bounding the growth of the worker table on farms with ephemeral nodes.
// Disabled and online workers are never auto-removed. A value <= 0 disables
// the sweep entirely. Default: 24 h.
OfflineWorkerRetention time.Duration
// JobRetention is how long a terminal job is kept before the retention
// sweep hard-deletes it and all of its data (steps, tasks, attempts, logs).
// completed and canceled jobs are always eligible; failed jobs only when
// JobRetentionIncludeFailed is set. A value <= 0 disables the sweep.
// Default: 168 h (7 days).
JobRetention time.Duration
// JobRetentionIncludeFailed extends the retention sweep to failed jobs.
// Default: false (failed jobs are kept for debugging).
JobRetentionIncludeFailed bool
}
Config holds tuning parameters for the Scheduler.
func DefaultConfig ¶
func DefaultConfig() Config
DefaultConfig returns a Config with conservative production-safe defaults.
type HeartbeatMsg ¶
HeartbeatMsg is the JSON payload workers publish to worker.heartbeat.
type RegisterMsg ¶
type RegisterMsg struct {
// WorkerID is the stable unique identifier for this worker instance.
// Workers MUST use the same ID across restarts so that re-registration
// updates the existing record rather than creating a duplicate.
WorkerID string `json:"worker_id"`
// FarmID is the farm this worker belongs to. Required.
FarmID string `json:"farm_id"`
// QueueID restricts the worker to a single queue if non-empty.
// An empty value means the worker accepts tasks from any queue in FarmID.
QueueID string `json:"queue_id,omitempty"`
// Name is the worker's human-readable display label. Persisted so the UI
// can distinguish multiple workers running on a single host.
Name string `json:"name,omitempty"`
// Hostname and IPAddress are the worker's network identity.
Hostname string `json:"hostname"`
IPAddress string `json:"ip_address,omitempty"`
// ComputeLocation is the named compute location this worker belongs to
// (e.g. "onprem_linux", "cloud_aws_us_east"). Used for path translation
// and task affinity.
ComputeLocation string `json:"compute_location,omitempty"`
// OS and OSVersion are the worker's operating system identity, used for
// capability matching.
OS string `json:"os"`
OSVersion string `json:"os_version,omitempty"`
// WorkerVersion is the sqi-worker build version the worker self-reports,
// distinct from the protocol Version field above.
WorkerVersion string `json:"worker_version,omitempty"`
// CPUCount and RAMMb are the worker's hardware capacity, used for
// resource-aware scheduling in future phases.
CPUCount int `json:"cpu_count,omitempty"`
RAMMb int `json:"ram_mb,omitempty"`
// GPUInfo describes any GPU(s) installed on the worker host.
// omitempty is omitted intentionally: it has no effect on struct fields.
GPUInfo store.GPUInfo `json:"gpu_info"`
// Tags holds arbitrary key/value capability tags the worker self-reports.
// Job requirements reference these tags when specifying worker constraints.
Tags map[string]string `json:"tags,omitempty"`
}
RegisterMsg is the JSON payload workers publish to worker.register. It carries the worker's self-reported identity and capability data. A later protocol revision may add formal versioning; this struct matches the minimal information the server needs for task matching.
type Scheduler ¶
type Scheduler struct {
// contains filtered or unexported fields
}
Scheduler owns the assignment loop, worker registry, and heartbeat sweep. Create it with New and drive it with [Run].
func New ¶
func New(cfg Config, st store.Store, busClient busClient, m *metrics.Metrics, logger *slog.Logger, notifier ws.Notifier, diagBuf *diag.Buffer) *Scheduler
New creates a Scheduler. Call [Run] to start its goroutines.
notifier receives live-event notifications after each state change so the WebSocket hub can fan them out to subscribed clients. Pass ws.NoopNotifier (or nil — treated as NoopNotifier) when no WebSocket hub is wired. diagBuf, when non-nil, enables worker diagnostic-log ingestion: the scheduler subscribes to worker.diag.> and appends records under "worker:<id>". Pass nil to disable diagnostics.
func (*Scheduler) CancelJob ¶
CancelJob transitions all non-terminal tasks for the given job to canceled, closes any running task attempts, releases held usage pool slots, and publishes a task.cancel.<taskID> signal to each worker that was actively executing a task at the time of cancellation.
CancelJob does NOT update the job's own status; that is the caller's responsibility (typically the REST handler that also calls store.JobStore.UpdateJobStatus).
The method is idempotent: if all tasks are already in terminal states the store operations are no-ops and no NATS messages are published.
func (*Scheduler) CancelTask ¶
CancelTask transitions a single task to store.TaskStatusCanceled, closes its running attempt if any, releases held usage pool slots, and publishes a task.cancel.<taskID> signal to the assigned worker.
If the task is already in a terminal state (succeeded, failed, canceled), CancelTask returns nil without modifying any state.
func (*Scheduler) ReleaseTaskUsage ¶
ReleaseTaskUsage releases all active usage-pool claims for the given task attempt. It is called when a task attempt transitions to a terminal state (succeeded, failed, or canceled), freeing the usage pool slots for other tasks.
This method is safe to call with an empty attemptID — it returns nil without querying the store. It is idempotent: releasing an already-released claim is a no-op in the underlying SQL.
The worker wire protocol calls this method when terminal task status messages arrive from workers.
func (*Scheduler) RetryJob ¶
RetryJob revives every failed/canceled task in the job. It returns the number of tasks revived (0 when none were eligible — an idempotent no-op).
func (*Scheduler) RetryTask ¶
RetryTask revives a single failed/canceled task, looking up its job to drive dependency re-resolution.
func (*Scheduler) Run ¶
Run starts all scheduler goroutines and blocks until ctx is canceled. It returns after all goroutines have exited cleanly.
func (*Scheduler) Stop ¶
func (s *Scheduler) Stop()
Stop cancels the scheduler's internal context, triggering a clean shutdown. It does not wait for goroutines to finish — that happens inside [Run].
func (*Scheduler) WakeQueue ¶
WakeQueue wakes any parked lease waiters on queueID. Called by the API job handler after a successful submission so newly-ready tasks are leased without waiting out the long-poll hold. It also wakes queue-unaffiliated workers parked under bus.WildcardQueueToken, since they are eligible for any queue's work.
func (*Scheduler) WorkerTimeout ¶
WorkerTimeout returns the effective heartbeat-timeout threshold after normalization. The API layer uses it to decide whether a disabled worker is dead (and therefore removable) from its last-heartbeat age.