scheduler

type CandidateGangIterator[T LegacySchedulerJob] struct {
	SchedulingConstraints
	SchedulingRoundReport *SchedulingRoundReport[T]
	// contains filtered or unexported fields
}

type Configuration struct {
	// Database configuration
	Postgres configuration.PostgresConfig
	// Metrics configuration
	Metrics configuration.MetricsConfig
	// General Pulsar configuration
	Pulsar configuration.PulsarConfig
	// Pulsar subscription name
	SubscriptionName string
	// Maximum time since the last batch before a batch will be inserted into the database
	BatchDuration time.Duration
	// Time for which the pulsar consumer will wait for a new message before retrying
	PulsarReceiveTimeout time.Duration
	// Time for which the pulsar consumer will back off after receiving an error on trying to receive a message
	PulsarBackoffTime time.Duration
}

type ExecutorApi struct {
	// contains filtered or unexported fields
}

type JobDb struct {
	// In-memory database. Stores *SchedulerJob.
	// Used to efficiently iterate over jobs in sorted order.
	Db *memdb.MemDB
}

type JobDbAdapter struct {
	// contains filtered or unexported fields
}

type JobIterator[T LegacySchedulerJob] interface {
	Next() (T, error)
}

type JobQueueIterator struct {
	// contains filtered or unexported fields
}

type JobRepository interface {
	GetQueueJobIds(queueName string) ([]string, error)
	GetExistingJobsByIds(ids []string) ([]*api.Job, error)
}

type JobRepositoryAdapter struct {
	// contains filtered or unexported fields
}

type JobRun struct {
	// Unique identifier for the run
	RunID uuid.UUID
	// The name of the executor this run has been leased to
	Executor string
	// True if the job has been reported as pending by the executor
	Pending bool
	// True if the job has been reported as running by the executor
	Running bool
	// True if the job has been reported as succeeded by the executor
	Succeeded bool
	// True if the job has been reported as failed by the executor
	Failed bool
	// True if the job has been reported as cancelled by the executor
	Cancelled bool
	// True if the job has been returned by the executor
	Returned bool
	// True if the job has been expired by the scheduler
	Expired bool
}

type JobSchedulingReport[T LegacySchedulerJob] struct {
	// Time at which this report was created.
	Timestamp time.Time
	// Id of the job this pod corresponds to.
	JobId uuid.UUID
	// Job spec.
	Job T
	// Scheduling requirements of this job.
	// We currently require that each job contains exactly one pod spec.
	Req *schedulerobjects.PodRequirements
	// Executor this job was attempted to be assigned to.
	ExecutorId string
	// Reason for why the job could not be scheduled.
	// Empty if the job was scheduled successfully.
	UnschedulableReason string
	// Scheduling reports for the individual pods that make up the job.
	PodSchedulingReports []*PodSchedulingReport
}

type KubernetesLeaderController struct {
	// contains filtered or unexported fields
}

type LeaderConfig struct {
	// Name of the K8s Lock Object
	LeaseLockName string
	// Namespace of the K8s Lock Object
	LeaseLockNamespace string
	// The name of the pod
	PodName string
	// How long the lease is held for.
	// Non leaders much wait this long before trying to acquire the lease
	LeaseDuration time.Duration
	// RenewDeadline is the duration that the acting leader will retry refreshing leadership before giving up.
	RenewDeadline time.Duration
	// RetryPeriod is the duration the LeaderElector clients should waite between tries of actions.
	RetryPeriod time.Duration
}

type LeaderController interface {
	// GetToken returns a LeaderToken which allows you to determine if you are leader or not
	GetToken() LeaderToken
	// ValidateToken allows a caller to determine whether a previously obtained token is still valid.
	// Returns true if the token is a leader and false otherwise
	ValidateToken(tok LeaderToken) bool
}

type LeaderToken struct {
	// contains filtered or unexported fields
}

type LeaseListener interface {
	// contains filtered or unexported methods
}

type LegacyScheduler[T LegacySchedulerJob] struct {
	SchedulingConstraints
	SchedulingRoundReport *SchedulingRoundReport[T]
	CandidateGangIterator *CandidateGangIterator[T]
	// Contains all nodes to be considered for scheduling.
	// Used for matching pods with nodes.
	NodeDb *NodeDb
	// Used to request jobs from Redis and to mark jobs as leased.
	JobRepository SchedulerJobRepository[T]
	// Jobs are grouped into gangs by this annotation.
	GangIdAnnotation string
	// Jobs in a gang specify the number of jobs in the gang via this annotation.
	GangCardinalityAnnotation string
	// contains filtered or unexported fields
}

type LegacySchedulerJob interface {
	GetQueue() string
	GetAnnotations() map[string]string
	GetId() string
}

type LegacySchedulingAlgo struct {
	// contains filtered or unexported fields
}

type NodeDb struct {
	// contains filtered or unexported fields
}

type NodeItemAvailableResourceIndex struct {
	// Resource name, e.g., "cpu", "gpu", or "memory".
	Resource string
	// Job priority.
	Priority int32
}

type NodeTypeResourceIterator struct {
	// contains filtered or unexported fields
}

type NodeTypesResourceIterator struct {
	// contains filtered or unexported fields
}

type NodeTypesResourceIteratorItem struct {
	// contains filtered or unexported fields
}

type NodeTypesResourceIteratorPQ []*NodeTypesResourceIteratorItem

type NodesIterator struct {
	// contains filtered or unexported fields
}

type PodSchedulingReport struct {
	// Time at which this report was created.
	Timestamp time.Time
	// Pod scheduling requirements.
	Req *schedulerobjects.PodRequirements
	// Resource type determined by the scheduler to be the hardest to satisfy
	// the scheduling requirements for.
	DominantResourceType string
	// Node the pod was assigned to.
	// If nil, the pod could not be assigned to any Node.
	Node *schedulerobjects.Node
	// Score indicates how well the pod fits on the selected Node.
	Score int
	// Number of Node types that
	NumMatchedNodeTypes int
	// Number of Node types excluded by reason.
	NumExcludedNodeTypesByReason map[string]int
	// Number of nodes excluded by reason.
	NumExcludedNodesByReason map[string]int
}

type Publisher interface {
	// PublishMessages will publish the supplied messages. A LeaderToken is provided and the
	// implementor may decide whether to publish based on the status of this token
	PublishMessages(ctx context.Context, events []*armadaevents.EventSequence, shouldPublish func() bool) error

	// PublishMarkers publishes a single marker message for each Pulsar partition.  Each marker
	// massage contains the supplied group id, which allows all marker messages for a given call
	// to be identified.  The uint32 returned is the number of messages published
	PublishMarkers(ctx context.Context, groupId uuid.UUID) (uint32, error)
}

type PulsarPublisher struct {
	// contains filtered or unexported fields
}

type QueueCandidateGangIterator[T LegacySchedulerJob] struct {
	SchedulingConstraints
	QueueSchedulingRoundReport *QueueSchedulingRoundReport[T]
	// contains filtered or unexported fields
}

type QueueSchedulingReport[T LegacySchedulerJob] struct {
	// Queue name.
	Name                                      string
	MostRecentSuccessfulJobSchedulingReport   *JobSchedulingReport[T]
	MostRecentUnsuccessfulJobSchedulingReport *JobSchedulingReport[T]
}

type QueueSchedulingRoundReport[T LegacySchedulerJob] struct {
	// These factors influence the fraction of resources assigned to each queue.
	PriorityFactor float64
	// Resources assigned to the queue across all clusters at the start of the scheduling cycle.
	InitialResourcesByPriority schedulerobjects.QuantityByPriorityAndResourceType
	// Resources assigned to this queue during this scheduling cycle.
	ScheduledResourcesByPriority schedulerobjects.QuantityByPriorityAndResourceType
	// Reports for all successful job scheduling attempts.
	SuccessfulJobSchedulingReports map[uuid.UUID]*JobSchedulingReport[T]
	// Reports for all unsuccessful job scheduling attempts.
	UnsuccessfulJobSchedulingReports map[uuid.UUID]*JobSchedulingReport[T]
	// Total number of jobs successfully scheduled in this round for this queue.
	NumScheduledJobs int
	// contains filtered or unexported fields
}

type QueuedGangIterator[T LegacySchedulerJob] struct {
	// contains filtered or unexported fields
}

type QueuedJobsIterator struct {
	// contains filtered or unexported fields
}

type Scheduler struct {
	// contains filtered or unexported fields
}

type SchedulerJob struct {
	// String representation of the job id
	JobId string
	// Name of the queue this job belongs to.
	Queue string
	// Jobset the job belongs to
	// We store this as it's needed for sending job event messages
	Jobset string
	// Per-queue priority of this job.
	Priority uint32
	// Logical timestamp indicating the order in which jobs are submitted.
	// Jobs with identical Queue and Priority
	// are sorted by timestamp.
	Timestamp int64
	// Name of the executor to which this job has been assigned.
	// Empty if this job has not yet been assigned.
	Executor string
	// Name of the node to which this job has been assigned.
	// Empty if this job has not yet been assigned.
	Node string
	// True if the job is currently queued.
	// If this is set then the job will not be considered for scheduling
	Queued bool

	// True if the user has requested this job be cancelled
	CancelRequested bool
	// True if the scheduler has cancelled the job
	Cancelled bool
	// True if the scheduler has failed the job
	Failed bool
	// True if the scheduler has marked the job as succeeded
	Succeeded bool
	// Job Runs in the order they were received.
	// For now there can be only one active job run which will be the last element of the slice
	Runs []*JobRun
	// contains filtered or unexported fields
}

type SchedulerJobRepository[T LegacySchedulerJob] interface {
	// GetJobIterator returns a iterator over queued jobs for a given queue.
	GetJobIterator(ctx context.Context, queue string) (JobIterator[T], error)
	// TryLeaseJobs tries to create jobs leases and returns the jobs that were successfully leased.
	// Leasing may fail, e.g., if the job was concurrently leased to another executor.
	TryLeaseJobs(clusterId string, queue string, jobs []T) ([]T, error)
}

type SchedulingAlgo interface {
	// Schedule should assign jobs to nodes
	// Any jobs that are scheduled should be marked as such in the JobDb using the transaction provided
	// It should return a slice containing all scheduled jobs.
	Schedule(txn *memdb.Txn, jobDb *JobDb) ([]*SchedulerJob, error)
}

type SchedulingConstraints struct {
	Priorities      []int32
	PriorityClasses map[string]configuration.PriorityClass
	// Executor for which we're currently scheduling jobs.
	ExecutorId string
	// Resource pool of this executor.
	Pool string
	// Weights used when computing total resource usage.
	ResourceScarcity map[string]float64
	// Max number of jobs to scheduler per lease jobs call.
	MaximumJobsToSchedule uint
	// Max number of consecutive unschedulable jobs to consider for a queue before giving up.
	MaxConsecutiveUnschedulableJobs uint
	// Jobs leased to this executor must be at least this large.
	// Used, e.g., to avoid scheduling CPU-only jobs onto clusters with GPUs.
	MinimumJobSize schedulerobjects.ResourceList
	// Per-queue resource limits.
	// Map from resource type to the limit for that resource.
	MaximalResourceFractionPerQueue map[string]float64
	// Limit- as a fraction of total resources across worker clusters- of resource types at each priority.
	// The limits are cumulative, i.e., the limit at priority p includes all higher levels.
	MaximalCumulativeResourceFractionPerQueueAndPriority map[int32]map[string]float64
	// Max resources to schedule per queue at a time.
	MaximalResourceFractionToSchedulePerQueue map[string]float64
	// Max resources to schedule at a time.
	MaximalResourceFractionToSchedule map[string]float64
	// Total resources across all worker clusters.
	// Used when computing resource limits.
	TotalResources schedulerobjects.ResourceList
}

type SchedulingReportsRepository[T LegacySchedulerJob] struct {
	// Scheduling reports for the jobs that were most recently attempted to be scheduled.
	MostRecentJobSchedulingReports *lru.Cache
	// Scheduling reports for the most recently seen queues.
	MostRecentQueueSchedulingReports *lru.Cache
}

type SchedulingRoundReport[T LegacySchedulerJob] struct {
	// Time at which the scheduling cycle started.
	Started time.Time
	// Time at which the scheduling cycle finished.
	Finished time.Time
	// Executor for which the scheduler was invoked.
	Executor string
	// Per-queue scheduling reports.
	QueueSchedulingRoundReports map[string]*QueueSchedulingRoundReport[T]
	// Total resources across all clusters available at the start of the scheduling cycle.
	TotalResources schedulerobjects.ResourceList
	// Resources assigned across all queues during this scheduling cycle.
	ScheduledResourcesByPriority schedulerobjects.QuantityByPriorityAndResourceType
	// Total number of jobs successfully scheduled in this round.
	NumScheduledJobs int
	// Reason for why the scheduling round finished.
	TerminationReason string
	// contains filtered or unexported fields
}

type StandaloneLeaderController struct {
	// contains filtered or unexported fields
}

type SubmitChecker struct {
	// contains filtered or unexported fields
}