fix: sub_path slices with with_sequence

fix: slices of HDF5Datasets

Signed-off-by: zjgemi <liuxin_zijian@163.com>

Author: zjgemi

src/dflow/python/utils.py

@@ -78,14 +78,15 @@ def handle_input_artifact(name, sign, slices=None, data_root="/tmp",
         path_object = assemble_path_object(art_path)
         path_object = get_slices(path_object, prefix)
 
-    path_object = get_slices(path_object, slices)
-
     sign_type = sign.type
     if getattr(sign_type, "__origin__", None) == Union:
         args = sign_type.__args__
         if HDF5Datasets in args:
-            if isinstance(path_object, list) and all([isinstance(
-                    p, str) and p.endswith(".h5") for p in path_object]):
+            if isinstance(path_object, list) and len(path_object) > 0 and all([
+                isinstance(p, str) and p.endswith(".h5")
+                    for p in path_object]):
+                sign_type = HDF5Datasets
+            elif art_path.endswith(".h5"):
                 sign_type = HDF5Datasets
             elif args[0] == HDF5Datasets:
                 sign_type = args[1]
@@ -94,6 +95,8 @@ def handle_input_artifact(name, sign, slices=None, data_root="/tmp",
 
     if sign_type == HDF5Datasets:
         import h5py
+        if os.path.isfile(art_path):
+            path_object = [art_path]
         assert isinstance(path_object, list)
         res = None
         for path in path_object:
@@ -108,6 +111,10 @@ def handle_input_artifact(name, sign, slices=None, data_root="/tmp",
                 if res is None:
                     res = {}
                 res.update(datasets)
+        res = get_slices(res, slices)
+    else:
+        path_object = get_slices(path_object, slices)
+
     if sign_type in [str, Path]:
         if path_object is None or isinstance(path_object, str):
             res = path_object
@@ -146,7 +153,7 @@ def handle_input_artifact(name, sign, slices=None, data_root="/tmp",
         return None
 
     _cls = res.__class__
-    res = artifact_classes[_cls](res)
+    res = artifact_classes.get(_cls, lambda x: x)(res)
     res.art_root = root
     return res
 
@@ -230,6 +237,9 @@ def handle_output_artifact(name, value, sign, slices=None, data_root="/tmp",
                         d.attrs["type"] = "dir"
                         d.attrs["path"] = str(v)
                         d.attrs["dtype"] = "binary"
+                elif isinstance(v, HDF5Dataset):
+                    d = f.create_dataset(s, data=v.dataset[()])
+                    d.attrs.update(v.dataset.attrs)
                 else:
                     d = f.create_dataset(s, data=v)
                     d.attrs["type"] = "data"

src/dflow/step.py

@@ -484,7 +484,7 @@ def __init__(
                 auto_loop_artifacts:
             self.template = self.template.deepcopy()
             sequence_format = self.with_sequence.format \
-                if self.with_sequence is not None else "%d"
+                if self.with_sequence is not None else None
             init_template = InitArtifactForSlices(
                 self.template, self.util_image, self.util_command,
                 self.util_image_pull_policy, self.key, sliced_output_artifact,
@@ -598,6 +598,13 @@ def merge_output_artifact(art, parent, layer=0):
                         elif v is not None:
                             self.prepare_step.set_artifacts({name: v})
                             self.inputs.artifacts[name].sp = "{{item.%s}}" % name
+                        if self.with_sequence is not None:
+                            for par in self.inputs.parameters.values():
+                                if hasattr(par, "value") and isinstance(
+                                        par.value, str):
+                                    par.value = par.value.replace(
+                                        "{{item}}", "{{item.order}}")
+                            self.with_sequence = None
                         self.with_param = self.prepare_step.outputs.parameters[
                             "dflow_slices_path"]
 

src/dflow/util_ops.py

@@ -10,7 +10,7 @@ class InitArtifactForSlices(PythonScriptOPTemplate):
     def __init__(self, template, image, command, image_pull_policy, key,
                  sliced_output_artifact=None, sliced_input_artifact=None,
                  sum_var=None, concat_var=None, auto_loop_artifacts=None,
-                 group_size=None, format="%d", post_script="",
+                 group_size=None, format=None, post_script="",
                  tmp_root="/tmp"):
         name = template.name
         super().__init__(name="%s-init-artifact" % name, image=image,
@@ -171,7 +171,7 @@ def render_script(self):
                     script += "        continue\n"
                     script += "    group_dir = r'%s/outputs/artifacts/%s/"\
                         "group_' + ('%s' %% i)\n" % (
-                            self.tmp_root, name, self.format)
+                            self.tmp_root, name, self.format or "%d")
                     script += "    os.makedirs(group_dir, exist_ok=True)\n"
                     if self.template.slices.shuffle:
                         script += "    path_list = [path_list_%s[j] for j in "\
@@ -206,7 +206,10 @@ def render_script(self):
             else:
                 script += "slices_path = []\n"
                 script += "for i in range(len(path_list_%s)):\n" % required[0]
-                script += "    item = {'order': i}\n"
+                if self.format:
+                    script += "    item = {'order': '%s' %% i}\n" % self.format
+                else:
+                    script += "    item = {'order': i}\n"
                 for i, name in enumerate(self.sliced_input_artifact):
                     script += "    item['%s'] = path_list_%s[i]"\
                         "['dflow_list_item'] if path_list_%s else None\n" % (