DRAFT: Sticky Step Assignment (assigning steps to workers with specific state)

[grutt]

DRAFT: This Proposal is a rough brain dump and is not in a final state. Feedback and comments will help us ensure this feature is complete!

Problem

Hatchet customers often have steps that require specific worker states or resources, such as loaded machine learning models in memory. Currently, when executing steps, Hatchet does not consider the state of the workers, leading to potential issues such as:

• Inefficient resource utilization: Steps are assigned to workers without considering their current state, resulting in unnecessary resource loading and initialization.
• Increased latency: If a worker needs to load a resource (e.g., a machine learning model) for each step, it introduces additional latency in step execution.
• Suboptimal worker assignment: Steps are not prioritized to workers that already have the required resources loaded, leading to suboptimal worker assignment and reduced performance.

To address these problems, we propose a feature called "Sticky Steps" that allows steps to express their resource requirements and favors assigning them to workers that are already in a compatible state.

Proposed Solution

Step Resource Requirements

Introduce a new field called requiredResources in the step definition object. This field will be an array of strings representing the resources or states required by the step.

Note: This is very much a 1st draft as it does not afford a resource key with multiple states. In other words, it assumes there's only one slot per resource which is likely not correct.

interface Step {
  // ...existing fields...
  requiredResources?: Record<string, string>;
}

Example usage:

const step1: Step = {
  // ...
  requiredResources: {"model": "modelA", "dataset": "datasetX"},
};

const step2: Step = {
  // ...
  requiredResources: {"model": "modelB", "dataset": "datasetY"},
};

In this example, step1 requires the resources "modelA" and "datasetX", while step2 requires the resources "modelB" and "datasetY".

Worker State Management

Introduce a new method called updateWorkerResourceState in the context object passed to the step function. This method allows the step to update the worker's state, indicating the resources or states that the worker has acquired during the step execution.

interface Context {
  // ...existing fields...
  updateWorkerResourceState: (state: Record<string, any>) => void;
}

Example usage:

const stepFunction = async (context: StepContext) => {
  // Load the required resources
  const modelA = await loadModelA();
  const datasetX = await loadDatasetX();

  // Update the worker state
  context.updateWorkerResourceState({
    modelA: 'modelA',
    datasetX: 'datasetX',
  });

  // Use the loaded resources in the step execution
  // ...
};

In this example, the step function loads the required resources (e.g., modelA and datasetX) and updates the worker state using the context.updateWorkerResourceState method. The worker state is represented as an object where the keys are the resource names and the values are the loaded resources.

Sticky Step Scheduling

When a step is ready to be executed, the Hatchet engine will follow these steps:

1. Check if the step has any requiredResources specified.
2. If requiredResources are specified, search for workers that have all the required resources available in their state.
   • If one or more compatible workers are found, prioritize assigning the step to one of those workers.
   • If no compatible worker is found, assign the step to the next available worker slot.
   • QUESTION: do we favor partial state?
3. If no requiredResources are specified, follow the default worker assignment logic.

When a step is assigned to a worker that doesn't have the required resources in its state, the worker will load and initialize the necessary resources before executing the step. After the step is completed, the worker will update its state using the context.updateWorkerResourceState method to reflect the acquired resources.

Resource Eviction and Replacement

If a worker needs to execute a step that requires resources different from its current state, it will replace the existing resources with the newly required ones. This ensures that workers can adapt to changing step requirements while optimizing for resource reuse when possible.

Risks and Considerations

• Increased complexity in the Hatchet engine to track worker states and make scheduling decisions based on resource requirements.
• Potential resource fragmentation if steps with different resource requirements are interleaved, leading to frequent resource eviction and reloading.
• Need for steps to accurately update the worker state using the context.updateWorkerResourceState method and handle resource initialization and cleanup efficiently.
  • This does not account for resources that may remain in memory after another resource is loaded... this is poorly assuming that there is only one resource per key which may be an incorrect assumption.
• Balancing the benefits of resource reuse with the overall system throughput and fairness in step assignment.
• Handling scenarios where multiple compatible workers are available and deciding whether to concentrate steps on fewer workers or distribute them evenly.
• Considering the impact of worker failures or slowdowns on steps that rely on specific worker states.
• Providing monitoring and logging capabilities to track resource utilization and optimize sticky step scheduling.

[Cixelyn]

Dynamic sticky states seems great, and seems to solve a lot of challenges that current orchestrators have with modern ML workloads!

Even the static variant alone is quite useful, this enables a bunch of use cases, e.g.:

• Prioritizing workers running the latest codebase during a large-scale rolling restart.
• Prioritizing faster workers in cases of a heterogenous cluster.
• Prioritizing on-prem machines vs. pre-emptible cloud machines in a hybrid environment.
• Prioritizing workers in a particular geographic region for e.g. GDPR compliance or latency reasons.

---

Conceptually, I think the design seems similar to k8s's node affinity feature using the preferredDuringSchedulingIgnoredDuringExecution flag.

I do quite like the Record<string, any> since feels natural and more flexible. However if you did want to simplify the design, a simple tag list of resources currently available to the worker might be sufficient, e.g. Array<string>. The API could then be modified accordingly to ctx.updateWorkerResourceState([modelA", "modelB", "datasetX"]), or potentially just setWorkerResourceState(key, boolean).

Unless I'm misunderstanding the current proposal, since everything is done via string matching on the resource, a tag list design should be equivalent in selection power. I think where the Record<string, any> becomes more powerful is if once could do fancier selectors, like requiredResources: {"FreeMemory": "(x) => x > 100"}. But I think right now the proposed issue w/ multi-state keys can be translated into a flat-list of tags, with only a little bit of extra burden on the workflow creator to design the tag space correctly.

QUESTION: do we favor partial state?

K8S's design allows for a node affinity weight. If the design was modified to the tag list, the API call could then be requiredResources: <string, number> where number indicates the weighting of finding a particular tag.

More ambitious would be to modify to requiredResources <string, number | boolean> where a value of true is a hard constraint on scheduling, similar to requiredDuringSchedulingIgnoredDuringExecution