integrated generate_graph into the bee code

beeflow/common/deps/neo4j_manager.py

@@ -24,6 +24,7 @@
 certs_dir = mount_dir + '/certificates'
 confs_dir = mount_dir + "/conf"
 dags_dir = os.path.join(bee_workdir, 'dags')
+graphmls_dir = dags_dir + "/graphmls"
 container_path = container_manager.get_container_dir('neo4j')
 log = bee_logging.setup('neo4j')
 
@@ -61,6 +62,7 @@ def setup_mounts():
     os.makedirs(certs_dir, exist_ok=True)
     os.makedirs(run_dir, exist_ok=True)
     os.makedirs(dags_dir, exist_ok=True)
+    os.makedirs(graphmls_dir, exist_ok=True)
 
 
 def setup_configs(bolt_port, http_port, https_port):
@@ -120,7 +122,7 @@ def create_credentials():
             "-b", data_dir + ":/data",
             "-b", logs_dir + ":/logs",
             "-b", run_dir + ":/var/lib/neo4j/run", container_path,
-            "-W", "-b", dags_dir + ":/var/lib/neo4j/import",
+            "-W", "-b", graphmls_dir + ":/var/lib/neo4j/import",
             "--", *command
         ], check=True)
     except subprocess.CalledProcessError:
@@ -140,7 +142,7 @@ def create_database():
             "-b", logs_dir + ":/logs",
             "-b", run_dir + ":/var/lib/neo4j/run",
             "-b", certs_dir + ":/var/lib/neo4j/certificates",
-            "-W", "-b", dags_dir + ":/var/lib/neo4j/import",
+            "-W", "-b", graphmls_dir + ":/var/lib/neo4j/import",
             container_path, "--", *command
         ], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
         wait_gdb()

generate_graph.py → beeflow/common/gdb/generate_graph.py

@@ -1,7 +1,18 @@
+import os
 import networkx as nx
 import graphviz
 
-def graphml_to_graphviz(graphml_path, output_path):
+from beeflow.common import paths
+
+bee_workdir = paths.workdir()
+dags_dir = os.path.join(bee_workdir, 'dags')
+graphmls_dir = dags_dir + "/graphmls"
+
+def generate_viz(wf_id):
+    short_id = wf_id[:6]
+    graphml_path = graphmls_dir + "/" + short_id + ".graphml"
+    output_path = dags_dir + "/" + short_id
+
     # Load the GraphML file using NetworkX
     G = nx.read_graphml(graphml_path)
 
@@ -11,7 +22,7 @@ def graphml_to_graphviz(graphml_path, output_path):
     # Add nodes to the Graphviz graph
     for node in G.nodes(data=True):
         node_id = node[0]
-        label = node[1].get('labels', node_id)  # Use label if available, otherwise node_id
+        label = node[1].get('labels', node_id)
         if label == ":Workflow":
             node_label = "Workflow"
             color ='steelblue'
@@ -34,7 +45,7 @@ def graphml_to_graphviz(graphml_path, output_path):
     for edge in G.edges(data=True):
         source = edge[0]
         target = edge[1]
-        edge_label = edge[2].get('label', '')  # Use edge label if available
+        edge_label = edge[2].get('label', '')
         if edge_label == "INPUT_OF" or edge_label == "DESCRIBES":
             dot.edge(source, target, label=edge_label, fontsize="10")
         else:
@@ -43,14 +54,3 @@ def graphml_to_graphviz(graphml_path, output_path):
     # Set the output format to PNG and render the graph
     dot.format = 'png'
     dot.render(output_path, view=False)
-
-if __name__ == "__main__":
-    # Path to your GraphML file
-    graphml_file = '/vast/home/leahh/.beeflow/dags/bebd6f.graphml'
-    output_file = 'hierarchical_graph_clamr'
-
-    # Generate the hierarchical graph
-    graphml_to_graphviz(graphml_file, output_file)
-
-    print(f"Graph has been saved as {output_file}.png")
-

beeflow/wf_manager/resources/wf_utils.py

@@ -8,6 +8,7 @@
 from beeflow.common import log as bee_logging
 from beeflow.common.config_driver import BeeConfig as bc
 from beeflow.common.gdb import neo4j_driver
+from beeflow.common.gdb.generate_graph import generate_viz
 from beeflow.common.wf_interface import WorkflowInterface
 from beeflow.common.connection import Connection
 from beeflow.common import paths
@@ -301,6 +302,7 @@ def export_dag(wf_id):
     """Export the DAG of the workflow."""
     wfi = get_workflow_interface(wf_id)
     wfi.export_graphml()
+    generate_viz(wf_id)
 
 
 def start_workflow(wf_id):