From bc3f76a57c7771e3bb3a8c34c05b8373d62d7e82 Mon Sep 17 00:00:00 2001
From: John Kirkham <jakirkham@gmail.com>
Date: Mon, 8 Feb 2021 19:44:28 -0800
Subject: [PATCH] Make `_unknown_durations` an ordinary `dict`

As `defaultdict`s are difficult for Cython to optimize and usually a
`dict` will suffice, change `_unknown_duractions` to a `dict` to allow
Cython to use Python C API specific to `dict`s.
---
 distributed/scheduler.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/distributed/scheduler.py b/distributed/scheduler.py
index 40d84b9b2f9..a444ae1337d 100644
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -1593,7 +1593,7 @@ class SchedulerState:
     _task_metadata: dict
     _total_nthreads: Py_ssize_t
     _total_occupancy: double
-    _unknown_durations: object
+    _unknown_durations: dict
     _unrunnable: set
     _validate: bint
     _workers: object
@@ -1645,7 +1645,7 @@ def __init__(
         self._task_metadata = dict()
         self._total_nthreads = 0
         self._total_occupancy = 0
-        self._unknown_durations = defaultdict(set)
+        self._unknown_durations = dict()
         if unrunnable is not None:
             self._unrunnable = unrunnable
         else:
@@ -2645,8 +2645,11 @@ def get_task_duration(self, ts: TaskState, default: double = -1) -> double:
         """
         duration: double = ts._prefix._duration_average
         if duration < 0:
-            s: set = self._unknown_durations[ts._prefix._name]
+            s: set = self._unknown_durations.get(ts._prefix._name)
+            if s is None:
+                self._unknown_durations[ts._prefix._name] = s = set()
             s.add(ts)
+
             if default < 0:
                 duration = UNKNOWN_TASK_DURATION
             else: