refactor decide_worker method out of transition code
@@ -2290,24 +2290,7 @@ class Scheduler(ServerNode):
pdb.set_trace()
raise
- def transition_waiting_processing(self, key):
- try:
- if self.validate:
- assert key in self.waiting
- assert not self.waiting[key]
- assert key not in self.who_has
- assert key not in self.exceptions_blame
- assert key not in self.rprocessing
- # assert key not in self.readyset
- assert key not in self.unrunnable
- assert all(dep in self.who_has
- for dep in self.dependencies[key])
-
- if any(not self.who_has[dep] for dep in self.dependencies[key]):
- return {}
-
- del self.waiting[key]
-
+ def decide_worker(self, key):
valid_workers = self.valid_workers(key)
if not valid_workers and key not in self.loose_restrictions and self.ncores:
@@ -2331,6 +2314,30 @@ class Scheduler(ServerNode):
worker = self.workers[self.n_tasks % len(self.workers)]
assert worker
+ return worker
+
+
+ def transition_waiting_processing(self, key):
+ try:
+ if self.validate:
+ assert key in self.waiting
+ assert not self.waiting[key]
+ assert key not in self.who_has
+ assert key not in self.exceptions_blame
+ assert key not in self.rprocessing
+ # assert key not in self.readyset
+ assert key not in self.unrunnable
+ assert all(dep in self.who_has
+ for dep in self.dependencies[key])
+
+ if any(not self.who_has[dep] for dep in self.dependencies[key]):
+ return {}
+
+ del self.waiting[key]
+
+ worker = self.decide_worker(key)
+ if not worker:
+ return worker
ks = key_split(key)
