Enable mypy in CI 1/2 #5328

crusaderky · 2021-09-17T13:02:48Z

Partial work towards #2803.
This PR covers everything but scheduler.py and actually running mypy in CI.
Continues in #5348.

crusaderky · 2021-09-23T20:14:13Z

The good: all files except scheduler.py are now passing mypy checks.
The bad: putting # type: ignore on top of scheduler.py breaks everything else. So I'll have to go through it and it's probably going to end up in a heap of # type: ignore on individual lines.

jrbourbeau

Thanks @crusaderky!

As this is a relatively large change (though one I think there's interest in), we should be sure to get feedback from folks who regularly contribute to distributed. cc @fjetter, @jakirkham, @gjoseph92, @mrocklin, @jacobtomlinson, @madsbk thoughts on adding type annotations? (also feel free to rope others into this discussion)

madsbk · 2021-09-24T11:09:30Z

I think this is a great idea, thanks @crusaderky!

crusaderky · 2021-09-24T15:57:25Z

As discussed in the latest coiled meeting, I have reverted all changes to scheduler.py and the pre-commit hook, to publish them in a second iterative PR.
This PR is now ready for review.

jakirkham · 2021-09-24T16:06:42Z

Generally ok with this. However I do have one concern.

In particular I'd like us to be aware that scheduler.py has very particular annotations based on what Cython will accept and work with efficiently (it can take some other things, but run poorly instead, which make finding/fixing harder). My guess is Cython annotations do not map very well to how MyPy understands things. So would recommend we leave scheduler.py alone. If we opt to Cythonize more things in the future, we may need to deviate from MyPy there as well.

ian-r-rose

Thanks @crusaderky, this is awesome

ian-r-rose · 2021-09-24T16:31:35Z

distributed/_ipython_utils.py

@@ -135,7 +129,7 @@ def remote_magic(line, cell=None):


 # cache clients for re-use in remote magic
-remote_magic._clients = {}
+remote_magic._clients = {}  # type: ignore


Suggested change

remote_magic._clients = {} # type: ignore

remote_magic._clients: Dict[str, BlockingKernelClient] = {}

Thanks, fixed. We should always use lowercase dict[...], list[...], etc. and put from __future__ import annotations at the top of the module.

Correction: can't.

distributed/_ipython_utils.py:133: error: Type cannot be declared in assignment to non-self attribute distributed/_ipython_utils.py:133: error: "Callable[[Any, Any], Any]" has no attribute "_clients"

ian-r-rose · 2021-09-24T16:33:49Z

distributed/_version.py

@@ -50,7 +50,7 @@ class NotThisMethod(Exception):
    """Exception raised if a method is not valid for the current scenario."""


-LONG_VERSION_PY = {}
+LONG_VERSION_PY: dict = {}


I think this file is missing a from __future__ import __annotations to make this work.

that's only necessary if you use bracket-stile dict[...] or if the annotations create circular dependencies

ian-r-rose · 2021-09-24T16:40:29Z

distributed/client.py

@@ -105,7 +110,7 @@
 NO_DEFAULT_PLACEHOLDER = "_no_default_"


-def _get_global_client():
+def _get_global_client() -> Client | None:


Note to myself and others: 3.10 syntax, but backwards compatible with from __future__ import annotations, AFAIK.

correct. 3.10 syntax works on all python versions as long as your version of mypy is recent.
Notable exception: cast, since it is executed at runtime, needs Python 3.10 style annotations wrapped in a string.

distributed/dashboard/components/scheduler.py

distributed/dashboard/utils.py

ian-r-rose · 2021-09-24T16:48:34Z

distributed/deploy/adaptive_core.py

    def __init__(
        self,
        minimum: int = 0,
-        maximum: int = math.inf,
+        maximum: int | float = math.inf,


Not a huge deal, but I wonder if we could narrow further to int | Literal[math.inf]

Sadly, no: https://www.python.org/dev/peps/pep-0586/

The following are provisionally disallowed for simplicity. We can consider allowing them in future extensions of this PEP.

Floats: e.g. Literal[3.14]. Representing Literals of infinity or NaN in a clean way is tricky; real-world APIs are unlikely to vary their behavior based on a float parameter.

FYI: python/mypy#11208

ian-r-rose · 2021-09-24T16:52:28Z

distributed/protocol/serialize.py

 import importlib
 import traceback
 from array import array
 from enum import Enum
 from functools import partial
+from types import ModuleType


TIL about this type

ian-r-rose · 2021-09-24T16:54:51Z

distributed/utils.py

@@ -603,7 +607,7 @@ def key_split(s):
        return "Other"


-def key_split_group(x):
+def key_split_group(x) -> str:


Does it work to narrow x to str | bytes | tuple[Hashable] or similar? I'd love to have written down the real scope of a key.

No, because (for some reason I didn't investigate) unexpected types are accepted and return a hardcoded "Other".

ian-r-rose · 2021-09-24T16:56:53Z

distributed/utils.py

@@ -634,7 +638,7 @@ def key_split_group(x):
    elif typ is bytes:
        return key_split_group(x.decode())
    else:
-        return key_split(x)
+        return "Other"


Note for reviewers: this avoids some redundant checks in key_split, which would itself just return "Other" here.

ian-r-rose · 2021-09-24T17:57:23Z

distributed/preloading.py

    if isinstance(preload, str):
        preload = [preload]
+    if preload_argv and isinstance(preload_argv[0], str):
+        preload_argv = [cast("list[str]", preload_argv)] * len(preload)


I played around with this one as well to try to avoid the cast. Annoying that mypy is unable to narrow based on the above check.

jakirkham · 2021-09-27T22:17:23Z

@jakirkham

So would recommend we leave scheduler.py alone.

Would like to, but can't. Adding # type: ignore at the top of the file breaks everything else. I could not make the --exclude parameter of mypy work either. So the only alternative to changing the actual annotations in scheduler.py is to add a wealth of # type: ignore on its individual lines.

As noted above, this is going to cause issues with Cythonization in the scheduler. If we can't leave that file alone, then unfortunately I don't think we are prepared to adopt MyPy yet.

cc @mrocklin @quasiben (for awareness)

crusaderky · 2021-09-28T16:11:56Z

Please move all conversation specific to scheduler.py to #5348

crusaderky · 2021-09-29T11:58:47Z

distributed/deploy/local.py

@@ -258,12 +256,3 @@ def _repr_html_(self, cluster_status=None):
            cluster_status=cluster_status,
        )
        return super()._repr_html_(cluster_status=cluster_status)
-
-
-clusters_to_close = weakref.WeakSet()


never populated

Good catch! I believe we now handled here

distributed/distributed/deploy/spec.py

Lines 654 to 660 in 43d3866

@atexit.register

def close_clusters():

for cluster in list(SpecCluster._instances):

if cluster.shutdown_on_close:

with suppress(gen.TimeoutError, TimeoutError):

if cluster.status != Status.closed:

cluster.close(timeout=10)

crusaderky · 2021-09-29T12:01:12Z

distributed/diagnostics/tests/test_scheduler_plugin.py

-        pass
+    async with Scheduler(dashboard_address=":0") as s:
+        s.add_plugin(plugin)
+        await plugin.start(s)


#5348 removes the plugins= parameter from Scheduler and Worker. It was exclusively used here.

Maybe use Scheduler.register_scheduler_plugin instead? Just a suggestion, not meant to be a blocking comment

#5348 removes the plugins= parameter from Scheduler and Worker. It was exclusively used here.

There may be user code which uses plugins=

crusaderky · 2021-09-29T12:05:05Z

distributed/protocol/serialize.py

@@ -15,12 +18,10 @@
 from .compression import decompress, maybe_compress
 from .utils import frame_split_size, msgpack_opts, pack_frames_prelude, unpack_frames

-lazy_registrations = {}
-


crusaderky · 2021-09-29T12:09:11Z

This for me is done - @ian-r-rose are there further comments?
Any objections to merging?

jrbourbeau

Thanks @crusaderky! I left several small comments / questions, but overall this looks good to me

Also, just checking but there's an empty distributed/py.typed file being added here, is that intentional?

jrbourbeau · 2021-09-29T20:12:45Z

setup.cfg

+[options]
+zip_safe = False  # https://mypy.readthedocs.io/en/latest/installed_packages.html
+


I'm wondering how this is different from zip_safe=False in setup.py

distributed/setup.py

Line 110 in 43d3866

zip_safe=False,

Ah. I didn't notice it.

jrbourbeau · 2021-09-29T20:28:21Z

distributed/utils.py

@@ -1058,15 +1062,15 @@ def reset_logger_locks():

        # TODO: Use tornado's AnyThreadEventLoopPolicy, instead of class below,
        # once tornado > 6.0.3 is available.
-        if WINDOWS and hasattr(asyncio, "WindowsSelectorEventLoopPolicy"):
+        if WINDOWS:


Just noting for future readers that asyncio.WindowsSelectorEventLoopPolicy was added in Python 3.7 https://docs.python.org/3/whatsnew/3.7.html

jrbourbeau · 2021-09-29T20:32:25Z

distributed/tests/test_preload.py

@@ -75,7 +75,7 @@ def dask_teardown(worker):
    worker.foo = 'teardown'
 """
    with dask.config.set(
-        {"distributed.worker.preload": text, "distributed.nanny.preload": text}
+        {"distributed.worker.preload": [text], "distributed.nanny.preload": [text]}


Why was this change needed?

It violated distributed-schema.yaml. It's purely conincidental that process_preloads() will accept a single string.

jrbourbeau · 2021-09-29T20:36:58Z

distributed/protocol/compression.py

@@ -23,7 +27,13 @@

 from ..utils import ensure_bytes

-compressions = {None: {"compress": identity, "decompress": identity}}
+if TYPE_CHECKING:
+    from typing_extensions import Literal


Should we include typing_extensions in our conda environment file(s)?

No, because we are running everything through pre-commit. typing_extensions shouldn't be in the env files for the same reason flake8, black, isort and mypy aren't.

To clarify: if you look at #5348, typing_extensions is installed there exclusively in the mypy virtualenv by .pre-conmmit-config.yaml.

That makes sense, thanks for clarifying

distributed/preloading.py

jrbourbeau · 2021-09-29T20:54:35Z

distributed/deploy/local.py

@@ -258,12 +256,3 @@ def _repr_html_(self, cluster_status=None):
            cluster_status=cluster_status,
        )
        return super()._repr_html_(cluster_status=cluster_status)
-
-
-clusters_to_close = weakref.WeakSet()


Good catch! I believe we now handled here

distributed/distributed/deploy/spec.py

Lines 654 to 660 in 43d3866

@atexit.register

def close_clusters():

for cluster in list(SpecCluster._instances):

if cluster.shutdown_on_close:

with suppress(gen.TimeoutError, TimeoutError):

if cluster.status != Status.closed:

cluster.close(timeout=10)

jrbourbeau · 2021-09-29T20:55:49Z

distributed/deploy/adaptive.py

+        if self.cluster:
+            return self.cluster.loop
+        else:
+            return IOLoop.current()


Is this another "mypy caught a bug" case?

No. It caught an antipattern where the parent class expected properties that only exist in the children classes.

jrbourbeau · 2021-09-29T21:00:09Z

distributed/compatibility.py

+logging_names: dict[str | int, int | str] = {}
+logging_names.update(logging._levelToName)  # type: ignore
+logging_names.update(logging._nameToLevel)  # type: ignore


We're loosing a copy() here, just wanted to double check that's okay

We're replacing a copy() with an update() for the sake of readability. There's no functional difference.

jrbourbeau · 2021-09-29T21:07:31Z

distributed/comm/ucx.py

+if TYPE_CHECKING:
+    try:
+        import ucp
+        from ucp import create_endpoint as ucx_create_endpoint
+        from ucp import create_listener as ucx_create_listener
+    except ImportError:
+        pass
+else:
+    ucp = None  # type: ignore
+    ucx_create_endpoint = None  # type: ignore
+    ucx_create_listener = None  # type: ignore
+


@pentschev @madsbk @quasiben could one of you review the changes to this module?

Thanks for the ping here @jrbourbeau . We can't import ucp here, we have to do it from init_once, otherwise UCX will be imported before we can set any of its configurations. However, it seems it's only being imported here when TYPE_CHECKING is defined, when is that going to happen in Distributed?

However, it seems it's only being imported here when TYPE_CHECKING is defined, when is that going to happen in Distributed?

Today we don't run any type checking as part of the development process, but after #5348 is merged we'll start running mypy as a pre-commit hook. Given that this ucp import will only happen during type checking, do you think the proposed changes here are okay?

If it's guaranteed that TYPE_CHECKING is False for all runtime situations, then I think this is no problem.

That should be the case. From https://docs.python.org/3/library/typing.html#typing.TYPE_CHECKING:

It is False at runtime

cc @crusaderky for confirmation

In that case, we should be good for now. If there's any reason why TYPE_CHECKING is True at runtime, we'll need to make changes to cover that.

jrbourbeau · 2021-09-29T21:29:19Z

distributed/client.py

-    _shutdown = _close
-


Just noting to future readers that there's a _shutdown method defined a few lines below

ian-r-rose · 2021-09-29T23:48:37Z

distributed/preloading.py

+if TYPE_CHECKING:
+    from .core import Server
+


That's usually the reason for this pattern, yes. However, as far as I can tell it's not actually needed here, since Server doesn't depend on preloading, and everything that imports preloading has already imported Server somewhere.

ian-r-rose · 2021-09-29T23:56:36Z

distributed/active_memory_manager.py

@@ -215,7 +217,7 @@ def _find_recipient(
        candidates -= pending_repl
        if not candidates:
            return None
-        return min(candidates, key=self.workers_memory.get)
+        return min(candidates, key=self.workers_memory.get)  # type: ignore


I cringe a bit to suggest this, but self.workers_memory.__getitem__ passes muster here, since it will raise if the key doesn't exist

ian-r-rose · 2021-09-30T00:03:40Z

distributed/deploy/spec.py

-
-        if not self.new_spec:
-            raise ValueError("To scale by cores= you must specify cores per worker")
+        assert False, "unreachable"


This should be triggered if none of the above keys are in new_spec. Without this assertion, the proper return type of this function would be Optional[int]. Perhaps a more helpful error would be something similar to the ValueError above?

ian-r-rose · 2021-09-30T00:09:07Z

distributed/deploy/tests/test_spec_cluster.py

+    @status.setter
+    def status(self, value):
+        raise NotImplementedError()
+


Was curious about this as well -- it's because this inherits from Worker, where status isn't read-only. Without this it complains that we are swapping a read-write property with a read-only one, which, fair enough.

Co-authored-by: James Bourbeau <jrbourbeau@users.noreply.github.com>

crusaderky · 2021-09-30T11:09:50Z

All review comments have been incorporated

crusaderky · 2021-09-30T11:24:08Z

Also, just checking but there's an empty distributed/py.typed file being added here, is that intentional?

Yes, this file is used by mypy when it validates third party code that imports distributed, and it signals that mypy should fetch type annotations from the library itself (as opposed to look for a separate stubs package).

jrbourbeau · 2021-09-30T14:30:22Z

@jakirkham earlier you expressed some hesitation around adopting mypy due to its potential impact on the Cythonized scheduler (#5328 (comment)). @crusaderky has isolated the changes to scheduler.py to #5348 where he is, I think, taking a conservative approach with modifications. The current plan is to merge this PR (which doesn't touch scheduler.py) and then have someone review the scheduler.py changes in #5348 for performance regressions (you're clearly the most qualified person to do this). Just wanted to check to see if you're okay with this plan before moving forward.

jakirkham · 2021-09-30T19:52:36Z

Yeah I'm ok with that James. Thanks for checking :)

jrbourbeau

Sounds good!

Thanks @crusaderky for your work on this and addressing all the review comments from @ian-r-rose and me

ian-r-rose · 2021-09-30T20:00:41Z

Woo!

crusaderky marked this pull request as draft September 17, 2021 13:03

crusaderky force-pushed the mypy branch 3 times, most recently from b3e14ab to 62ded0f Compare September 18, 2021 18:07

Enable mypy in CI

47d144b

crusaderky force-pushed the mypy branch from 62ded0f to 47d144b Compare September 18, 2021 18:15

Merge branch 'main' into mypy

2b2de7e

crusaderky force-pushed the mypy branch from 9a41fb7 to f64ddb3 Compare September 23, 2021 12:39

crusaderky force-pushed the mypy branch 2 times, most recently from 8bda7e9 to 3de84a1 Compare September 23, 2021 20:21

annotations

a48f68b

crusaderky force-pushed the mypy branch from 4f75db3 to a48f68b Compare September 23, 2021 20:38

jrbourbeau reviewed Sep 23, 2021

View reviewed changes

fix regressions

fca1f7a

crusaderky force-pushed the mypy branch from ab7e71f to fca1f7a Compare September 24, 2021 10:08

addressing

7d2072a

crusaderky added 6 commits September 24, 2021 13:05

scheduler.py

a10c034

Merge branch 'main' into mypy

a7985cb

scheduler.py

c216653

SchedulerPlugins

74091d9

scheduler.py

4b6ad55

Revert changes to scheduler.py and pre-commit hooks

8a0ca70

crusaderky changed the title ~~Enable mypy in CI~~ Enable mypy in CI 1/2 Sep 24, 2021

crusaderky marked this pull request as ready for review September 24, 2021 15:57

crusaderky mentioned this pull request Sep 24, 2021

Enable mypy in CI 2/2 #5348

Merged

ian-r-rose reviewed Sep 24, 2021

View reviewed changes

crusaderky force-pushed the mypy branch from 2728850 to 5802463 Compare September 27, 2021 20:59

code review

ec485f3

crusaderky added 3 commits September 28, 2021 15:20

Merge branch 'main' into mypy

bc67e27

merge from main

4ef6c46

Merge branch 'main' into mypy

831cdae

Merge branch 'main' into mypy

fa8d737

crusaderky commented Sep 29, 2021

View reviewed changes

jrbourbeau reviewed Sep 29, 2021

View reviewed changes

ian-r-rose reviewed Sep 30, 2021

View reviewed changes

crusaderky and others added 4 commits September 30, 2021 11:37

Merge branch 'main' into mypy

86e6cbf

Update distributed/preloading.py

9d82ae6

Co-authored-by: James Bourbeau <jrbourbeau@users.noreply.github.com>

Code review

1fa20a2

Merge remote-tracking branch 'origin/mypy' into mypy

1853f60

crusaderky added 2 commits September 30, 2021 16:56

Merge branch 'main' into mypy

8e9c69f

revert

edcd206

jrbourbeau approved these changes Sep 30, 2021

View reviewed changes

jrbourbeau merged commit 7a3ea4c into dask:main Sep 30, 2021

crusaderky deleted the mypy branch September 30, 2021 20:02

jrbourbeau mentioned this pull request Oct 12, 2021

Add columns argument to read_gbq coiled/dask-bigquery#15

Merged

Illviljan mentioned this pull request Nov 24, 2022

mypy - Remove some ignored packages and modules pydata/xarray#7319

Merged

	remote_magic._clients = {} # type: ignore
	remote_magic._clients: Dict[str, BlockingKernelClient] = {}

	@atexit.register
	def close_clusters():
	for cluster in list(SpecCluster._instances):
	if cluster.shutdown_on_close:
	with suppress(gen.TimeoutError, TimeoutError):
	if cluster.status != Status.closed:
	cluster.close(timeout=10)

		[options]
		zip_safe = False # https://mypy.readthedocs.io/en/latest/installed_packages.html

Enable mypy in CI 1/2 #5328

Enable mypy in CI 1/2 #5328

Conversation

crusaderky commented Sep 17, 2021 • edited Loading

crusaderky commented Sep 23, 2021

jrbourbeau left a comment

Choose a reason for hiding this comment

madsbk commented Sep 24, 2021

crusaderky commented Sep 24, 2021 • edited Loading

jakirkham commented Sep 24, 2021 • edited Loading

ian-r-rose left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

crusaderky Sep 27, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jakirkham commented Sep 27, 2021

crusaderky commented Sep 28, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

crusaderky commented Sep 29, 2021

jrbourbeau left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jrbourbeau Sep 30, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

crusaderky commented Sep 30, 2021

crusaderky commented Sep 30, 2021

jrbourbeau commented Sep 30, 2021

jakirkham commented Sep 30, 2021

jrbourbeau left a comment

Choose a reason for hiding this comment

ian-r-rose commented Sep 30, 2021

crusaderky commented Sep 17, 2021 •

edited

Loading

crusaderky commented Sep 24, 2021 •

edited

Loading

jakirkham commented Sep 24, 2021 •

edited

Loading

crusaderky Sep 27, 2021 •

edited

Loading

jrbourbeau Sep 30, 2021 •

edited

Loading