Skip to content

Defer pandas import on worker in P2P shuffle #5695

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Feb 2, 2022
Merged

Defer pandas import on worker in P2P shuffle #5695

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
338d0be
Defer pandas import on worker in P2P shuffle
gjoseph92 Jan 24, 2022
ce8e9d2
Test unnecessary imports in nanny tests
gjoseph92 Jan 28, 2022
6a6fc64
Update distributed/tests/test_nanny.py
crusaderky Jan 28, 2022
31c33c6
Update distributed/tests/test_nanny.py
crusaderky Jan 28, 2022
9d0cd65
Update distributed/tests/test_nanny.py
crusaderky Jan 28, 2022
a591f9a
Merge remote-tracking branch 'upstream/main' into p2p-shuffle/no-pand…
gjoseph92 Feb 1, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 2 additions & 11 deletions distributed/shuffle/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,16 +1,7 @@
try:
import pandas
except ImportError:
SHUFFLE_AVAILABLE = False
else:
del pandas
SHUFFLE_AVAILABLE = True

from .shuffle import rearrange_by_column_p2p
from .shuffle_extension import ShuffleId, ShuffleMetadata, ShuffleWorkerExtension
from .shuffle import rearrange_by_column_p2p
from .shuffle_extension import ShuffleId, ShuffleMetadata, ShuffleWorkerExtension

__all__ = [
"SHUFFLE_AVAILABLE",
"rearrange_by_column_p2p",
"ShuffleId",
"ShuffleMetadata",
Expand Down
5 changes: 4 additions & 1 deletion distributed/shuffle/shuffle.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
from typing import TYPE_CHECKING

from dask.base import tokenize
from dask.dataframe import DataFrame
from dask.delayed import Delayed, delayed
from dask.highlevelgraph import HighLevelGraph

Expand All @@ -12,6 +11,8 @@
if TYPE_CHECKING:
import pandas as pd

from dask.dataframe import DataFrame


def get_ext() -> ShuffleWorkerExtension:
from distributed import get_worker
Expand Down Expand Up @@ -53,6 +54,8 @@ def rearrange_by_column_p2p(
column: str,
npartitions: int | None = None,
):
from dask.dataframe import DataFrame
Copy link
Collaborator Author

@gjoseph92 gjoseph92 Jan 24, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't love how dask.dataframe.DataFrame is also imported if TYPE_CHECKING. It means I could remove this line here, and the code would still pass mypy and flake8, even though it would fail at runtime. Tests would still fail, but I could see making a similar mistake in the future in some place that didn't have test coverage, and unwittingly introducing a runtime error.

However, because I want to use DataFrame as a type annotation, it does need to be imported for type-checkers.

I don't really think there's a way around this. Just curious if anyone had ideas, particularly @crusaderky or @ian-r-rose

Copy link
Collaborator

@crusaderky crusaderky Jan 25, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yep, this issue is well-known upstream: PyCQA/pyflakes#530
The solution is remembering that linters are a nice and fast way to detect issues but cannot replace unit tests :)


npartitions = npartitions or df.npartitions
token = tokenize(df, column, npartitions)

Expand Down
6 changes: 4 additions & 2 deletions distributed/shuffle/shuffle_extension.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,12 +6,12 @@
from dataclasses import dataclass
from typing import TYPE_CHECKING, NewType

import pandas as pd

from distributed.protocol import to_serialize
from distributed.utils import sync

if TYPE_CHECKING:
import pandas as pd

from distributed.worker import Worker

ShuffleId = NewType("ShuffleId", str)
Expand Down Expand Up @@ -103,6 +103,8 @@ async def add_partition(self, data: pd.DataFrame) -> None:
await asyncio.gather(*tasks)

def get_output_partition(self, i: int) -> pd.DataFrame:
import pandas as pd

assert self.transferred, "`get_output_partition` called before barrier task"

assert self.metadata.worker_for(i) == self.worker.address, (
Expand Down
29 changes: 29 additions & 0 deletions distributed/tests/test_nanny.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import multiprocessing as mp
import os
import random
import sys
from contextlib import suppress
from time import sleep
from unittest import mock
Expand Down Expand Up @@ -610,3 +611,31 @@ async def test_environ_plugin(c, s, a, b):
assert results[a.worker_address] == "123"
assert results[b.worker_address] == "123"
assert results[n.worker_address] == "123"


@pytest.mark.parametrize(
"modname",
[
pytest.param(
"numpy",
marks=pytest.mark.xfail(reason="distributed#5723, distributed#5729"),
),
"scipy",
pytest.param("pandas", marks=pytest.mark.xfail(reason="distributed#5723")),
],
)
@gen_cluster(client=True, Worker=Nanny, nthreads=[("", 1)])
async def test_no_unnecessary_imports_on_worker(c, s, a, modname):
"""
Regression test against accidentally importing unnecessary modules at worker startup.

Importing modules like pandas slows down worker startup, especially if workers are
loading their software environment from NFS or other non-local filesystems.
It also slightly increases memory footprint.
"""

def assert_no_import(dask_worker):
assert modname not in sys.modules

await c.wait_for_workers(1)
await c.run(assert_no_import)
7 changes: 3 additions & 4 deletions distributed/worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@
typename,
)

from . import comm, preloading, profile, shuffle, system, utils
from . import comm, preloading, profile, system, utils
from .batched import BatchedSend
from .comm import Comm, connect, get_address_host
from .comm.addressing import address_from_user_args, parse_address
Expand All @@ -74,6 +74,7 @@
from .protocol import pickle, to_serialize
from .pubsub import PubSubWorkerExtension
from .security import Security
from .shuffle import ShuffleWorkerExtension
from .sizeof import safe_sizeof as sizeof
from .threadpoolexecutor import ThreadPoolExecutor
from .threadpoolexecutor import secede as tpe_secede
Expand Down Expand Up @@ -121,9 +122,7 @@
# Worker.status subsets
RUNNING = {Status.running, Status.paused, Status.closing_gracefully}

DEFAULT_EXTENSIONS: list[type] = [PubSubWorkerExtension]
if shuffle.SHUFFLE_AVAILABLE:
DEFAULT_EXTENSIONS.append(shuffle.ShuffleWorkerExtension)
DEFAULT_EXTENSIONS: list[type] = [PubSubWorkerExtension, ShuffleWorkerExtension]

DEFAULT_METRICS: dict[str, Callable[[Worker], Any]] = {}

Expand Down