Skip to content

Fix mypy for utilities.xla_device #8755

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 2 commits into from
Aug 6, 2021
Merged

Fix mypy for utilities.xla_device #8755

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f58f917
Fix mypy for utilities.xla_device
stancld Aug 5, 2021
205b71f
Add explicit type hint for _XLA_AVAILABLE in utilities.imports
stancld Aug 5, 2021
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
1 change: 1 addition & 0 deletions pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,5 +70,6 @@ module = [
"pytorch_lightning.utilities.device_parser",
"pytorch_lightning.utilities.distributed",
"pytorch_lightning.utilities.parsing",
"pytorch_lightning.utilities.xla_device",
]
ignore_errors = "False"
2 changes: 1 addition & 1 deletion pytorch_lightning/utilities/imports.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,7 @@ def _compare_version(package: str, op, version) -> bool:
_TORCHVISION_AVAILABLE = _module_available("torchvision")
_TORCHMETRICS_LOWER_THAN_0_3 = _compare_version("torchmetrics", operator.lt, "0.3.0")
_TORCHMETRICS_GREATER_EQUAL_0_3 = _compare_version("torchmetrics", operator.ge, "0.3.0")
_XLA_AVAILABLE = _module_available("torch_xla")
_XLA_AVAILABLE: bool = _module_available("torch_xla")
Copy link
Contributor

@carmocca carmocca Aug 6, 2021

Choose a reason for hiding this comment

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

Why is this necessary?

_module_available is annotated with -> bool

Copy link
Contributor Author

@stancld stancld Aug 6, 2021

Choose a reason for hiding this comment

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

Unfortunately, mypy is not able to figure that out without this explicit type hint.

Copy link
Collaborator

@SkafteNicki SkafteNicki Aug 6, 2021

Choose a reason for hiding this comment

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

can confirm, we had the same problem in metrics:
https://github.com/PyTorchLightning/metrics/blob/525642df8c3f3f5f93cdd22b6b9cd802a3854634/torchmetrics/utilities/imports.py#L79-L80


from pytorch_lightning.utilities.xla_device import XLADeviceUtils # noqa: E402

Expand Down
9 changes: 5 additions & 4 deletions pytorch_lightning/utilities/xla_device.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
import queue as q
import traceback
from multiprocessing import Process, Queue
from typing import Any, Callable, Union

from pytorch_lightning.utilities.imports import _XLA_AVAILABLE

Expand All @@ -26,7 +27,7 @@
TPU_CHECK_TIMEOUT = 60


def inner_f(queue, func, *args, **kwargs): # pragma: no cover
def inner_f(queue: Queue, func: Callable, *args: Any, **kwargs: Any) -> None: # pragma: no cover
try:
queue.put(func(*args, **kwargs))
# todo: specify the possible exception
Expand All @@ -35,10 +36,10 @@ def inner_f(queue, func, *args, **kwargs): # pragma: no cover
queue.put(None)


def pl_multi_process(func):
def pl_multi_process(func: Callable) -> Callable:
@functools.wraps(func)
def wrapper(*args, **kwargs):
queue = Queue()
def wrapper(*args: Any, **kwargs: Any) -> Union[bool, Any]:
queue: Queue = Queue()
proc = Process(target=inner_f, args=(queue, func, *args), kwargs=kwargs)
proc.start()
proc.join(TPU_CHECK_TIMEOUT)
Expand Down