Resource Management and Quotas

When Flyte tasks interact with external services that have inherent concurrency limits, it's crucial to manage resource allocation to prevent overloading these services. Flyte addresses this by providing a robust resource management system centered around the ResourceManager interface, which allows plugins to claim and release "tokens" representing units of a limited "resource." This ensures system-wide concurrency control and prevents resource exhaustion.

At its core, Flyte's resource management operates on two key concepts:

• Resource: An abstraction for anything with a limited quota that can be claimed in single or multiple units. In practice, this often represents a logical separation of an external service, such as a specific cluster, that can handle a finite number of outstanding requests.
• Token: A placeholder representing a single unit of a resource. The ResourceManager manages resources by tracking these tokens.

Claiming and Releasing Resources with ResourceManager

When a Flyte task needs to perform an operation that consumes a unit of a limited external resource, its plugin must first claim a token from the ResourceManager. This ensures that the operation proceeds only if the resource is available and within its defined quotas. Once the operation is complete, the token must be released.

To claim a resource, a plugin calls the AllocateResource method on the ResourceManager instance:

AllocateResource(ctx context.Context, namespace ResourceNamespace, allocationToken string, constraintsSpec ResourceConstraintsSpec) (AllocationStatus, error)

• ctx: The standard Go context.
• namespace: Identifies the specific resource being requested (e.g., a cluster name like "default_cluster").
• allocationToken: A unique string identifying this specific allocation request. This token is used to track and later release the resource.
• constraintsSpec: An optional ResourceConstraintsSpec that allows specifying additional, granular resource capping constraints at different levels (e.g., project or namespace scope).

The AllocateResource method returns an AllocationStatus indicating whether the request was granted and an error if something went wrong. A common usage pattern, as seen in the flyteplugins/go/tasks/pluginmachinery/internal/webapi/allocation_token.go file, involves checking this status:

allocationStatus, err := tCtx.ResourceManager().AllocateResource(ctx, ns, token, constraints)
	if err != nil {
		logger.Errorf(ctx, "Failed to allocate resources for task. Error: %v", err)
		return nil, core.PhaseInfo{}, err
	}

	switch allocationStatus {
	case core.AllocationStatusGranted:
		// Resource successfully allocated, proceed with the task
		metrics.AllocationGranted.Inc(ctx)
		// ... task logic ...
	case core.AllocationStatusNamespaceQuotaExceeded:
		// The request exceeded the namespace-level quota
		metrics.AllocationNotGranted.Inc(ctx)
		logger.Infof(ctx, "Couldn't allocate token because allocation status is [%v].", allocationStatus.String())
		// ... handle waiting or retry ...
	case core.AllocationStatusExhausted:
		// No resources are available globally
		metrics.AllocationNotGranted.Inc(ctx)
		logger.Infof(ctx, "Couldn't allocate token because allocation status is [%v].", allocationStatus.String())
		// ... handle waiting or retry ...
	}

The AllocationStatus enum provides detailed feedback:

• AllocationStatusGranted: The resource token was successfully allocated, and the task can proceed.
• AllocationStatusExhausted: No resources are available globally, indicating the overall capacity for the resource has been reached.
• AllocationStatusNamespaceQuotaExceeded: The request exceeded the quota defined for the specific namespace, even if global resources might still be available.

Once the task has completed its interaction with the external service, it must release the allocated token by calling ReleaseResource:

ReleaseResource(ctx context.Context, namespace ResourceNamespace, allocationToken string) error

• ctx: The standard Go context.
• namespace: The same resource namespace used during allocation.
• allocationToken: The unique token string that was previously allocated.

For example, a plugin interacting with a Qubole cluster might allocate and release resources like this:

// Claim a token for "default_cluster" with a unique ID
status, err := AllocateResource(ctx, "default_cluster", "flkgiwd13-akjdoe-0", ResourceConstraintsSpec{})

// ... perform Hive command ...

// Release the token once the command finishes
err := ReleaseResource(ctx, "default_cluster", "flkgiwd13-akjdoe-0")

Defining System-Wide Quotas with ResourceRegistrar

Before the ResourceManager can allocate tokens, the system needs to know what resources are available and what their maximum capacities are. This is handled during the Flyte Propeller setup phase by the ResourceRegistrar interface.

Plugins or system components use ResourceRegistrar.RegisterResourceQuota to define the limits for a given resource:

err := iCtx.ResourceRegistrar().RegisterResourceQuota(ctx, ns, quota)

As seen in flyteplugins/go/tasks/pluginmachinery/internal/webapi/core.go, this call associates a ResourceNamespace (ns) with a specific quota (which defines the maximum number of concurrent tokens for that resource). This registration process builds the underlying ResourceManager instance, informing it of the total capacity for each managed resource.

Granular Constraints with ResourceConstraintsSpec

Beyond system-wide quotas, Flyte allows for more granular control over resource allocation using ResourceConstraintsSpec. This structure enables plugins to specify additional capping constraints that apply at different organizational levels, such as per project or per namespace.

The ResourceConstraintsSpec contains pointers to ResourceConstraint objects:

type ResourceConstraintsSpec struct {
	ProjectScopeResourceConstraint   *ResourceConstraint
	NamespaceScopeResourceConstraint *ResourceConstraint
}

type ResourceConstraint struct {
	Value int64
}

• ProjectScopeResourceConstraint: If set, this constraint limits the number of concurrent tokens that can be allocated across all tasks within a specific project for the given resource.
• NamespaceScopeResourceConstraint: If set, this constraint limits the number of concurrent tokens that can be allocated within a specific namespace (e.g., a domain or workflow execution) for the given resource.

Each ResourceConstraint simply holds an int64 Value representing the maximum allowed concurrency for that scope. It