Export dag

beeflow/client/bee_client.py

@@ -39,6 +39,7 @@
 
 
 logging.basicConfig(level=logging.WARNING)
+logging.getLogger("neo4j").setLevel(logging.WARNING)
 WORKFLOW_MANAGER = 'bee_wfm/v1/jobs/'
 
 
@@ -593,6 +594,13 @@ def reexecute(wf_name: str = typer.Argument(..., help='The workflow name'),
     return wf_id
 
 
+@app.command()
+def dag(wf_id: str = typer.Argument(..., callback=match_short_id)):
+    """Export a DAG of the workflow to a GraphML file."""
+    wf_utils.export_dag(wf_id)
+    typer.echo(f"DAG for workflow {wf_id} has been exported successfully.")
+
+
 @app.callback(invoke_without_command=True)
 def version_callback(version: bool = False):
     """Beeflow."""

beeflow/common/deps/neo4j_manager.py

@@ -23,6 +23,7 @@
 run_dir = mount_dir + '/run'
 certs_dir = mount_dir + '/certificates'
 confs_dir = mount_dir + "/conf"
+dags_dir = os.path.join(bee_workdir, 'dags')
 container_path = container_manager.get_container_dir('neo4j')
 log = bee_logging.setup('neo4j')
 
@@ -59,6 +60,7 @@ def setup_mounts():
     os.makedirs(mount_dir, exist_ok=True)
     os.makedirs(certs_dir, exist_ok=True)
     os.makedirs(run_dir, exist_ok=True)
+    os.makedirs(dags_dir, exist_ok=True)
 
 
 def setup_configs(bolt_port, http_port, https_port):
@@ -86,6 +88,24 @@ def setup_configs(bolt_port, http_port, https_port):
     with open(gdb_configfile, "wt", encoding="utf8") as cfile:
         cfile.write(data)
 
+    apoc_configfile = os.path.join(confs_dir, "apoc.conf")
+    if not os.path.exists(apoc_configfile):
+        with open(apoc_configfile, "wt", encoding="utf8") as afile:
+            afile.write("apoc.export.file.enabled=true\n")
+        log.info(f"Created {apoc_configfile} with apoc.export.file.enabled=true")
+    else:
+        with open(apoc_configfile, "rt", encoding="utf8") as afile:
+            apoc_data = afile.read()
+
+        apoc_config = r'#?(apoc.export.file.enabled=)[^\n]*'
+        if re.search(apoc_config, apoc_data):
+            apoc_data = re.sub(apoc_config, r'apoc.export.file.enabled=true', apoc_data)
+        else:
+            apoc_data += "\napoc.export.file.enabled=true\n"
+
+        with open(apoc_configfile, "wt", encoding="utf8") as afile:
+            afile.write(apoc_data)
+
 
 def create_credentials():
     """Create the password and set the logfiles in environment."""
@@ -95,10 +115,12 @@ def create_credentials():
         subprocess.run([
             "ch-run",
             "--set-env=" + container_path + "/ch/environment",
+            "--set-env=apoc.export.file.enabled=true",
             "-b", confs_dir + ":/var/lib/neo4j/conf",
             "-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",
             "--", *command
         ], check=True)
     except subprocess.CalledProcessError:
@@ -112,11 +134,13 @@ def create_database():
         proc = subprocess.Popen([ #noqa can't use with because returning
             "ch-run",
             "--set-env=" + container_path + "/ch/environment",
+            "--set-env=apoc.export.file.enabled=true",
             "-b", confs_dir + ":/var/lib/neo4j/conf",
             "-b", data_dir + ":/data",
             "-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",
             container_path, "--", *command
         ], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
         wait_gdb()

beeflow/common/gdb/gdb_driver.py

@@ -287,3 +287,7 @@ def workflow_completed(self):
     @abstractmethod
     def close(self):
         """Close the connection to the graph database."""
+
+    @abstractmethod
+    def export_dag(self):
+        """Export a BEE workflow as a graphml."""

beeflow/common/gdb/neo4j_cypher.py

@@ -725,3 +725,19 @@ def cleanup(tx):
     cleanup_query = "MATCH (n) DETACH DELETE n"
 
     tx.run(cleanup_query)
+
+
+def export_dag(tx, wf_id):
+    """Export BEE workflow as graphml."""
+    short_id = wf_id[:6]
+    export_query = (
+        "WITH \"MATCH (n1)-[r]->(n2) "
+        f"WHERE n1.workflow_id = '{wf_id}' OR n2.workflow_id = '{wf_id}' "
+        "RETURN r, n1, n2\" AS query "
+        f"CALL apoc.export.graphml.query(query, '{short_id}.graphml', {{useTypes: true}}) "
+        "YIELD file, source, format, nodes, relationships, properties, time, rows, batchSize, "
+        "batches, done, data "
+        "RETURN file, source, format, nodes, relationships, properties, time, rows, batchSize, "
+        "batches, done, data"
+    )
+    tx.run(export_query)

beeflow/common/gdb/neo4j_driver.py

@@ -472,6 +472,11 @@ def _write_transaction(self, tx_fun, **kwargs):
         with self._driver.session() as session:
             session.write_transaction(tx_fun, **kwargs)
 
+    def export_dag(self, workflow_id):
+        """Export a BEE workflow as a graphml."""
+        with self._driver.session() as session:
+            session.write_transaction(tx.export_dag, wf_id=workflow_id)
+
 
 def _reconstruct_requirements(req_records):
     """Reconstruct requirements by their records retrieved from Neo4j.

beeflow/common/wf_interface.py

@@ -304,3 +304,7 @@ def workflow_completed(self):
         :rtype: bool
         """
         return self._gdb_driver.workflow_completed(self._workflow_id)
+
+    def export_dag(self):
+        """Export a BEE workflow as a graphml."""
+        self._gdb_driver.export_dag(self._workflow_id)

beeflow/wf_manager/resources/wf_utils.py

@@ -297,6 +297,12 @@ def setup_workflow(wf_id, wf_name, wf_dir, wf_workdir, no_start, workflow=None,
         start_workflow(wf_id)
 
 
+def export_dag(wf_id):
+    """Export the DAG of the workflow."""
+    wfi = get_workflow_interface(wf_id)
+    wfi.export_dag()
+
+
 def start_workflow(wf_id):
     """Attempt to start the workflow, returning True if successful."""
     db = connect_db(wfm_db, get_db_path())

generate_graph.py

@@ -0,0 +1,56 @@
+import networkx as nx
+import graphviz
+
+def graphml_to_graphviz(graphml_path, output_path):
+    # Load the GraphML file using NetworkX
+    G = nx.read_graphml(graphml_path)
+
+    # Create a new directed graph for Graphviz
+    dot = graphviz.Digraph(comment='Hierarchical Graph')
+
+    # 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
+        if label == ":Workflow":
+            node_label = "Workflow"
+            color ='steelblue'
+        if label == ":Output":
+            node_label = node[1].get('glob', node_id)
+            color = 'mediumseagreen'
+        if label == ":Metadata":
+            node_label = node[1].get('state', node_id)
+            color = 'skyblue'
+        if label == ":Task":
+            node_label = node[1].get('name', node_id)
+            color ='lightcoral'
+        if label == ":Input":
+            node_label = node[1].get('source', node_id)
+            color ='sandybrown'
+
+        dot.node(node_id, label=node_label, style='filled', fillcolor=color)
+
+    # Add edges to the Graphviz graph
+    for edge in G.edges(data=True):
+        source = edge[0]
+        target = edge[1]
+        edge_label = edge[2].get('label', '')  # Use edge label if available
+        if edge_label == "INPUT_OF" or edge_label == "DESCRIBES":
+            dot.edge(source, target, label=edge_label, fontsize="10")
+        else:
+            dot.edge(target, source, label=edge_label, fontsize="10")
+
+    # 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")
+

pyproject.toml

@@ -77,6 +77,7 @@ docutils = "0.18.1"
 google-api-python-client = { version = "^2.66.0", optional = true }
 python-openstackclient = { version = "^6.0.0", optional = true }
 python-heatclient = { version = "^3.1.0", optional = true }
+graphviz = "^0.20.3"
 
 [tool.poetry.extras]
 cloud_extras = ["google-api-python-client", "python-openstackclient", "python-heatclient"]