feat: dpdisp run (#456)

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

- **New Features**
- Introduced a new `run` command to execute Python scripts with
associated PEP 723 metadata.
- Added documentation on how to use the `run` command in `dpdispatcher`.
- Expanded the glossary in the documentation to include the term "run."

- **Documentation**
  - Corrected a spelling error in the `README.md`.
- Added a new guide `doc/run.md` for running Python scripts using
`dpdispatcher`.

- **Bug Fixes**
  - Ensured the `run` command is included in CLI tests.

- **Chores**
- Updated dependencies in `pyproject.toml` to include `tomli` for Python
versions below 3.11.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

Author: 3 people

README.md

@@ -36,4 +36,4 @@ See [Contributing Guide](CONTRIBUTING.md) to become a contributor! 🤓
 
 ## References
 
-DPDispatcher is derivated from the [DP-GEN](https://github.com/deepmodeling/dpgen) package. To mention DPDispatcher in a scholarly publication, please read Section 3.3 in the [DP-GEN paper](https://doi.org/10.1016/j.cpc.2020.107206).
+DPDispatcher is derived from the [DP-GEN](https://github.com/deepmodeling/dpgen) package. To mention DPDispatcher in a scholarly publication, please read Section 3.3 in the [DP-GEN paper](https://doi.org/10.1016/j.cpc.2020.107206).

doc/index.rst

@@ -22,6 +22,7 @@ DPDispatcher will monitor (poke) until these jobs finish and download the result
    machine
    resources
    task
+   run
    cli
    api/api
 

doc/pep723.rst

@@ -0,0 +1,3 @@
+.. dargs::
+   :module: dpdispatcher.run
+   :func: pep723_args

doc/run.md

@@ -0,0 +1,21 @@
+# Run Python scripts
+
+DPDispatcher can be used to directly run a single Python script:
+
+```sh
+dpdisp run script.py
+```
+
+The script must include [inline script metadata](https://packaging.python.org/en/latest/specifications/inline-script-metadata/) compliant with [PEP 723](https://peps.python.org/pep-0723/).
+An example of the script is shown below.
+
+```{literalinclude} ../examples/dpdisp_run.py
+:language: py
+:linenos:
+```
+
+The PEP 723 metadata entries for `tool.dpdispatcher` are defined as follows:
+
+```{eval-rst}
+.. include:: pep723.rst
+```

dpdispatcher/dpdisp.py

@@ -3,6 +3,7 @@
 from typing import List, Optional
 
 from dpdispatcher.entrypoints.gui import start_dpgui
+from dpdispatcher.entrypoints.run import run
 from dpdispatcher.entrypoints.submission import handle_submission
 
 
@@ -81,6 +82,18 @@ def main_parser() -> argparse.ArgumentParser:
             "to the network on both IPv4 and IPv6 (where available)."
         ),
     )
+    ##########################################
+    # run
+    parser_run = subparsers.add_parser(
+        "run",
+        help="Run a Python script.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser_run.add_argument(
+        "filename",
+        type=str,
+        help="Python script to run. PEP 723 metadata should be contained in this file.",
+    )
     return parser
 
 
@@ -117,6 +130,8 @@ def main():
             port=args.port,
             bind_all=args.bind_all,
         )
+    elif args.command == "run":
+        run(filename=args.filename)
     elif args.command is None:
         pass
     else:

dpdispatcher/entrypoints/run.py

@@ -0,0 +1,9 @@
+"""Run PEP 723 script."""
+
+from dpdispatcher.run import run_pep723
+
+
+def run(*, filename: str):
+    with open(filename) as f:
+        script = f.read()
+    run_pep723(script)

dpdispatcher/run.py

@@ -0,0 +1,172 @@
+import os
+import re
+import sys
+from glob import glob
+from hashlib import sha1
+
+from dpdispatcher.machine import Machine
+from dpdispatcher.submission import Resources, Submission, Task
+
+if sys.version_info >= (3, 11):
+    import tomllib
+else:
+    import tomli as tomllib
+from typing import List, Optional
+
+from dargs import Argument
+
+from dpdispatcher.arginfo import machine_dargs, resources_dargs, task_dargs
+
+REGEX = r"(?m)^# /// (?P<type>[a-zA-Z0-9-]+)$\s(?P<content>(^#(| .*)$\s)+)^# ///$"
+
+
+def read_pep723(script: str) -> Optional[dict]:
+    """Read a PEP 723 script metadata from a script string.
+
+    Parameters
+    ----------
+    script : str
+        Script content.
+
+    Returns
+    -------
+    dict
+        PEP 723 metadata.
+    """
+    name = "script"
+    matches = list(
+        filter(lambda m: m.group("type") == name, re.finditer(REGEX, script))
+    )
+    if len(matches) > 1:
+        # TODO: Add tests for scenarios where multiple script blocks are found
+        raise ValueError(f"Multiple {name} blocks found")
+    elif len(matches) == 1:
+        content = "".join(
+            line[2:] if line.startswith("# ") else line[1:]
+            for line in matches[0].group("content").splitlines(keepends=True)
+        )
+        return tomllib.loads(content)
+    else:
+        # TODO: Add tests for scenarios where no metadata is found
+        return None
+
+
+def pep723_args() -> Argument:
+    """Return the argument parser for PEP 723 metadata."""
+    machine_args = machine_dargs()
+    machine_args.fold_subdoc = True
+    machine_args.doc = "Machine configuration. See related documentation for details."
+    resources_args = resources_dargs(detail_kwargs=False)
+    resources_args.fold_subdoc = True
+    resources_args.doc = (
+        "Resources configuration. See related documentation for details."
+    )
+    task_args = task_dargs()
+    command_arg = task_args["command"]
+    command_arg.doc = (
+        "Python interpreter or launcher. No need to contain the Python script filename."
+    )
+    command_arg.default = "python"
+    command_arg.optional = True
+    task_args["task_work_path"].doc += " Can be a glob pattern."
+    task_args.name = "task_list"
+    task_args.doc = "List of tasks to execute."
+    task_args.repeat = True
+    task_args.dtype = (list,)
+    return Argument(
+        "pep723",
+        dtype=dict,
+        doc="PEP 723 metadata",
+        sub_fields=[
+            Argument(
+                "work_base",
+                dtype=str,
+                optional=True,
+                default="./",
+                doc="Base directory for the work",
+            ),
+            Argument(
+                "forward_common_files",
+                dtype=List[str],
+                optional=True,
+                default=[],
+                doc="Common files to forward to the remote machine",
+            ),
+            Argument(
+                "backward_common_files",
+                dtype=List[str],
+                optional=True,
+                default=[],
+                doc="Common files to backward from the remote machine",
+            ),
+            machine_args,
+            resources_args,
+            task_args,
+        ],
+    )
+
+
+def create_submission(metadata: dict, hash: str) -> Submission:
+    """Create a Submission instance from a PEP 723 metadata.
+
+    Parameters
+    ----------
+    metadata : dict
+        PEP 723 metadata.
+    hash : str
+        Submission hash.
+
+    Returns
+    -------
+    Submission
+        Submission instance.
+    """
+    base = pep723_args()
+    metadata = base.normalize_value(metadata, trim_pattern="_*")
+    base.check_value(metadata, strict=False)
+
+    tasks = []
+    for task in metadata["task_list"]:
+        task = task.copy()
+        task["command"] += f" $REMOTE_ROOT/script_{hash}.py"
+        task_work_path = os.path.join(
+            metadata["machine"]["local_root"],
+            metadata["work_base"],
+            task["task_work_path"],
+        )
+        if os.path.isdir(task_work_path):
+            tasks.append(Task.load_from_dict(task))
+        elif glob(task_work_path):
+            for file in glob(task_work_path):
+                tasks.append(Task.load_from_dict({**task, "task_work_path": file}))
+        # TODO: Add tests for scenarios where the task work path is a glob pattern
+        else:
+            # TODO: Add tests for scenarios where the task work path is not found
+            raise FileNotFoundError(f"Task work path {task_work_path} not found.")
+    return Submission(
+        work_base=metadata["work_base"],
+        forward_common_files=metadata["forward_common_files"],
+        backward_common_files=metadata["backward_common_files"],
+        machine=Machine.load_from_dict(metadata["machine"]),
+        resources=Resources.load_from_dict(metadata["resources"]),
+        task_list=tasks,
+    )
+
+
+def run_pep723(script: str):
+    """Run a PEP 723 script.
+
+    Parameters
+    ----------
+    script : str
+        Script content.
+    """
+    metadata = read_pep723(script)
+    if metadata is None:
+        raise ValueError("No PEP 723 metadata found.")
+    dpdispatcher_metadata = metadata["tool"]["dpdispatcher"]
+    script_hash = sha1(script.encode("utf-8")).hexdigest()
+    submission = create_submission(dpdispatcher_metadata, script_hash)
+    submission.machine.context.write_file(f"script_{script_hash}.py", script)
+    # write script
+    submission.run_submission()

examples/dpdisp_run.py

@@ -0,0 +1,28 @@
+# /// script
+# # dpdispatcher doesn't use `requires-python` and `dependencies`
+# requires-python = ">=3"
+# dependencies = [
+# ]
+# [tool.dpdispatcher]
+# work_base = "./"
+# forward_common_files=[]
+# backward_common_files=[]
+# [tool.dpdispatcher.machine]
+# batch_type = "Shell"
+# local_root = "./"
+# context_type = "LazyLocalContext"
+# [tool.dpdispatcher.resources]
+# number_node = 1
+# cpu_per_node = 1
+# gpu_per_node = 0
+# group_size = 0
+# [[tool.dpdispatcher.task_list]]
+# # no need to contain the script filename
+# command = "python"
+# # can be a glob pattern
+# task_work_path = "./"
+# forward_files = []
+# backward_files = ["log"]
+# ///
+
+print("hello world!")

pyproject.toml

@@ -28,6 +28,7 @@ dependencies = [
     'tqdm>=4.9.0',
     'typing_extensions; python_version < "3.7"',
     'pyyaml',
+    'tomli >= 1.1.0; python_version < "3.11"',
 ]
 requires-python = ">=3.7"
 readme = "README.md"

tests/context.py

@@ -14,6 +14,7 @@
 
 # test backward compatibility with dflow
 from dpdispatcher.dpcloudserver.client import RequestInfoException as _  # noqa: F401
+from dpdispatcher.entrypoints.run import run  # noqa: F401
 from dpdispatcher.entrypoints.submission import handle_submission  # noqa: F401
 from dpdispatcher.machine import Machine  # noqa: F401
 from dpdispatcher.machines.distributed_shell import DistributedShell  # noqa: F401