fix: support set output before OP returns

Signed-off-by: zjgemi <[email protected]>

Author: zjgemi

CHANGELOG.md

@@ -1,5 +1,45 @@
 # CHANGELOG
 
+## [1.8.98](https://github.com/deepmodeling/dflow/compare/v1.8.97...v1.8.98) (2024-10-31)
+
+
+### Bug Fixes
+
+* support specify hooks after defining step ([3c016b7](https://github.com/deepmodeling/dflow/commit/3c016b7d42d1c547cce20888959339ae4d019735))
+
+## [1.8.97](https://github.com/deepmodeling/dflow/compare/v1.8.96...v1.8.97) (2024-10-25)
+
+
+### Bug Fixes
+
+* add raise_for_group to Slices ([a3e7ee1](https://github.com/deepmodeling/dflow/commit/a3e7ee11413eb9c0bebed7d59ce84efe52b7b196))
+* add sort_by_generation to query_step ([a3e7ee1](https://github.com/deepmodeling/dflow/commit/a3e7ee11413eb9c0bebed7d59ce84efe52b7b196))
+* hooks of task ([67b5876](https://github.com/deepmodeling/dflow/commit/67b58766845dd8d98443e89a5b0a56558c21291d))
+* support get sub steps of a specific step ([7446003](https://github.com/deepmodeling/dflow/commit/74460034816a7a2ba3c7af44ce9581cda8c923f2))
+
+## [1.8.96](https://github.com/deepmodeling/dflow/compare/v1.8.95...v1.8.96) (2024-10-23)
+
+
+### Bug Fixes
+
+* add lifecycle hooks to step/task ([9030946](https://github.com/deepmodeling/dflow/commit/9030946e0f0edcffe5e806197bba4d3afa74b6e3))
+* add onExit hook ([acbd12d](https://github.com/deepmodeling/dflow/commit/acbd12d23e083d739290a57d04b241b522b9e063))
+
+## [1.8.95](https://github.com/deepmodeling/dflow/compare/v1.8.94...v1.8.95) (2024-10-18)
+
+
+### Bug Fixes
+
+* pass None for HDF5Datasets ([1310554](https://github.com/deepmodeling/dflow/commit/1310554ead008acee1754112ef6032c9a062dfd9))
+
+## [1.8.94](https://github.com/deepmodeling/dflow/compare/v1.8.93...v1.8.94) (2024-10-17)
+
+
+### Bug Fixes
+
+* HDF5Datasets with grouped slices ([1bfae67](https://github.com/deepmodeling/dflow/commit/1bfae67ab6533c9dddf8066a40919c56bb7c53b5))
+* None in HDF5Datasets ([1bfae67](https://github.com/deepmodeling/dflow/commit/1bfae67ab6533c9dddf8066a40919c56bb7c53b5))
+
 ## [1.8.93](https://github.com/deepmodeling/dflow/compare/v1.8.92...v1.8.93) (2024-09-20)
 
 

VERSION

@@ -1 +1 @@
-1.8.93
+1.8.98

src/dflow/plugins/dispatcher.py

@@ -58,6 +58,8 @@ class DispatcherExecutor(Executor):
         remote_root: remote root path for working
         retry_on_submission_error: max retries on submission error
         merge_sliced_step: handle multi slices in a single dispatcher job
+        terminate_grace_period: grace period in seconds after termination
+            for collecting outputs
     """
 
     def __init__(self,
@@ -89,6 +91,7 @@ def __init__(self,
                  remove_scheduling_strategies: bool = True,
                  envs: Optional[Dict[str, str]] = None,
                  merge_bohrium_job_group: bool = False,
+                 terminate_grace_period: Optional[int] = None,
                  ) -> None:
         self.host = host
         self.queue_name = queue_name
@@ -127,6 +130,7 @@ def __init__(self,
         self.remove_scheduling_strategies = remove_scheduling_strategies
         self.envs = envs
         self.merge_bohrium_job_group = merge_bohrium_job_group
+        self.terminate_grace_period = terminate_grace_period
 
         conf = {}
         if json_file is not None:
@@ -541,6 +545,10 @@ def render(self, template):
         new_template.script += "    print('Got SIGTERM, kill unfinished tasks"\
             "!')\n"
         new_template.script += "    submission.remove_unfinished_tasks()\n"
+        if self.terminate_grace_period:
+            new_template.script += "    import time\n"
+            new_template.script += "        time.sleep(%s)\n" % \
+                self.terminate_grace_period
         new_template.script += "import signal\n"
         new_template.script += "signal.signal(signal.SIGTERM, sigterm_handler"\
             ")\n"

src/dflow/python/op.py

@@ -20,6 +20,7 @@
                   OutputParameter, type_to_str)
 from ..utils import dict2list, get_key, randstr, s3_config
 from .opio import OPIO, Artifact, BigParameter, OPIOSign, Parameter
+from .utils import handle_output_artifact, handle_output_parameter
 from .vendor.typeguard import check_type
 
 iwd = os.getcwd()
@@ -48,6 +49,10 @@ class OP(ABC):
     progress_current = 0
     key = None
     workflow_name = None
+    outputs = {}
+    slices = {}
+    tmp_root = "/tmp"
+    create_slice_dir = False
 
     def __init__(
             self,
@@ -400,6 +405,26 @@ def from_graph(cls, graph):
             op = getattr(mod, graph["name"])
         return op
 
+    def set_output(self, name, value):
+        self.outputs[name] = value
+        self.handle_outputs({name: value}, symlink=True)
+
+    def handle_outputs(self, outputs, symlink=False):
+        os.makedirs("%s/outputs/parameters" % self.tmp_root, exist_ok=True)
+        os.makedirs("%s/outputs/artifacts" % self.tmp_root, exist_ok=True)
+        output_sign = self.get_output_sign()
+        for name in outputs:
+            sign = output_sign[name]
+            if isinstance(sign, Artifact):
+                slices = self.slices.get(name)
+                handle_output_artifact(
+                    name, outputs[name], sign, slices, self.tmp_root,
+                    self.create_slice_dir and slices, symlink=symlink)
+            else:
+                slices = self.slices.get(name)
+                handle_output_parameter(
+                    name, outputs[name], sign, slices, self.tmp_root)
+
 
 def type2opiosign(t):
     from typing import Tuple

src/dflow/python/python_op_template.py

@@ -593,6 +593,21 @@ def render_script(self):
                         "%s, r'%s')\n" % (name, name, name, name, slices,
                                           self.tmp_root)
 
+        script += "    op_obj.tmp_root = '%s'\n" % self.tmp_root
+        script += "    op_obj.create_slice_dir = %s\n" % (
+            self.create_slice_dir and getattr(self.slices, "pool_size", None)
+            is None)
+        for name, sign in output_sign.items():
+            if isinstance(sign, Artifact):
+                slices = self.get_slices(output_artifact_slices, name)
+                script += "    op_obj.slices['%s'] = %s\n" % (name, slices)
+
+        script += "    import signal\n"
+        script += "    def sigterm_handler(signum, frame):\n"
+        script += "        print('Got SIGTERM')\n"
+        script += "        raise RuntimeError('Got SIGTERM')\n"
+        script += "    signal.signal(signal.SIGTERM, sigterm_handler)\n"
+
         if self.slices is not None and self.slices.pool_size is not None:
             sliced_inputs = self.slices.input_artifact + \
                 self.slices.input_parameter
@@ -651,7 +666,7 @@ def render_script(self):
             script += "        if o is not None:\n"
             script += "            output = o\n"
             for name in sliced_outputs:
-                script += "    output['%s'] = [o['%s'] if o is not None"\
+                script += "    output['%s'] = [o.get('%s') if o is not None"\
                     " else None for o in output_list]\n" % (name, name)
                 if isinstance(output_sign[name], Artifact):
                     if output_sign[name].type == str:
@@ -662,31 +677,21 @@ def render_script(self):
                             "]\n" % name
         else:
             script += "    try:\n"
-            script += "        output = op_obj.execute(input)\n"
+            script += "        try:\n"
+            script += "            output = op_obj.execute(input)\n"
+            script += "        except Exception as e:\n"
+            script += "            if op_obj.outputs:\n"
+            script += "                op_obj.handle_outputs(op_obj.outputs)\n"
+            script += "            raise e\n"
             script += "    except TransientError:\n"
             script += "        traceback.print_exc()\n"
             script += "        sys.exit(1)\n"
             script += "    except FatalError:\n"
             script += "        traceback.print_exc()\n"
             script += "        sys.exit(2)\n"
 
-        script += "    os.makedirs(r'%s/outputs/parameters', exist_ok=True)\n"\
-            % self.tmp_root
-        script += "    os.makedirs(r'%s/outputs/artifacts', exist_ok=True)\n" \
-            % self.tmp_root
-        for name, sign in output_sign.items():
-            if isinstance(sign, Artifact):
-                slices = self.get_slices(output_artifact_slices, name)
-                script += "    handle_output_artifact('%s', output['%s'], "\
-                    "output_sign['%s'], %s, r'%s', %s)\n" % (
-                        name, name, name, slices, self.tmp_root,
-                        slices is not None and self.create_slice_dir and
-                        getattr(self.slices, "pool_size", None) is None)
-            else:
-                slices = self.get_slices(output_parameter_slices, name)
-                script += "    handle_output_parameter('%s', output['%s'], "\
-                    "output_sign['%s'], %s, r'%s')\n" % (name, name, name,
-                                                         slices, self.tmp_root)
+        script += "    op_obj.handle_outputs(output)\n"
+
         if config["register_tasks"]:
             if self.slices is not None and self.slices.register_first_only:
                 if "{{item}}" in self.dflow_vars:

src/dflow/python/utils.py

@@ -210,7 +210,9 @@ def slice_to_dir(slice):
 
 
 def handle_output_artifact(name, value, sign, slices=None, data_root="/tmp",
-                           create_dir=False):
+                           create_dir=False, symlink=False):
+    if os.path.isdir(data_root + '/outputs/artifacts/' + name):
+        shutil.rmtree(data_root + '/outputs/artifacts/' + name)
     path_list = []
     if sign.type == HDF5Datasets:
         import h5py
@@ -268,15 +270,16 @@ def handle_output_artifact(name, value, sign, slices=None, data_root="/tmp",
             slices = 0
         path_list.append(copy_results_and_return_path_item(
             value, name, slices, data_root,
-            slice_to_dir(slices) if create_dir else None))
+            slice_to_dir(slices) if create_dir else None, symlink=symlink))
     elif sign.type in [List[str], List[Path], Set[str], Set[Path]]:
         os.makedirs(data_root + '/outputs/artifacts/' + name, exist_ok=True)
         if slices is not None:
             if isinstance(slices, int):
                 for path in value:
                     path_list.append(copy_results_and_return_path_item(
                         path, name, slices, data_root,
-                        slice_to_dir(slices) if create_dir else None))
+                        slice_to_dir(slices) if create_dir else None,
+                        symlink=symlink))
             else:
                 assert len(slices) == len(value)
                 for path, s in zip(value, slices):
@@ -285,25 +288,26 @@ def handle_output_artifact(name, value, sign, slices=None, data_root="/tmp",
                             path_list.append(
                                 copy_results_and_return_path_item(
                                     p, name, s, data_root, slice_to_dir(
-                                        s) if create_dir else None))
+                                        s) if create_dir else None,
+                                    symlink=symlink))
                     else:
                         path_list.append(copy_results_and_return_path_item(
                             path, name, s, data_root, slice_to_dir(
-                                s) if create_dir else None))
+                                s) if create_dir else None, symlink=symlink))
         else:
             for s, path in enumerate(value):
                 path_list.append(copy_results_and_return_path_item(
-                    path, name, s, data_root))
+                    path, name, s, data_root, symlink=symlink))
     elif sign.type in [Dict[str, str], Dict[str, Path]]:
         os.makedirs(data_root + '/outputs/artifacts/' + name, exist_ok=True)
         for s, path in value.items():
             path_list.append(copy_results_and_return_path_item(
-                path, name, s, data_root))
+                path, name, s, data_root, symlink=symlink))
     elif sign.type in [NestedDict[str], NestedDict[Path]]:
         os.makedirs(data_root + '/outputs/artifacts/' + name, exist_ok=True)
         for s, path in flatten(value).items():
             path_list.append(copy_results_and_return_path_item(
-                path, name, s, data_root))
+                path, name, s, data_root, symlink=symlink))
 
     os.makedirs(data_root + "/outputs/artifacts/%s/%s" % (name, config[
         "catalog_dir_name"]), exist_ok=True)
@@ -343,15 +347,16 @@ def handle_output_parameter(name, value, sign, slices=None, data_root="/tmp"):
 
 
 def copy_results_and_return_path_item(path, name, order, data_root="/tmp",
-                                      slice_dir=None):
-    if path and os.path.exists(str(path)):
+                                      slice_dir=None, symlink=False):
+    if (path and os.path.exists(str(path))) or symlink:
         return {"dflow_list_item": copy_results(
-                    path, name, data_root, slice_dir), "order": order}
+                    path, name, data_root, slice_dir, symlink), "order": order}
     else:
         return {"dflow_list_item": None, "order": order}
 
 
-def copy_results(source, name, data_root="/tmp", slice_dir=None):
+def copy_results(source, name, data_root="/tmp", slice_dir=None,
+                 symlink=False):
     source = str(source)
     # if refer to input artifact
     if source.find(data_root + "/inputs/artifacts/") == 0:
@@ -364,7 +369,12 @@ def copy_results(source, name, data_root="/tmp", slice_dir=None):
         if slice_dir is not None:
             rel_path = "%s/%s" % (slice_dir, rel_path)
         target = data_root + "/outputs/artifacts/%s/%s" % (name, rel_path)
-        copy_file(source, target, shutil.copy)
+        if symlink:
+            os.makedirs(os.path.abspath(os.path.dirname(target)),
+                        exist_ok=True)
+            os.symlink(source, target)
+        else:
+            copy_file(source, target, shutil.copy)
         if rel_path[:1] == "/":
             rel_path = rel_path[1:]
         return rel_path
@@ -378,7 +388,12 @@ def copy_results(source, name, data_root="/tmp", slice_dir=None):
         if slice_dir is not None:
             rel_path = "%s/%s" % (slice_dir, rel_path)
         target = data_root + "/outputs/artifacts/%s/%s" % (name, rel_path)
-        copy_file(source, target)
+        if symlink:
+            os.makedirs(os.path.abspath(os.path.dirname(target)),
+                        exist_ok=True)
+            os.symlink(source, target)
+        else:
+            copy_file(source, target)
         return rel_path
 
 
@@ -456,7 +471,10 @@ def try_to_execute(input, slice_dir, op_obj, output_sign, cwd, timeout=None):
     except Exception as e:
         traceback.print_exc()
         os.chdir(cwd)
-        return None, e
+        if op_obj.outputs:
+            return op_obj.outputs, e
+        else:
+            return None, e
     finally:
         if timeout is not None:
             signal.alarm(0)