NF: Add Environment class, with initial Native/Docker implementations

pydra/__init__.py

@@ -15,14 +15,13 @@
 import attr
 
 from . import mark
-from .engine import AuditFlag, DockerTask, ShellCommandTask, Submitter, Workflow, specs
+from .engine import AuditFlag, ShellCommandTask, Submitter, Workflow, specs
 
 __all__ = (
     "Submitter",
     "Workflow",
     "AuditFlag",
     "ShellCommandTask",
-    "DockerTask",
     "specs",
     "mark",
 )

pydra/engine/__init__.py

@@ -1,12 +1,11 @@
 """The core of the workflow engine."""
 from .submitter import Submitter
 from .core import Workflow
-from .task import AuditFlag, ShellCommandTask, DockerTask
+from .task import AuditFlag, ShellCommandTask
 from . import specs
 
 __all__ = [
     "AuditFlag",
-    "DockerTask",
     "ShellCommandTask",
     "Submitter",
     "Workflow",

pydra/engine/core.py

@@ -429,7 +429,13 @@ def cont_dim(self, cont_dim):
             self._cont_dim = cont_dim
 
     def __call__(
-        self, submitter=None, plugin=None, plugin_kwargs=None, rerun=False, **kwargs
+        self,
+        submitter=None,
+        plugin=None,
+        plugin_kwargs=None,
+        rerun=False,
+        environment=None,
+        **kwargs,
     ):
         """Make tasks callable themselves."""
         from .submitter import Submitter
@@ -449,9 +455,9 @@ def __call__(
         if submitter:
             with submitter as sub:
                 self.inputs = attr.evolve(self.inputs, **kwargs)
-                res = sub(self)
+                res = sub(self, environment=environment)
         else:  # tasks without state could be run without a submitter
-            res = self._run(rerun=rerun, **kwargs)
+            res = self._run(rerun=rerun, environment=environment, **kwargs)
         return res
 
     def _modify_inputs(self):
@@ -501,7 +507,7 @@ def _populate_filesystem(self, checksum, output_dir):
             shutil.rmtree(output_dir)
         output_dir.mkdir(parents=False, exist_ok=self.can_resume)
 
-    def _run(self, rerun=False, **kwargs):
+    def _run(self, rerun=False, environment=None, **kwargs):
         self.inputs = attr.evolve(self.inputs, **kwargs)
         self.inputs.check_fields_input_spec()
 
@@ -518,6 +524,7 @@ def _run(self, rerun=False, **kwargs):
                     return result
             cwd = os.getcwd()
             self._populate_filesystem(checksum, output_dir)
+            os.chdir(output_dir)
             orig_inputs = self._modify_inputs()
             result = Result(output=None, runtime=None, errored=False)
             self.hooks.pre_run_task(self)
@@ -526,7 +533,7 @@ def _run(self, rerun=False, **kwargs):
                 self.audit.audit_task(task=self)
             try:
                 self.audit.monitor()
-                self._run_task()
+                self._run_task(environment=environment)
                 result.output = self._collect_outputs(output_dir=output_dir)
             except Exception:
                 etype, eval, etr = sys.exc_info()
@@ -538,7 +545,7 @@ def _run(self, rerun=False, **kwargs):
                 self.hooks.post_run_task(self, result)
                 self.audit.finalize_audit(result)
                 save(output_dir, result=result, task=self)
-                self.output_ = None
+                # self.output_ = None
                 # removing the additional file with the chcksum
                 (self.cache_dir / f"{self.uid}_info.json").unlink()
                 # # function etc. shouldn't change anyway, so removing

pydra/engine/environments.py

@@ -0,0 +1,84 @@
+from .helpers import execute
+
+from pathlib import Path
+
+
+class Environment:
+    def setup(self):
+        pass
+
+    def execute(self, task):
+        raise NotImplementedError
+
+    def teardown(self):
+        pass
+
+
+class Native(Environment):
+    def execute(self, task):
+        # breakpoint()
+        # args = task.render_arguments_in_root()
+        keys = ["return_code", "stdout", "stderr"]
+        values = execute(task.command_args(), strip=task.strip)
+        output = dict(zip(keys, values))
+        if output["return_code"]:
+            msg = f"Error running '{task.name}' task with {task.command_args()}:"
+            if output["stderr"]:
+                msg += "\n\nstderr:\n" + output["stderr"]
+            if output["stdout"]:
+                msg += "\n\nstdout:\n" + output["stdout"]
+            raise RuntimeError(msg)
+        return output
+
+
+class Docker(Environment):
+    def __init__(
+        self, image, tag="latest", binds=None, output_cpath="/output_pydra", xargs=None
+    ):
+        self.image = image
+        self.tag = tag
+        self.xargs = xargs
+        self.output_cpath = output_cpath
+
+    @staticmethod
+    def bind(loc, mode="ro", root="/mnt/pydra"):  # TODO
+        # XXX Failure mode: {loc} overwrites a critical directory in image
+        # To fix, we'll need to update any args within loc to a new location
+        # such as /mnt/pydra/loc
+        loc_abs = Path(loc).absolute()
+        return f"{loc_abs}:{root}{loc_abs}:{mode}"  # TODO: moving entire path?
+
+    def execute(self, task, root="/mnt/pydra"):
+        # XXX Need to mount all input locations
+        docker_img = f"{self.image}:{self.tag}"
+        # TODO ?
+        # Skips over any inputs in task.cache_dir
+        # Needs to include `out_file`s when not relative to working dir
+        # Possibly a `TargetFile` type to distinguish between `File` and `str`?
+        mounts = task.get_bindings(root=root)
+
+        # todo adding xargsy etc
+        docker_args = ["docker", "run", "-v", self.bind(task.cache_dir, "rw")]
+        docker_args.extend(
+            " ".join(
+                [f"-v {key}:{val[0]}:{val[1]}" for (key, val) in mounts.items()]
+            ).split()
+        )
+        docker_args.extend(["-w", f"{root}{task.output_dir}"])
+        keys = ["return_code", "stdout", "stderr"]
+        # print("\n Docker args", docker_args)
+
+        values = execute(
+            docker_args + [docker_img] + task.command_args(root="/mnt/pydra"),
+            strip=task.strip,
+        )
+        output = dict(zip(keys, values))
+        if output["return_code"]:
+            if output["stderr"]:
+                raise RuntimeError(output["stderr"])
+            else:
+                raise RuntimeError(output["stdout"])
+        # Any outputs that have been created with a re-rooted path need
+        # to be de-rooted
+        # task.finalize_outputs("/mnt/pydra") TODO: probably don't need it
+        return output

pydra/engine/specs.py

@@ -693,13 +693,6 @@ class ContainerSpec(ShellSpec):
     )
 
 
-@attr.s(auto_attribs=True, kw_only=True)
-class DockerSpec(ContainerSpec):
-    """Particularize container specifications to the Docker engine."""
-
-    container: str = attr.ib("docker", metadata={"help_string": "container"})
-
-
 @attr.s(auto_attribs=True, kw_only=True)
 class SingularitySpec(ContainerSpec):
     """Particularize container specifications to Singularity."""

pydra/engine/submitter.py

@@ -35,14 +35,16 @@ def __init__(self, plugin="cf", **kwargs):
             raise NotImplementedError(f"No worker for {self.plugin}")
         self.worker.loop = self.loop
 
-    def __call__(self, runnable, cache_locations=None, rerun=False):
+    def __call__(self, runnable, cache_locations=None, rerun=False, environment=None):
         """Submitter run function."""
         if cache_locations is not None:
             runnable.cache_locations = cache_locations
-        self.loop.run_until_complete(self.submit_from_call(runnable, rerun))
+        self.loop.run_until_complete(
+            self.submit_from_call(runnable, rerun, environment)
+        )
         return runnable.result()
 
-    async def submit_from_call(self, runnable, rerun):
+    async def submit_from_call(self, runnable, rerun, environment):
         """
         This coroutine should only be called once per Submitter call,
         and serves as the bridge between sync/async lands.
@@ -56,7 +58,7 @@ async def submit_from_call(self, runnable, rerun):
         Once Python 3.10 is the minimum, this should probably be refactored into using
         structural pattern matching.
         """
-        if is_workflow(runnable):
+        if is_workflow(runnable):  # TODO: env to wf
             # connect and calculate the checksum of the graph before running
             runnable._connect_and_propagate_to_tasks(override_task_caches=True)
             # 0
@@ -74,10 +76,11 @@ async def submit_from_call(self, runnable, rerun):
             # 2
             if runnable.state is None:
                 # run_el should always return a coroutine
-                await self.worker.run_el(runnable, rerun=rerun)
+                print("in SUBM", environment)
+                await self.worker.run_el(runnable, rerun=rerun, environment=environment)
             # 3
             else:
-                await self.expand_runnable(runnable, wait=True, rerun=rerun)
+                await self.expand_runnable(runnable, wait=True, rerun=rerun)  # TODO
         return True
 
     async def expand_runnable(self, runnable, wait=False, rerun=False):