From fe6fdd21089baf1684f87ec20a4e948dc4e2086b Mon Sep 17 00:00:00 2001 From: zjgemi Date: Wed, 7 Aug 2024 10:32:22 +0800 Subject: [PATCH] fix: multiprocessing in debug mode Signed-off-by: zjgemi --- src/dflow/python/python_op_template.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/dflow/python/python_op_template.py b/src/dflow/python/python_op_template.py index ba00b9ee..af0bec21 100644 --- a/src/dflow/python/python_op_template.py +++ b/src/dflow/python/python_op_template.py @@ -504,7 +504,8 @@ def render_script(self): " handle_input_parameter\n" script += "from dflow.python.utils import handle_output_artifact," \ " handle_output_parameter, handle_lineage\n" - script += f"from {mod} import {class_name}\n\n" + if mod not in ["__main__", "__mp_main__"]: + script += f"from {mod} import {class_name}\n\n" if hasattr(op_class, "func"): script += "op_obj = %s\n" % class_name elif op is None: