Skip to content

Commit

Permalink
feat: add plugin entry point for PT (deepmodeling#3965)
Browse files Browse the repository at this point in the history
Currently, we have it for TF, for the backend, but not for PT.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **Documentation**
- Updated guidelines for packaging new codes into a Python package for
DeePMD-kit, including changes to use `pyproject.toml` instead of
`setup.py`.

- **Refactor**
- Improved entry point loading by centralizing the logic into a new
utility function, enhancing code maintainability and readability.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
  • Loading branch information
njzjz authored and Mathieu Taillefumier committed Sep 18, 2024
1 parent d15ad5e commit f2a9655
Show file tree
Hide file tree
Showing 6 changed files with 60 additions and 30 deletions.
13 changes: 5 additions & 8 deletions deepmd/backend/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,12 +7,15 @@
# copy from dpdata
from importlib import (
import_module,
metadata,
)
from pathlib import (
Path,
)

from deepmd.utils.entry_point import (
load_entry_point,
)

PACKAGE_BASE = "deepmd.backend"
NOT_LOADABLE = ("__init__.py",)

Expand All @@ -21,10 +24,4 @@
module_name = f".{module_file.stem}"
import_module(module_name, PACKAGE_BASE)

# https://setuptools.readthedocs.io/en/latest/userguide/entry_point.html
try:
eps = metadata.entry_points(group="deepmd.backend")
except TypeError:
eps = metadata.entry_points().get("deepmd.backend", [])
for ep in eps:
plugin = ep.load()
load_entry_point("deepmd.backend")
5 changes: 5 additions & 0 deletions deepmd/pt/model/__init__.py
Original file line number Diff line number Diff line change
@@ -1 +1,6 @@
# SPDX-License-Identifier: LGPL-3.0-or-later
from deepmd.utils.entry_point import (
load_entry_point,
)

load_entry_point("deepmd.pt")
17 changes: 4 additions & 13 deletions deepmd/tf/__init__.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,10 @@
# SPDX-License-Identifier: LGPL-3.0-or-later
"""Root of the deepmd package, exposes all public classes and submodules."""

try:
from importlib import (
metadata,
)
except ImportError: # for Python<3.8
import importlib_metadata as metadata

import deepmd.tf.utils.network as network
from deepmd.utils.entry_point import (
load_entry_point,
)

from . import (
cluster,
Expand Down Expand Up @@ -39,12 +35,7 @@
)

# load third-party plugins
try:
eps = metadata.entry_points(group="deepmd")
except TypeError:
eps = metadata.entry_points().get("deepmd", [])
for ep in eps:
ep.load()
load_entry_point("deepmd")

__all__ = [
"__version__",
Expand Down
25 changes: 25 additions & 0 deletions deepmd/utils/entry_point.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
# SPDX-License-Identifier: LGPL-3.0-or-later
from importlib import (
metadata,
)


def load_entry_point(group: str) -> list:
"""Load entry points from a group.
Parameters
----------
group : str
The group name.
Returns
-------
list
A list of loaded entry points.
"""
# https://setuptools.readthedocs.io/en/latest/userguide/entry_point.html
try:
eps = metadata.entry_points(group=group)
except TypeError:
eps = metadata.entry_points().get(group, [])
return [ep.load() for ep in eps]
16 changes: 16 additions & 0 deletions doc/development/create-a-model-pt.md
Original file line number Diff line number Diff line change
Expand Up @@ -168,6 +168,22 @@ allows one to use your new descriptor as below:

The arguments here should be consistent with the class arguments of your new component.

## Package new codes

You may package new codes into a new Python package if you don't want to contribute it to the main DeePMD-kit repository.
It's crucial to add your new component to `project.entry-points."deepmd.pt"` in `pyproject.toml`:

```toml
[project.entry-points."deepmd.pt"]
some_descrpt = "deepmd_some_descrtpt:SomeDescript"
```

where `deepmd_some_descrtpt` is the module of your codes. It is equivalent to `from deepmd_some_descrtpt import SomeDescript`.

If you place `SomeDescript` and `descrpt_some_args` into different modules, you are also expected to add `descrpt_some_args` to `entry_points`.

After you install your new package, you can now use `dp train` to run your new model.

## Unit tests

### Universal tests
Expand Down
14 changes: 5 additions & 9 deletions doc/development/create-a-model-tf.md
Original file line number Diff line number Diff line change
Expand Up @@ -58,16 +58,12 @@ The arguments here should be consistent with the class arguments of your new com

## Package new codes

You may use `setuptools` to package new codes into a new Python package. It's crucial to add your new component to `entry_points['deepmd']` in `setup.py`:
You may package new codes into a new Python package if you don't want to contribute it to the main DeePMD-kit repository.
It's crucial to add your new component to `project.entry-points."deepmd"` in `pyproject.toml`:

```py
entry_points = (
{
"deepmd": [
"some_descrpt=deepmd_some_descrtpt:SomeDescript",
],
},
)
```toml
[project.entry-points."deepmd"]
some_descrpt = "deepmd_some_descrtpt:SomeDescript"
```

where `deepmd_some_descrtpt` is the module of your codes. It is equivalent to `from deepmd_some_descrtpt import SomeDescript`.
Expand Down

0 comments on commit f2a9655

Please sign in to comment.