Add `AcceleratorRegistry` #12180

kaushikb11 · 2022-03-02T12:02:35Z

What does this PR do?

Follow up to #12030

accelerator_name = "custom_accelerator"

class TestAccelerator(Accelerator):
    @staticmethod
    def parse_devices(devices):
        return devices

    @staticmethod
    def get_parallel_devices(devices):
        return ["foo"] * devices

    @staticmethod
    def auto_device_count():
        return 3

    @staticmethod
    def is_available():
        return True

AcceleratorRegistry.register(
    accelerator_name, CustomAccelerator, description=accelerator_description, param1="abc", param2=123
)

assert accelerator_name in AcceleratorRegistry
assert isinstance(AcceleratorRegistry.get(accelerator_name), TestAccelerator)

trainer = Trainer(accelerator="custom_accelerator", devices="auto")
assert isinstance(trainer.accelerator, TestAccelerator)
assert trainer._accelerator_connector.parallel_devices == ["foo"] * 3

Does your PR introduce any breaking changes? If yes, please list them.

Before submitting

Was this discussed/approved via a GitHub issue? (not for typos and docs)
Did you read the contributor guideline, Pull Request section?
Did you make sure your PR does only one thing, instead of bundling different changes together?
Did you make sure to update the documentation with your changes? (if necessary)
Did you write any new necessary tests? (not for typos and docs)
Did you verify new and existing tests pass locally with your changes?
Did you list all the breaking changes introduced by this pull request?
Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:

Is this pull request ready for review? (if not, please submit in draft mode)
Check that all items from Before submitting are resolved
Make sure the title is self-explanatory and the description concisely explains the PR
Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

Make sure you had fun coding 🙃

cc @Borda @akihironitta @rohitgr7 @justusschock

pytorch_lightning/accelerators/registry.py

tests/accelerators/test_accelerator_registry.py

four4fish

I like this! Provides more flexibilities for customized accelerator and it's consistent with Strategy behavior.

Do you think we can remove AcceleratorType after this PR?

pytorch_lightning/accelerators/registry.py

pytorch_lightning/accelerators/accelerator.py

pytorch_lightning/accelerators/registry.py

pytorch_lightning/accelerators/accelerator.py

pytorch_lightning/accelerators/registry.py

pytorch_lightning/trainer/connectors/accelerator_connector.py

Co-authored-by: ananthsub <[email protected]>

pytorch_lightning/trainer/connectors/accelerator_connector.py

pytorch_lightning/accelerators/registry.py

codecov · 2022-03-24T17:49:12Z

Codecov Report

Merging #12180 (096afb9) into master (c099c8b) will decrease coverage by 0%.
The diff coverage is 89%.

@@           Coverage Diff           @@
##           master   #12180   +/-   ##
=======================================
- Coverage      88%      88%   -0%     
=======================================
  Files         207      209    +2     
  Lines       17704    17747   +43     
=======================================
+ Hits        15524    15561   +37     
- Misses       2180     2186    +6

carmocca · 2022-03-24T19:24:51Z

pytorch_lightning/accelerators/gpu.py

+        accelerator_registry.register(
+            "gpu",
+            cls,
+            description=f"{cls.__class__.__name__}",


__class__.__name__ is already a string so the fstring is redundant.

carmocca · 2022-03-25T17:13:10Z

pytorch_lightning/accelerators/registry.py

+        @AcceleratorRegistry.register("sota", description="Custom sota accelerator", a=1, b=True)
+        class SOTAAccelerator(Accelerator):
+            def __init__(self, a, b):
+                ...


This doesn't work:

Traceback (most recent call last): File "/Users/carmocca/git/pytorch-lightning/thing.py", line 5, in <module> class SOTAAccelerator: TypeError: do_register() missing 1 required positional argument: 'accelerator'

carmocca · 2022-03-25T17:13:56Z

pytorch_lightning/accelerators/accelerator.py

@@ -75,7 +75,6 @@ def auto_device_count() -> int:
    def is_available() -> bool:
        """Detect if the hardware is available."""

-    @staticmethod
-    @abstractmethod
-    def name() -> str:


I don't get the reasoning for removing the name.

If we kept it, the registry could use it to automatically define the name for the class

Related discussion #12180 (comment)

carmocca · 2022-03-25T17:15:00Z

pytorch_lightning/utilities/registry.py

+from typing import Any
+
+
+def _is_register_method_overridden(mod: type, base_cls: Any, method: str) -> bool:


Why not use is_overridden?

carmocca · 2022-03-25T17:15:44Z

pytorch_lightning/accelerators/registry.py

+    def register(
+        self,
+        name: str,
+        accelerator: Optional[Callable] = None,


The type should be Optional[Type], not Callable

carmocca · 2022-03-25T17:20:56Z

pytorch_lightning/accelerators/__init__.py

 from pytorch_lightning.accelerators.tpu import TPUAccelerator  # noqa: F401
+
+ACCELERATORS_BASE_MODULE = "pytorch_lightning.accelerators"


Why expose this string?

If it was to avoid registering these by default, it would not work because to change the path you'd need to import the variable

And at import time, they would get registered anyways.

carmocca · 2022-03-25T17:21:27Z

pytorch_lightning/accelerators/cpu.py

-        """Name of the Accelerator."""
-        return "cpu"
+    @classmethod
+    def register_accelerators(cls, accelerator_registry: Dict) -> None:


The argument should be typed as _AcceleratorRegistry, not Dict

awaelchli · 2022-03-27T15:37:05Z

pytorch_lightning/accelerators/accelerator.py

-    def name() -> str:
-        """Name of the Accelerator."""
+    @classmethod
+    def register_accelerators(cls, accelerator_registry: Dict) -> None:


Typing is wrong. should be the actual accelerator registry

awaelchli · 2022-03-27T15:38:40Z

pytorch_lightning/accelerators/registry.py

+    an accelerator, e.g., "gpu". It also returns Optional description and
+    parameters to initialize the Accelerator, which were defined during the
+    registration.


... an optional description ...

awaelchli · 2022-03-27T15:39:25Z

pytorch_lightning/accelerators/registry.py

+    registration.
+
+    The motivation for having a AcceleratorRegistry is to make it convenient
+    for the Users to try different accelerators by passing mapped aliases


users capitalized

kaushikb11 self-assigned this Mar 2, 2022

kaushikb11 added the accelerator label Mar 2, 2022

kaushikb11 added this to the 1.6 milestone Mar 2, 2022

kaushikb11 marked this pull request as ready for review March 2, 2022 12:29

kaushikb11 requested review from tchaton, SeanNaren, carmocca, Borda, williamFalcon, awaelchli, justusschock and rohitgr7 as code owners March 2, 2022 12:29

kaushikb11 changed the title ~~Add AcceleratorRegistry~~ Add AcceleratorRegistry Mar 2, 2022

rohitgr7 reviewed Mar 2, 2022

View reviewed changes

pytorch_lightning/accelerators/registry.py Outdated Show resolved Hide resolved

rohitgr7 reviewed Mar 2, 2022

View reviewed changes

tests/accelerators/test_accelerator_registry.py Outdated Show resolved Hide resolved

rohitgr7 mentioned this pull request Mar 2, 2022

add accelerator.is_available() check #12104

Merged

12 tasks

ananthsub reviewed Mar 2, 2022

View reviewed changes

tests/accelerators/test_accelerator_registry.py Outdated Show resolved Hide resolved

four4fish approved these changes Mar 3, 2022

View reviewed changes

pytorch_lightning/accelerators/registry.py Show resolved Hide resolved

mergify bot added the has conflicts label Mar 9, 2022

mergify bot removed the has conflicts label Mar 21, 2022

kaushikb11 commented Mar 22, 2022

View reviewed changes

pytorch_lightning/accelerators/accelerator.py Outdated Show resolved Hide resolved

kaushikb11 commented Mar 22, 2022

View reviewed changes

pytorch_lightning/accelerators/registry.py Outdated Show resolved Hide resolved

ananthsub reviewed Mar 22, 2022

View reviewed changes

pytorch_lightning/accelerators/accelerator.py Outdated Show resolved Hide resolved

ananthsub reviewed Mar 23, 2022

View reviewed changes

justusschock approved these changes Mar 23, 2022

View reviewed changes

pytorch_lightning/accelerators/registry.py Show resolved Hide resolved

ananthsub approved these changes Mar 23, 2022

View reviewed changes

mergify bot added the ready PRs ready to be merged label Mar 23, 2022

Borda approved these changes Mar 24, 2022

View reviewed changes

pytorch_lightning/trainer/connectors/accelerator_connector.py Show resolved Hide resolved

kaushikb11 mentioned this pull request Mar 24, 2022

Add support for Habana accelerator (HPU) #11808

Merged

12 tasks

kaushikb11 enabled auto-merge (squash) March 24, 2022 15:22

kaushikb11 and others added 12 commits March 24, 2022 22:19

Fix tests

ef188a1

Update exception message

eb22e01

Update test

cdb0572

Remove NotImplementedError

c8ea7f2

Update test

8c96d9c

Update pytorch_lightning/accelerators/registry.py

d18091d

Co-authored-by: ananthsub <[email protected]>

Update pytorch_lightning/accelerators/registry.py

a5925d9

Co-authored-by: ananthsub <[email protected]>

Update pytorch_lightning/accelerators/registry.py

75884c6

Co-authored-by: ananthsub <[email protected]>

Update pytorch_lightning/accelerators/registry.py

197dbaa

Co-authored-by: ananthsub <[email protected]>

Update pytorch_lightning/accelerators/registry.py

b0e1120

Co-authored-by: ananthsub <[email protected]>

Fix tests

d1de49e

Fix mypy

096afb9

kaushikb11 force-pushed the add/acc_registry branch from 68f1df5 to 096afb9 Compare March 24, 2022 16:50

ananthsub reviewed Mar 24, 2022

View reviewed changes

pytorch_lightning/trainer/connectors/accelerator_connector.py Show resolved Hide resolved

kaushikb11 marked this pull request as draft March 24, 2022 17:14

auto-merge was automatically disabled March 24, 2022 17:14
Pull request was converted to draft

kaushikb11 marked this pull request as ready for review March 24, 2022 17:14

kaushikb11 commented Mar 24, 2022

View reviewed changes

pytorch_lightning/accelerators/registry.py Show resolved Hide resolved

kaushikb11 force-pushed the add/acc_registry branch from 32dc47f to 096afb9 Compare March 24, 2022 17:31

kaushikb11 enabled auto-merge (squash) March 24, 2022 17:31

kaushikb11 merged commit dcc973e into master Mar 24, 2022

kaushikb11 deleted the add/acc_registry branch March 24, 2022 18:29

carmocca reviewed Mar 25, 2022

View reviewed changes

carmocca mentioned this pull request Mar 25, 2022

[WIP] Accelerator registry follow-up #12461

Closed

12 tasks

kaushikb11 mentioned this pull request Mar 26, 2022

Remove AcceleratorConnector.tpu_cores and Deprecate Trainer.tpu_cores #12437

Merged

12 tasks

awaelchli reviewed Mar 27, 2022

View reviewed changes

awaelchli mentioned this pull request Apr 10, 2022

Choosing an invalid accelerator does not raise an error #12698

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `AcceleratorRegistry` #12180

Add `AcceleratorRegistry` #12180

kaushikb11 commented Mar 2, 2022 •

edited by github-actions bot

Loading

four4fish left a comment

codecov bot commented Mar 24, 2022

carmocca Mar 24, 2022

carmocca Mar 25, 2022

carmocca Mar 25, 2022

kaushikb11 Mar 28, 2022

carmocca Mar 25, 2022

carmocca Mar 25, 2022

carmocca Mar 25, 2022

carmocca Mar 25, 2022

awaelchli Mar 27, 2022

awaelchli Mar 27, 2022

awaelchli Mar 27, 2022

		from typing import Any


		def _is_register_method_overridden(mod: type, base_cls: Any, method: str) -> bool:

		from pytorch_lightning.accelerators.tpu import TPUAccelerator # noqa: F401

		ACCELERATORS_BASE_MODULE = "pytorch_lightning.accelerators"

Add AcceleratorRegistry #12180

Add AcceleratorRegistry #12180

Conversation

kaushikb11 commented Mar 2, 2022 • edited by github-actions bot Loading

What does this PR do?

Does your PR introduce any breaking changes? If yes, please list them.

Before submitting

PR review

Did you have fun?

four4fish left a comment

Choose a reason for hiding this comment

codecov bot commented Mar 24, 2022

Codecov Report

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

Add `AcceleratorRegistry` #12180

Add `AcceleratorRegistry` #12180

kaushikb11 commented Mar 2, 2022 •

edited by github-actions bot

Loading