Worker Affinity Assignment (Beta)

It is often desirable to assign workflows to specific workers based on certain criteria, such as worker capabilities, resource availability, or location. Worker affinity allows you to specify that a workflow should be assigned to a specific worker based on worker label state. Labels can be set dynamically on workers to reflect their current state, such as a specific model loaded into memory or specific disk requirements.

Specific tasks can then specify desired label state to ensure that workflows are assigned to workers that meet specific criteria. If no worker meets the specified criteria, the task run will remain in a pending state until a suitable worker becomes available or the task is cancelled. (See Scheduling Timeouts)

Specifying Worker Labels

Labels can be set on workers when they are registered with Hatchet. Labels are key-value pairs that can be used to specify worker capabilities, resource availability, or other criteria that can be used to match workflows to workers. Values can be strings or numbers, and multiple labels can be set on a worker.

worker = hatchet.worker(
    "affinity-worker",
    slots=10,
    labels={
        "model": "fancy-ai-model-v2",
        "memory": 512,
    },
    workflows=[affinity_worker_workflow],
)
worker.start()

const worker1 = await hatchet.worker('affinity-worker-1', {
  labels: {
    model: 'abc',
    memory: 1024,
  },
});

	w, err := worker.NewWorker(
		worker.WithClient(
			c,
		),
		worker.WithLabels(map[string]interface{}{
			"model":  "fancy-ai-model-v2",
			"memory": 512,
		}),
	)

Specifying Step Desired Labels

You can specify desired worker label state for specific tasks in a workflow by setting the desired_worker_labels property on the task definition. This property is an object where the keys are the label keys and the values are objects with the following properties:

• value: The desired value of the label
• comparator (default: EQUAL): The comparison operator to use when matching the label value.
  • EQUAL: The label value must be equal to the desired value
  • NOT_EQUAL: The label value must not be equal to the desired value
  • GREATER_THAN: The label value must be greater than the desired value
  • GREATER_THAN_OR_EQUAL: The label value must be greater than or equal to the desired value
  • LESS_THAN: The label value must be less than the desired value
  • LESS_THAN_OR_EQUAL: The label value must be less than or equal to the desired value
• required (default: true): Whether the label is required for the task to run. If true, the task will remain in a pending state until a worker with the desired label state becomes available. If false, the worker will be prioritized based on the sum of the highest matching weights.
• weight (optional, default: 100): The weight of the label. Higher weights are prioritized over lower weights when selecting a worker for the task. If multiple workers have the same highest weight, the worker with the highest sum of weights will be selected. Ignored if required is true.

affinity_worker_workflow = hatchet.workflow(name="AffinityWorkflow")


@affinity_worker_workflow.task(
    desired_worker_labels={
        "model": DesiredWorkerLabel(value="fancy-ai-model-v2", weight=10),
        "memory": DesiredWorkerLabel(
            value=256,
            required=True,
            comparator=WorkerLabelComparator.LESS_THAN,
        ),
    },
)

const workflow: Workflow = {
  id: 'affinity-workflow',
  description: 'test',
  steps: [
    {
      name: 'step1',
      run: async (ctx) => {
        const results: Promise<any>[] = [];
        // eslint-disable-next-line no-plusplus
        for (let i = 0; i < 50; i++) {
          const result = await ctx.spawnWorkflow(childWorkflow.id, {});
          results.push(result.output);
        }
        console.log('Spawned 50 child workflows');
        console.log('Results:', await Promise.all(results));

        return { step1: 'step1 results!' };
      },
    },
  ],
};

	err = w.RegisterWorkflow(
		&worker.WorkflowJob{
			On:          worker.Events("user:create:affinity"),
			Name:        "affinity",
			Description: "affinity",
			Steps: []*worker.WorkflowStep{
				worker.Fn(func(ctx worker.HatchetContext) (result *taskOneOutput, err error) {
					return &taskOneOutput{
						Message: ctx.Worker().ID(),
					}, nil
				}).
					SetName("task-one").
					SetDesiredLabels(map[string]*types.DesiredWorkerLabel{
						"model": {
							Value:  "fancy-ai-model-v2",
							Weight: 10,
						},
						"memory": {
							Value:      512,
							Required:   true,
							Comparator: types.ComparatorPtr(types.WorkerLabelComparator_GREATER_THAN),
						},
					}),
			},
		},
	)

Dynamic Worker Labels

Labels can also be set dynamically on workers using the upsertLabels method. This can be useful when worker state changes over time, such as when a new model is loaded into memory or when a worker’s resource availability changes.

async def step(input: EmptyModel, ctx: Context) -> dict[str, str | None]:
    if ctx.worker.labels().get("model") != "fancy-ai-model-v2":
        ctx.worker.upsert_labels({"model": "unset"})
        # DO WORK TO EVICT OLD MODEL / LOAD NEW MODEL
        ctx.worker.upsert_labels({"model": "fancy-ai-model-v2"})

    return {"worker": ctx.worker.id()}

const affinity: Workflow = {
  id: "dynamic-affinity-workflow",
  description: "test",
  tasks: [
    {
      name: "child-task1",
      worker_labels: {
        model: {
          value: "fancy-vision-model",
          required: false,
        },
      },
      run: async (ctx) => {
        if (ctx.worker.labels().model !== "fancy-vision-model") {
          await ctx.worker.upsertLabels({ model: undefined });
          await evictModel();
          await loadNewModel("fancy-vision-model");
          await ctx.worker.upsertLabels({ model: "fancy-vision-model" });
        }
        // DO WORK
        return { childStep1: "childStep1 results!" };
      },
    },
  ],
};

	err = w.RegisterWorkflow(
		&worker.WorkflowJob{
			On:          worker.Events("user:create:affinity"),
			Name:        "affinity",
			Description: "affinity",
			Steps: []*worker.WorkflowStep{
				worker.Fn(func(ctx worker.HatchetContext) (result *taskOneOutput, err error) {
 
    				model := ctx.Worker().GetLabels()["model"]
 
    				if model != "fancy-vision-model" {
    					ctx.Worker().UpsertLabels(map[string]interface{}{
    						"model": nil,
    					})
    					// Do something to load the model
            evictModel();
            loadNewModel("fancy-vision-model");
    					ctx.Worker().UpsertLabels(map[string]interface{}{
    						"model": "fancy-vision-model",
    					})
    				}
 
    				return &taskOneOutput{
    					Message: ctx.Worker().ID(),
    				}, nil
    			}).
    				SetName("task-one").
    				SetDesiredLabels(map[string]*types.DesiredWorkerLabel{
    					"model": {
    						Value:  "fancy-vision-model",
    						Weight: 10,
    					},
    					"memory": {
    						Value:      512,
    						Required:   true,
    						Comparator: types.WorkerLabelComparator_GREATER_THAN,
    					},
    				}),
    		},
    	},
    )