fix: add dispatcher_debug

Signed-off-by: zjgemi <[email protected]>
deepmodeling · May 24, 2024 · a4eeb69 · a4eeb69
1 parent 978d0a9
commit a4eeb69
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/src/dflow/config.py b/src/dflow/config.py
@@ -79,6 +79,7 @@ def split_headers(s):
     "debug_save_copy_method": os.environ.get("DFLOW_DEBUG_SAVE_COPY_METHOD",
                                              "symlink"),
     "raise_for_group": boolize(os.environ.get("DFLOW_RAISE_FOR_GROUP", False)),
+    "dispatcher_debug": boolize(os.environ.get("DISPATCHER_DEBUG", False)),
 }
 
 

diff --git a/src/dflow/plugins/dispatcher.py b/src/dflow/plugins/dispatcher.py
@@ -336,6 +336,10 @@ def render(self, template):
 
         machine_dict["local_root"] = self.work_root
         new_template.script += "import json, shlex\n"
+        if config["dispatcher_debug"]:
+            new_template.script += "import logging\n"
+            new_template.script += "from dpdispatcher.dlog import dlog\n"
+            new_template.script += "dlog.setLevel(logging.DEBUG)\n"
         new_template.script += "from dpdispatcher import Machine, Resources,"\
             " Task, Submission\n"
         new_template.script += "machine = Machine.load_from_dict(json.loads("\