resource_manager.go

package core

import (
	"context"
)

//go:generate enumer -type=AllocationStatus -trimprefix=AllocationStatus

type AllocationStatus int

const (
	// This is the enum returned when there's an error
	AllocationUndefined AllocationStatus = iota

	// Go for it
	AllocationStatusGranted

	// This means that no resources are available globally.  This is the only rejection message we use right now.
	AllocationStatusExhausted

	// We're not currently using this - but this would indicate that things globally are okay, but that your
	// own namespace is too busy
	AllocationStatusNamespaceQuotaExceeded
)

const namespaceSeparator = ":"

type ResourceNamespace string

func (r ResourceNamespace) CreateSubNamespace(namespace ResourceNamespace) ResourceNamespace {
	return r + namespaceSeparator + namespace
}

type ResourceRegistrar interface {
	RegisterResourceQuota(ctx context.Context, namespace ResourceNamespace, quota int) error
}

// ResourceManager Interface
// 1. Terms and definitions
//
//   - Resource: resource is an abstraction of anything that has a limited quota of units and can be claimed in a
//     single unit or multiple units at once. At Flyte's current state, a resource means a logical
//     separation (e.g., a cluster) of an external service that allows a limited number of outstanding
//     requests to be sent to.
//
//   - Token: Flyte uses a token to serve as the placeholder to represent a unit of resource. Flyte resource manager
//     manages resources by managing the tokens of the resources.
//
//     2. Description
//     ResourceManager provides a task-type-specific pooling system for Flyte Tasks. Plugin writers can optionally
//     request for resources in their tasks, in single quantity.
//
//     3. Usage
//     A Flyte plugin registers the resources and the desired quota of each resource with ResourceRegistrar at the
//     setup time of Flyte Propeller. At the end of the setup time, Flyte Propeller builds a ResourceManager based on
//     these registration requests.
//
//     During runtime, the ResourceManager does two simple things: allocating tokens and releasing tokens. When a Flyte
//     task execution wants to send a request to an external service, the plugin should claim a unit of the corresponding
//     resource. Specifically, an execution needs to generate a unique token, and register the token with ResourceManager
//     by calling ResourceManager's AllocateResource() function. ResourceManager will check its current utilization and
//     the allocation policy to decide whether or not to grant the request. Only when receiving the "AllocationGranted"
//     status shall this execution move forward and send out the request. The granted token will be recorded in a token
//     pool corresponding to the resource and managed by ResourceManager. When the request is done, the plugin will ask
//     the resource manager to release the token by calling ResourceManager's ReleaseResource(), and the token will be
//     erased from the corresponding pool.
//
//     4. Example
//     Flyte has a built-on Qubole plugin that allows Flyte tasks to send out Hive commands to Qubole.
//     In the plugin, a single Qubole cluster is a resource, and sending out a single Hive command to a Qubole cluster consumes
//     a token of the corresponding resource. The resource allocation is achieved by the Qubole plugin calling
//     status, err := AllocateResource(ctx, <cluster name>, <token string>, <constraint spec>)
//     and the de-allocation is achieved by the plugin calling
//     status, err := AllocateResource(ctx, <cluster name>, <token string>, <constraint spec>)
//
//     For example,
//     status, err := AllocateResource(ctx, "default_cluster", "flkgiwd13-akjdoe-0", ResourceConstraintsSpec{})
//     When the corresponding Hive command finishes, the plugin needs to make the following function call to release
//     the corresponding token
//     err := ReleaseResource(ctx, "default_cluster", "flkgiwd13-akjdoe-0")
type ResourceManager interface {
	GetID() string
	// During execution time, plugins can call AllocateResource() to register a token to the token pool associated with a resource with the resource manager.
	// If granted an allocation, the token will be recorded in the corresponding token pool until the same plugin releases it.
	// When calling AllocateResource, the plugin needs to specify a ResourceConstraintsSpec which contains resource capping constraints at different levels.
	// The ResourceConstraint pointers in ResourceConstraintsSpec, however, can be set to nil to present a non-constraint at that level
	AllocateResource(ctx context.Context, namespace ResourceNamespace, allocationToken string, constraintsSpec ResourceConstraintsSpec) (AllocationStatus, error)
	// During execution time, after an outstanding request is completed, the plugin need to use ReleaseResource() to release the allocation of the corresponding token
	// from the token pool in order to gain back the quota taken by the token
	ReleaseResource(ctx context.Context, namespace ResourceNamespace, allocationToken string) error
}

type ResourceConstraint struct {
	Value int64
}

// ResourceConstraintsSpec is a contract that a plugin can specify with ResourceManager to force runtime quota-allocation constraints
// at different levels.
//
// Setting constraints in a ResourceConstraintsSpec to nil objects is valid, meaning there's no constraint at the corresponding level.
// For example, a ResourceConstraintsSpec with nil ProjectScopeResourceConstraint and a non-nil NamespaceScopeResourceConstraint means
// that it only poses a cap at the namespace level. A zero-value ResourceConstraintsSpec means there's no constraints posed at any level.
type ResourceConstraintsSpec struct {
	ProjectScopeResourceConstraint   *ResourceConstraint
	NamespaceScopeResourceConstraint *ResourceConstraint
}