Refactor some common transformations (PolusAI#285)

* add user args for workflows built with python api

* remove redundant io in main & remove irrelevant warning/exit in compiler

* Refactor some common transformations into post_compile

---------

Co-authored-by: Vasu Jaganath <[email protected]>

Author: vjaganat90

src/sophios/api/http/restapi.py

@@ -13,6 +13,7 @@
 from sophios.utils_graphs import get_graph_reps
 from sophios.utils_yaml import wic_loader
 from sophios import utils_cwl
+from sophios.post_compile import cwl_inline_runtag
 from sophios.cli import get_args
 from sophios.wic_types import CompilerInfo, Json, Tool, Tools, StepId, YamlTree, Cwl, NodeData
 from sophios.api.utils import converter
@@ -132,9 +133,8 @@ async def compile_wf(request: Request) -> Json:
                                                             tools_cwl, True, relative_run_path=True, testing=False)
 
     rose_tree = compiler_info.rose
-    if args.inline_cwl_runtag:
-        input_output.write_to_disk(rose_tree, Path('autogenerated/'), True, args.inputs_file)
-        rose_tree = plugins.cwl_update_inline_runtag_rosetree(rose_tree, Path('autogenerated/'), True)
+    input_output.write_to_disk(rose_tree, Path('autogenerated/'), True, args.inputs_file)
+    cwl_inline_runtag(args, rose_tree)
     # ======== OUTPUT PROCESSING ================
     # ========= PROCESS COMPILED OBJECT =========
     sub_node_data: NodeData = rose_tree.data

src/sophios/api/pythonapi.py

@@ -12,7 +12,7 @@
 from cwl_utils.parser import load_document_by_uri, load_document_by_yaml
 from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, field_validator
 
-from sophios import compiler, input_output, plugins, utils_cwl
+from sophios import compiler, input_output, plugins, utils_cwl, post_compile
 from sophios import run_local as run_local_module
 from sophios.cli import get_args
 from sophios.utils_graphs import get_graph_reps
@@ -756,32 +756,11 @@ def run(self) -> None:
         args = get_args(self.process_name, self.user_args)  # Use mock CLI args
         rose_tree: RoseTree = compiler_info.rose
 
-        # cwl-docker-extract recursively `docker pull`s all images in all subworkflows.
-        # This is important because cwltool only uses `docker run` when executing
-        # workflows, and if there is a local image available,
-        # `docker run` will NOT query the remote repository for the latest image!
-        # cwltool has a --force-docker-pull option, but this may cause multiple pulls in parallel.
-        if args.container_engine == 'singularity':
-            cmd = ['cwl-docker-extract', '-s', '--dir',
-                   f'{args.singularity_pull_dir}', f'autogenerated/{self.process_name}.cwl']
-        else:
-            cmd = ['cwl-docker-extract', '--force-download', f'autogenerated/{self.process_name}.cwl']
-        sub.run(cmd, check=True)
-
-        # If you don't like it, you can programmatically overwrite anything in args
-        # args.docker_remove_entrypoints = True
-        if args.docker_remove_entrypoints:
-            # Requires root, so guard behind CLI option
-            if args.container_engine == 'docker':
-                plugins.remove_entrypoints_docker()
-            if args.container_engine == 'podman':
-                plugins.remove_entrypoints_podman()
-
-            rose_tree = plugins.dockerPull_append_noentrypoint_rosetree(rose_tree)
-            input_output.write_to_disk(rose_tree, Path('autogenerated/'), True, args.inputs_file)
+        post_compile.cwl_docker_extract(args, self.process_name)
+        post_compile.remove_entrypoints(args, rose_tree)
 
         # Do NOT capture stdout and/or stderr and pipe warnings and errors into a black hole.
-        retval = run_local_module.run_local(args, rose_tree, args.cachedir, 'cwltool', True)
+        retval = run_local_module.run_local(args, rose_tree, args.cachedir, args.cwl_runner, True)
 
         # Finally, since there is an output file copying bug in cwltool,
         # we need to copy the output files manually. See comment above.

src/sophios/main.py

@@ -11,6 +11,7 @@
 
 from sophios.utils_yaml import wic_loader
 from . import input_output as io
+from . import post_compile as pc
 from . import ast, cli, compiler, inference, inlineing, plugins, run_local, utils  # , utils_graphs
 from .schemas import wic_schema
 from .wic_types import GraphData, GraphReps, Json, StepId, Yaml, YamlTree
@@ -170,11 +171,8 @@ def main() -> None:
 
     io.write_to_disk(rose_tree, Path('autogenerated/'), True, args.inputs_file)
 
-    # this has to happen after at least one write
-    # so we can copy from local cwl_dapters in autogenerated/
-    if args.inline_cwl_runtag:
-        rose_tree = plugins.cwl_update_inline_runtag_rosetree(rose_tree, Path('autogenerated/'), True)
-        io.write_to_disk(rose_tree, Path('autogenerated/'), True, args.inputs_file)
+    pc.cwl_inline_runtag(args, rose_tree)
+    io.write_to_disk(rose_tree, Path('autogenerated/'), True, args.inputs_file)
 
     if args.graphviz:
         if shutil.which('dot'):
@@ -195,27 +193,8 @@ def main() -> None:
             print("but not the graphviz system package.)")
 
     if args.run_local or args.generate_run_script:
-        # cwl-docker-extract recursively `docker pull`s all images in all subworkflows.
-        # This is important because cwltool only uses `docker run` when executing
-        # workflows, and if there is a local image available,
-        # `docker run` will NOT query the remote repository for the latest image!
-        # cwltool has a --force-docker-pull option, but this may cause multiple pulls in parallel.
-        if args.container_engine == 'singularity':
-            cmd = ['cwl-docker-extract', '-s', '--dir',
-                   f'{args.singularity_pull_dir}', f'autogenerated/{yaml_stem}.cwl']
-        else:
-            cmd = ['cwl-docker-extract', '--force-download', f'autogenerated/{yaml_stem}.cwl']
-        sub.run(cmd, check=True)
-
-        if args.docker_remove_entrypoints:
-            # Requires root, so guard behind CLI option
-            if args.container_engine == 'docker':
-                plugins.remove_entrypoints_docker()
-            if args.container_engine == 'podman':
-                plugins.remove_entrypoints_podman()
-
-            rose_tree = plugins.dockerPull_append_noentrypoint_rosetree(rose_tree)
-            io.write_to_disk(rose_tree, Path('autogenerated/'), True, args.inputs_file)
+        pc.cwl_docker_extract(args, yaml_stem)
+        pc.remove_entrypoints(args, rose_tree)
 
         run_local.run_local(args, rose_tree, args.cachedir, args.cwl_runner, False)
 

src/sophios/post_compile.py

@@ -0,0 +1,43 @@
+import argparse
+from pathlib import Path
+import subprocess as sub
+
+from . import plugins
+from . import input_output as io
+from .wic_types import RoseTree
+
+
+def cwl_docker_extract(args: argparse.Namespace, file_name: str) -> None:
+    """Helper function to do the cwl_docker_extract"""
+    # cwl-docker-extract recursively `docker pull`s all images in all subworkflows.
+    # This is important because cwltool only uses `docker run` when executing
+    # workflows, and if there is a local image available,
+    # `docker run` will NOT query the remote repository for the latest image!
+    # cwltool has a --force-docker-pull option, but this may cause multiple pulls in parallel.
+    if args.container_engine == 'singularity':
+        cmd = ['cwl-docker-extract', '-s', '--dir',
+               f'{args.singularity_pull_dir}', f'autogenerated/{file_name}.cwl']
+    else:
+        cmd = ['cwl-docker-extract', '--force-download', f'autogenerated/{file_name}.cwl']
+    sub.run(cmd, check=True)
+
+
+def cwl_inline_runtag(args: argparse.Namespace, rose_tree: RoseTree) -> None:
+    """Transform with cwl inline runtag"""
+    # this has to happen after at least one write
+    # so we can copy from local cwl_dapters in autogenerated/
+    if args.inline_cwl_runtag:
+        rose_tree = plugins.cwl_update_inline_runtag_rosetree(rose_tree, Path('autogenerated/'), True)
+
+
+def remove_entrypoints(args: argparse.Namespace, rose_tree: RoseTree) -> None:
+    """Remove entry points"""
+    if args.docker_remove_entrypoints:
+        # Requires root, so guard behind CLI option
+        if args.container_engine == 'docker':
+            plugins.remove_entrypoints_docker()
+        if args.container_engine == 'podman':
+            plugins.remove_entrypoints_podman()
+
+        rose_tree = plugins.dockerPull_append_noentrypoint_rosetree(rose_tree)
+        io.write_to_disk(rose_tree, Path('autogenerated/'), True, args.inputs_file)

tests/test_examples.py

@@ -22,6 +22,7 @@
 from sophios import auto_gen_header
 from sophios.cli import get_args
 from sophios.utils_yaml import wic_loader
+from sophios.post_compile import cwl_docker_extract
 from sophios.wic_types import NodeData, StepId, Yaml, YamlTree, Json
 from sophios.utils_graphs import get_graph_reps
 
@@ -212,18 +213,7 @@ def run_workflows(yml_path_str: str, yml_path: Path, cwl_runner: str, args: argp
     sophios.input_output.write_to_disk(rose_tree, Path('autogenerated/'), True, args.inputs_file)
 
     if docker_pull_only:
-        # cwl-docker-extract recursively `docker pull`s all images in all subworkflows.
-        # This is important because cwltool only uses `docker run` when executing
-        # workflows, and if there is a local image available,
-        # `docker run` will NOT query the remote repository for the latest image!
-        # cwltool has a --force-docker-pull option, but this may cause multiple pulls in parallel.
-        if args.container_engine == 'singularity':
-            cmd = ['cwl-docker-extract', '-s', '--dir',
-                   f'{args.singularity_pull_dir}', f'autogenerated/{Path(yml_path).stem}.cwl']
-        else:
-            cmd = ['cwl-docker-extract', '--force-download', f'autogenerated/{Path(yml_path).stem}.cwl']
-        sub.run(cmd, check=True)
-
+        cwl_docker_extract(args, Path(yml_path).stem)
         return
 
     if args.docker_remove_entrypoints: