Streamline FederatedRuntime Workspace by removing Unnecessary Files (#…

…1262) * Clean created dir-workspace Signed-off-by: refai06 <[email protected]> * Update cleanup dir-workspace Signed-off-by: refai06 <[email protected]> * Update docstring Signed-off-by: refai06 <[email protected]> * Improve code-structure Signed-off-by: refai06 <[email protected]> --------- Signed-off-by: refai06 <[email protected]>
securefederatedai · Jan 13, 2025 · db9c0bd · db9c0bd
1 parent c35d02c
commit db9c0bd
Showing 1 changed file with 17 additions and 1 deletion.
diff --git a/openfl/experimental/workflow/workspace_export/export.py b/openfl/experimental/workflow/workspace_export/export.py
@@ -88,7 +88,6 @@ def __init__(self, notebook_path: str, output_workspace: str) -> None:
                 f"{export_filename}.py",
             )
         ).resolve()
-        print_tree(self.created_workspace_path, level=2)
 
         # Generated python script name without .py extension
         self.script_name = self.script_path.name.split(".")[0].strip()
@@ -290,6 +289,8 @@ def export_federated(cls, notebook_path: str, output_workspace: str) -> Tuple[st
         instance = cls(notebook_path, output_workspace)
         instance.generate_requirements()
         instance.generate_plan_yaml()
+        instance._clean_generated_workspace()
+        print_tree(output_workspace, level=2)
         return instance.generate_experiment_archive()
 
     @classmethod
@@ -304,6 +305,7 @@ def export(cls, notebook_path: str, output_workspace: str) -> None:
         instance.generate_requirements()
         instance.generate_plan_yaml()
         instance.generate_data_yaml()
+        print_tree(output_workspace, level=2)
 
     def generate_experiment_archive(self) -> Tuple[str, str]:
         """
@@ -357,6 +359,20 @@ def generate_requirements(self) -> None:
                 if i not in line_nos:
                     f.write(line)
 
+    def _clean_generated_workspace(self) -> None:
+        """
+        Remove cols.yaml and data.yaml from the generated workspace
+        as these are not needed in FederatedRuntime (Director based workflow)
+
+        """
+        cols_file = self.output_workspace_path.joinpath("plan", "cols.yaml")
+        data_file = self.output_workspace_path.joinpath("plan", "data.yaml")
+
+        if cols_file.exists():
+            cols_file.unlink()
+        if data_file.exists():
+            data_file.unlink()
+
     def generate_plan_yaml(self) -> None:
         """
         Generates plan.yaml