Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
9 changes: 3 additions & 6 deletions .pre-commit-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -17,14 +17,11 @@ repos:
# Python
- repo: https://github.com/astral-sh/ruff-pre-commit
# Ruff version.
rev: v0.1.3
rev: v0.16.0
hooks:
- id: ruff-check
args: [--fix]
- id: ruff-format
- repo: https://github.com/PyCQA/isort
rev: 5.12.0
hooks:
- id: isort
files: \.py$
# numpydoc
- repo: https://github.com/Carreau/velin
rev: 0.0.12
Expand Down
2 changes: 1 addition & 1 deletion docs/conf.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
# -- Project information -----------------------------------------------------

project = "DPGEN2"
copyright = "2022-%d, DeepModeling" % date.today().year
copyright = f"2022-{date.today().year}, DeepModeling"
author = "DeepModeling"


Expand Down
25 changes: 15 additions & 10 deletions docs/exploration.md
Original file line number Diff line number Diff line change
Expand Up @@ -26,11 +26,11 @@ class StageScheduler(ABC):

@abstractmethod
def plan_next_iteration(
self,
hist_reports : List[ExplorationReport],
report : ExplorationReport,
confs : List[Path],
) -> Tuple[bool, ExplorationTaskGroup, ConfSelector] :
self,
hist_reports: List[ExplorationReport],
report: ExplorationReport,
confs: List[Path],
) -> Tuple[bool, ExplorationTaskGroup, ConfSelector]:
"""
Make the plan for the next iteration of the stage.

Expand Down Expand Up @@ -66,15 +66,15 @@ One may check more details on the [exploratin task group](#exploration-task-grou

DPGEN2 defines a python class `ExplorationTask` to manage all necessry files needed to run a exploration task. It can be used as the example provided in the doc string.
```python
class ExplorationTask():
class ExplorationTask:
"""Define the files needed by an exploration task.

Examples
--------
>>> # this example dumps all files needed by the task.
>>> files = exploration_task.files()
... for file_name, file_content in files.items():
... with open(file_name, 'w') as fp:
... with open(file_name, "w") as fp:
... fp.write(file_content)

"""
Expand All @@ -94,8 +94,8 @@ class ExplorationTaskGroup(Sequence):
...

def add_group(
self,
group : 'ExplorationTaskGroup',
self,
group: "ExplorationTaskGroup",
):
"""Add another group to the group."""
...
Expand Down Expand Up @@ -167,7 +167,9 @@ Then we define the exploration arguments schema in the dpgen2 input configuratio

```python
def xxx_task_group_args():
doc_xxx_task_grp = "XXX exploration tasks. dpgen2 will generate the XXX input script"
doc_xxx_task_grp = (
"XXX exploration tasks. dpgen2 will generate the XXX input script"
)
return Argument(
"task_group",
dict,
Expand All @@ -177,12 +179,14 @@ def xxx_task_group_args():
doc=doc_xxx_task_grp,
)


def xxx_normalize(config):
args = xxx_task_group_args()
config = args.normalize_value(config, trim_pattern="_*")
args.check_value(config, strict=False)
return config


def make_xxx_task_group_from_config(config):
config = xxx_normalize(config)
tgroup = XXXTaskGroup(**config)
Expand All @@ -197,6 +201,7 @@ def xxx_args():
# Argument(...),
]


def variant_explore():
# ...
doc_xxx = "The exploration by XXX"
Expand Down
169 changes: 91 additions & 78 deletions docs/operator.md
Original file line number Diff line number Diff line change
Expand Up @@ -92,31 +92,34 @@ The `prep_train` prepares a list of paths, each of which contains all necessary

The `run_train` slices the list of paths, and assign each item in the list to a DeePMD-kit task. The task is executed by `run_train_op`. This is a very nice feature of `dflow`, because the developer only needs to implement how one DeePMD-kit task is executed, and then all the items in the task list will be executed [in parallel](https://github.com/dptech-corp/dflow/blob/master/README.md#315-produce-parallel-steps-using-loop). See the following code to see how it works
```python
run_train = Step(
'run-train',
template=PythonOPTemplate(
run_train_op,
image=run_train_image,
slices = Slices(
"int('{{item}}')",
input_parameter = ["task_name"],
input_artifact = ["task_path", "init_model"],
output_artifact = ["model", "lcurve", "log", "script"],
),
run_train = Step(
"run-train",
template=PythonOPTemplate(
run_train_op,
image=run_train_image,
slices=Slices(
"int('{{item}}')",
input_parameter=["task_name"],
input_artifact=["task_path", "init_model"],
output_artifact=["model", "lcurve", "log", "script"],
),
parameters={
"config" : train_steps.inputs.parameters["train_config"],
"task_name" : prep_train.outputs.parameters["task_names"],
},
artifacts={
'task_path' : prep_train.outputs.artifacts['task_paths'],
"init_model" : train_steps.inputs.artifacts['init_models'],
"init_data": train_steps.inputs.artifacts['init_data'],
"iter_data": train_steps.inputs.artifacts['iter_data'],
},
with_sequence=argo_sequence(argo_len(prep_train.outputs.parameters["task_names"]), format=train_index_pattern),
key = step_keys['run-train'],
)
),
parameters={
"config": train_steps.inputs.parameters["train_config"],
"task_name": prep_train.outputs.parameters["task_names"],
},
artifacts={
"task_path": prep_train.outputs.artifacts["task_paths"],
"init_model": train_steps.inputs.artifacts["init_models"],
"init_data": train_steps.inputs.artifacts["init_data"],
"iter_data": train_steps.inputs.artifacts["iter_data"],
},
with_sequence=argo_sequence(
argo_len(prep_train.outputs.parameters["task_names"]),
format=train_index_pattern,
),
key=step_keys["run-train"],
)
```
The input parameter `"task_names"` and artifacts `"task_paths"` and `"init_model"` are sliced and supplied to each DeePMD-kit task. The output artifacts of the tasks (`"model"`, `"lcurve"`, `"log"` and `"script"`) are stacked in the same order as the input lists. These lists are assigned as the outputs of `train_steps` by
```python
Expand Down Expand Up @@ -144,74 +147,84 @@ from dflow.python import (
OPIOSign,
Artifact,
)


class RunDPTrain(OP):
@classmethod
def get_input_sign(cls):
return OPIOSign({
"config" : dict,
"task_name" : str,
"task_path" : Artifact(Path),
"init_model" : Artifact(Path),
"init_data" : Artifact(List[Path]),
"iter_data" : Artifact(List[Path]),
})
return OPIOSign(
{
"config": dict,
"task_name": str,
"task_path": Artifact(Path),
"init_model": Artifact(Path),
"init_data": Artifact(List[Path]),
"iter_data": Artifact(List[Path]),
}
)

@classmethod
def get_output_sign(cls):
return OPIOSign({
"script" : Artifact(Path),
"model" : Artifact(Path),
"lcurve" : Artifact(Path),
"log" : Artifact(Path),
})
return OPIOSign(
{
"script": Artifact(Path),
"model": Artifact(Path),
"lcurve": Artifact(Path),
"log": Artifact(Path),
}
)
```

All items not defined as `Artifact` are treated as parameters of the `OP`. The concept of parameter and artifact are explained in the [dflow document](https://github.com/dptech-corp/dflow/blob/master/README.md#Parametersandartifacts). To be short, the artifacts can be `pathlib.Path` or a list of `pathlib.Path`. The artifacts are passed by the file system. Other data structures are treated as parameters, they are passed as variables encoded in `str`. Therefore, a large amout of information should be stored in artifacts, otherwise they can be considered as parameters.

The operation of the `OP` is implemented in method `execute`, and are run in docker containers. Again taking the `execute` method of `RunDPTrain` as an example

```python
@OP.exec_sign_check
def execute(
self,
ip : OPIO,
) -> OPIO:
...
task_name = ip['task_name']
task_path = ip['task_path']
init_model = ip['init_model']
init_data = ip['init_data']
iter_data = ip['iter_data']
...
work_dir = Path(task_name)
...
# here copy all files in task_path to work_dir
...
with set_directory(work_dir):
fplog = open('train.log', 'w')
def clean_before_quit():
fplog.close()
# train model
command = ['dp', 'train', train_script_name]
ret, out, err = run_command(command)
if ret != 0:
clean_before_quit()
raise FatalError('dp train failed')
fplog.write(out)
# freeze model
ret, out, err = run_command(['dp', 'freeze', '-o', 'frozen_model.pb'])
if ret != 0:
clean_before_quit()
raise FatalError('dp freeze failed')
fplog.write(out)
@OP.exec_sign_check
def execute(
self,
ip: OPIO,
) -> OPIO:
...
task_name = ip["task_name"]
task_path = ip["task_path"]
init_model = ip["init_model"]
init_data = ip["init_data"]
iter_data = ip["iter_data"]
...
work_dir = Path(task_name)
...
# here copy all files in task_path to work_dir
...
with set_directory(work_dir):
fplog = open("train.log", "w")

def clean_before_quit():
fplog.close()

# train model
command = ["dp", "train", train_script_name]
ret, out, err = run_command(command)
if ret != 0:
clean_before_quit()

return OPIO({
"script" : work_dir / train_script_name,
"model" : work_dir / "frozen_model.pb",
"lcurve" : work_dir / "lcurve.out",
"log" : work_dir / "train.log",
})
raise FatalError("dp train failed")
fplog.write(out)
# freeze model
ret, out, err = run_command(["dp", "freeze", "-o", "frozen_model.pb"])
if ret != 0:
clean_before_quit()
raise FatalError("dp freeze failed")
fplog.write(out)
clean_before_quit()

return OPIO(
{
"script": work_dir / train_script_name,
"model": work_dir / "frozen_model.pb",
"lcurve": work_dir / "lcurve.out",
"log": work_dir / "train.log",
}
)
```

The inputs and outputs variables are recorded in data structure `dflow.python.OPIO`, which is initialized by a Python dict. The keys in the input/output `dict`, and the types of the input/output variables will be checked against their signatures by decorator `OP.exec_sign_check`. If any key or type does not match, an exception will be raised.
Expand Down
Loading
Loading