fix: required affinity (#812)

* fix: required affinity * chore: rm dead code
hatchet-dev · Aug 23, 2024 · ee5d867 · ee5d867
1 parent 53be615
commit ee5d867
Show file tree

Hide file tree

Showing 6 changed files with 75 additions and 7 deletions.
diff --git a/pkg/scheduling/affinity_test.go b/pkg/scheduling/affinity_test.go
@@ -75,6 +75,13 @@ func TestComputeWeight(t *testing.T) {
 		{
 			name: "No match for a required returns -1",
 			desiredLabels: []*dbsqlc.GetDesiredLabelsRow{
+				{
+					Key:        "memory",
+					Comparator: dbsqlc.WorkerLabelComparatorEQUAL,
+					StrValue:   sqlchelpers.TextFromStr("100"),
+					Weight:     15,
+					Required:   false,
+				},
 				{
 					Key:        "region",
 					Comparator: dbsqlc.WorkerLabelComparatorEQUAL,

diff --git a/pkg/scheduling/fixtures/affinity_hard.json b/pkg/scheduling/fixtures/affinity_hard.json
@@ -52,6 +52,18 @@
       "tenantId": "707d0855-80ab-4e1f-a156-f1c4546cbf52",
       "queue": "child:process2",
       "Order": 0
+    },
+    {
+      "id": 152260,
+      "stepRunId": "064e787a-6cfd-4f82-8b6d-d8031459fdef",
+      "stepId": "e2c744b8-bebb-4b05-8292-d2ce36f27929",
+      "actionId": "child:process2",
+      "scheduleTimeoutAt": "2024-08-07T21:12:49.56Z",
+      "stepTimeout": "60s",
+      "isQueued": true,
+      "tenantId": "707d0855-80ab-4e1f-a156-f1c4546cbf52",
+      "queue": "child:process2",
+      "Order": 0
     }
   ],
   "WorkerLabels": {
@@ -86,6 +98,15 @@
             "Required": true,
             "Weight": 1
           }
+        ],
+      "e2c744b8-bebb-4b05-8292-d2ce36f27929": [
+          {
+            "Key": "MODEL",
+            "Comparator": "EQUAL",
+            "StrValue": "C",
+            "Required": true,
+            "Weight": 1
+          }
         ]
     }
 }
diff --git a/pkg/scheduling/fixtures/affinity_hard_output.json b/pkg/scheduling/fixtures/affinity_hard_output.json
@@ -0,0 +1,45 @@
+{
+  "StepRunIds": [
+    "16fa711e-c03e-435c-88d9-62adf1591d98",
+    "064e787a-6cfd-4f82-8b6d-d8031459fdee"
+  ],
+  "StepRunTimeouts": [
+    "60s",
+    "60s"
+  ],
+  "SlotIds": [
+    "8cf68f09-b914-4f31-9777-8082b751a2d4",
+    "aed946be-43db-4b5f-876a-c6c71f4d52aa"
+  ],
+  "WorkerIds": [
+    "aaaaaaaa-43db-4b5f-876a-c6c71f4d52aa",
+    "bbbbbbbb-43db-4b5f-876a-c6c71f4d52aa"
+  ],
+  "UnassignedStepRunIds": [
+    "064e787a-6cfd-4f82-8b6d-d8031459fdef"
+  ],
+  "QueuedStepRuns": [
+    {
+      "StepRunId": "16fa711e-c03e-435c-88d9-62adf1591d98",
+      "WorkerId": "aaaaaaaa-43db-4b5f-876a-c6c71f4d52aa",
+      "DispatcherId": "9994a9eb-430d-46da-934d-d9dd953cfd21"
+    },
+    {
+      "StepRunId": "064e787a-6cfd-4f82-8b6d-d8031459fdee",
+      "WorkerId": "bbbbbbbb-43db-4b5f-876a-c6c71f4d52aa",
+      "DispatcherId": "9994a9eb-430d-46da-934d-d9dd953cfd21"
+    }
+  ],
+  "TimedOutStepRuns": [],
+  "RateLimitedStepRuns": [],
+  "RateLimitedQueues": {},
+  "QueuedItems": [
+    137295,
+    152259
+  ],
+  "ShouldContinue": false,
+  "MinQueuedIds": {
+    "child:process2": 0
+  },
+  "RateLimitUnitsConsumed": {}
+}
diff --git a/pkg/scheduling/fixtures/affinity_output.json → ...duling/fixtures/affinity_soft_output.json b/pkg/scheduling/fixtures/affinity_output.json → ...duling/fixtures/affinity_soft_output.json
diff --git a/pkg/scheduling/scheduling_test.go b/pkg/scheduling/scheduling_test.go
@@ -156,7 +156,7 @@ func TestGeneratePlan(t *testing.T) {
 			name: "GeneratePlan_Affinity_Soft",
 			args: args{
 				fixtureArgs:   "./fixtures/affinity_soft.json",
-				fixtureResult: "./fixtures/affinity_output.json",
+				fixtureResult: "./fixtures/affinity_soft_output.json",
 				noTimeout:     true,
 			},
 			want: func(s SchedulePlan, fixtureResult string) bool {
@@ -175,7 +175,7 @@ func TestGeneratePlan(t *testing.T) {
 			name: "GeneratePlan_Affinity_Hard",
 			args: args{
 				fixtureArgs:   "./fixtures/affinity_hard.json",
-				fixtureResult: "./fixtures/affinity_output.json",
+				fixtureResult: "./fixtures/affinity_hard_output.json",
 				noTimeout:     true,
 			},
 			want: func(s SchedulePlan, fixtureResult string) bool {

diff --git a/pkg/scheduling/worker_state_manager.go b/pkg/scheduling/worker_state_manager.go
@@ -38,11 +38,6 @@ func NewWorkerStateManager(
 		for workerId, worker := range workers {
 			weight := ComputeWeight(desired, worker.labels)
 
-			// skip workers that are not a match (i.e. required)
-			if weight < 0 {
-				continue
-			}
-
 			// cache the weight on the worker
 			workers[workerId].AddStepWeight(stepId, weight)