Skip to content

Commit

Permalink
adding xarg to env command, changing output_cpath to root used in the…
Browse files Browse the repository at this point in the history
… bindings
  • Loading branch information
djarecka committed Nov 3, 2023
1 parent 161635b commit 8d60dd1
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 19 deletions.
44 changes: 26 additions & 18 deletions pydra/engine/environments.py
Original file line number Diff line number Diff line change
Expand Up @@ -65,46 +65,53 @@ class Container(Environment):
Tag of the container image
output_cpath : str
Path to the output directory in the container
xargs : dict
xargs : Union[str, List[str]]
Extra arguments to be passed to the container
"""

def __init__(self, image, tag="latest", output_cpath="/output_pydra", xargs=None):
def __init__(self, image, tag="latest", root="/mnt/pydra", xargs=None):
self.image = image
self.tag = tag
if xargs is None:
xargs = []
elif isinstance(xargs, str):
xargs = xargs.split()
self.xargs = xargs
self.output_cpath = output_cpath
self.root = root

@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
def bind(loc, mode="ro", root="/mnt/pydra"):
loc_abs = Path(loc).absolute()
return f"{loc_abs}:{root}{loc_abs}:{mode}" # TODO: moving entire path?
return f"{loc_abs}:{root}{loc_abs}:{mode}"


class Docker(Container):
"""Docker environment."""

def execute(self, task, root="/mnt/pydra"):
def execute(self, task):
docker_img = f"{self.image}:{self.tag}"
# mounting all input locations
mounts = task.get_bindings(root=root)
mounts = task.get_bindings(root=self.root)

# todo adding xargsy etc
docker_args = ["docker", "run", "-v", self.bind(task.cache_dir, "rw")]
docker_args = [
"docker",
"run",
"-v",
self.bind(task.cache_dir, "rw", self.root),
]
docker_args.extend(self.xargs)
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}"])
docker_args.extend(["-w", f"{self.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"),
docker_args + [docker_img] + task.command_args(root=self.root),
strip=task.strip,
)
output = dict(zip(keys, values))
Expand All @@ -119,28 +126,29 @@ def execute(self, task, root="/mnt/pydra"):
class Singularity(Container):
"""Singularity environment."""

def execute(self, task, root="/mnt/pydra"):
def execute(self, task):
singularity_img = f"{self.image}:{self.tag}"
# mounting all input locations
mounts = task.get_bindings(root=root)
mounts = task.get_bindings(root=self.root)

# todo adding xargsy etc
singularity_args = [
"singularity",
"exec",
"-B",
self.bind(task.cache_dir, "rw"),
self.bind(task.cache_dir, "rw", self.root),
]
singularity_args.extend(self.xargs)
singularity_args.extend(
" ".join(
[f"-B {key}:{val[0]}:{val[1]}" for (key, val) in mounts.items()]
).split()
)
singularity_args.extend(["--pwd", f"{root}{task.output_dir}"])
singularity_args.extend(["--pwd", f"{self.root}{task.output_dir}"])
keys = ["return_code", "stdout", "stderr"]

values = execute(
singularity_args + [singularity_img] + task.command_args(root="/mnt/pydra"),
singularity_args + [singularity_img] + task.command_args(root=self.root),
strip=task.strip,
)
output = dict(zip(keys, values))
Expand Down
11 changes: 10 additions & 1 deletion pydra/engine/tests/test_environments.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from .utils import no_win, need_docker, need_singularity

import attr
import pytest


def makedir(path, name):
Expand Down Expand Up @@ -79,7 +80,15 @@ def test_docker_1(tmp_path):

@no_win
@need_docker
def test_docker_1_subm(tmp_path):
@pytest.mark.parametrize(
"docker",
[
Docker(image="busybox"),
Docker(image="busybox", tag="latest", xargs="--rm"),
Docker(image="busybox", xargs=["--rm"]),
],
)
def test_docker_1_subm(tmp_path, docker):
"""docker env with submitter: simple command, no arguments"""
newcache = lambda x: makedir(tmp_path, x)

Expand Down

0 comments on commit 8d60dd1

Please sign in to comment.