From d5b6a176c19aba31dc21de7d01e62d48b04af9f1 Mon Sep 17 00:00:00 2001 From: Xinzijian Liu Date: Tue, 25 Jun 2024 12:49:04 +0800 Subject: [PATCH] Add train optional files (#236) ## Summary by CodeRabbit - **New Features** - Introduced the ability to specify optional files for training through new parameters in various functions and scripts. - Enhanced the concurrent learning operation by allowing optional training files to be included. - **Improvements** - Improved flexibility in training workflows by supporting optional files in multiple components. --------- Signed-off-by: zjgemi Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- dpgen2/entrypoint/args.py | 16 ++++++++++++++++ dpgen2/entrypoint/submit.py | 4 ++++ dpgen2/op/run_dp_train.py | 5 +++++ dpgen2/superop/prep_run_dp_train.py | 6 ++++++ 4 files changed, 31 insertions(+) diff --git a/dpgen2/entrypoint/args.py b/dpgen2/entrypoint/args.py index 9959af0d..fc7decd8 100644 --- a/dpgen2/entrypoint/args.py +++ b/dpgen2/entrypoint/args.py @@ -46,6 +46,7 @@ def dp_dist_train_args(): doc_template_script = "File names of the template training script. It can be a `List[str]`, the length of which is the same as `numb_models`. Each template script in the list is used to train a model. Can be a `str`, the models share the same template training script. " dock_student_model_path = "The path of student model" doc_student_model_uri = "The URI of student model" + doc_optional_files = "Optional files for training" return [ Argument( @@ -67,6 +68,13 @@ def dp_dist_train_args(): default=None, doc=doc_student_model_uri, ), + Argument( + "optional_files", + list, + optional=True, + default=None, + doc=doc_optional_files, + ), ] @@ -76,6 +84,7 @@ def dp_train_args(): doc_template_script = "File names of the template training script. It can be a `List[str]`, the length of which is the same as `numb_models`. Each template script in the list is used to train a model. Can be a `str`, the models share the same template training script. " doc_init_models_paths = "the paths to initial models" doc_init_models_uri = "The URI of initial models" + doc_optional_files = "Optional files for training" return [ Argument( @@ -105,6 +114,13 @@ def dp_train_args(): default=None, doc=doc_init_models_uri, ), + Argument( + "optional_files", + list, + optional=True, + default=None, + doc=doc_optional_files, + ), ] diff --git a/dpgen2/entrypoint/submit.py b/dpgen2/entrypoint/submit.py index 6f966247..7e1b25ad 100644 --- a/dpgen2/entrypoint/submit.py +++ b/dpgen2/entrypoint/submit.py @@ -155,6 +155,7 @@ def make_concurrent_learning_op( cl_step_config: dict = default_config, upload_python_packages: Optional[List[os.PathLike]] = None, valid_data: Optional[S3Artifact] = None, + train_optional_files: Optional[List[str]] = None, ): if train_style in ("dp", "dp-dist"): prep_run_train_op = PrepRunDPTrain( @@ -165,6 +166,7 @@ def make_concurrent_learning_op( run_config=run_train_config, upload_python_packages=upload_python_packages, valid_data=valid_data, + optional_files=train_optional_files, ) else: raise RuntimeError(f"unknown train_style {train_style}") @@ -488,6 +490,7 @@ def workflow_concurrent_learning( collect_data_config = config["step_configs"]["collect_data_config"] cl_step_config = config["step_configs"]["cl_step_config"] upload_python_packages = config.get("upload_python_packages", None) + train_optional_files = config["train"].get("optional_files", None) if train_style == "dp": init_models_paths = config["train"].get("init_models_paths", None) @@ -537,6 +540,7 @@ def workflow_concurrent_learning( cl_step_config=cl_step_config, upload_python_packages=upload_python_packages, valid_data=valid_data, + train_optional_files=train_optional_files, ) scheduler = make_naive_exploration_scheduler(config) diff --git a/dpgen2/op/run_dp_train.py b/dpgen2/op/run_dp_train.py index 183e68f5..e1d7391a 100644 --- a/dpgen2/op/run_dp_train.py +++ b/dpgen2/op/run_dp_train.py @@ -194,6 +194,7 @@ def get_input_sign(cls): "init_data": Artifact(NestedDict[Path]), "iter_data": Artifact(List[Path]), "valid_data": Artifact(List[Path], optional=True), + "optional_files": Artifact(List[Path], optional=True), } ) @@ -331,6 +332,10 @@ def clean_before_quit(): with open(train_script_name, "w") as fp: json.dump(train_dict, fp, indent=4) + if ip["optional_files"] is not None: + for f in ip["optional_files"]: + Path(f.name).symlink_to(f) + # train model command = _make_train_command( dp_command, diff --git a/dpgen2/superop/prep_run_dp_train.py b/dpgen2/superop/prep_run_dp_train.py index a244c689..a2a42cdc 100644 --- a/dpgen2/superop/prep_run_dp_train.py +++ b/dpgen2/superop/prep_run_dp_train.py @@ -135,6 +135,7 @@ def __init__( upload_python_packages: Optional[List[os.PathLike]] = None, finetune: bool = False, valid_data: Optional[S3Artifact] = None, + optional_files: Optional[List[str]] = None, ): self._input_parameters = { "block_id": InputParameter(type=str, value=""), @@ -196,6 +197,7 @@ def __init__( upload_python_packages=upload_python_packages, finetune=finetune, valid_data=valid_data, + optional_files=optional_files, ) @property @@ -230,6 +232,7 @@ def _prep_run_dp_train( upload_python_packages: Optional[List[os.PathLike]] = None, finetune: bool = False, valid_data: Optional[S3Artifact] = None, + optional_files: Optional[List[str]] = None, ): prep_config = deepcopy(prep_config) run_config = deepcopy(run_config) @@ -285,6 +288,9 @@ def _prep_run_dp_train( "init_data": train_steps.inputs.artifacts["init_data"], "iter_data": train_steps.inputs.artifacts["iter_data"], "valid_data": valid_data, + "optional_files": upload_artifact(optional_files) + if optional_files is not None + else None, }, with_sequence=argo_sequence( argo_len(prep_train.outputs.parameters["task_names"]),