split up environments module

Author: tclose

pydra/engine/submitter.py

@@ -11,7 +11,7 @@
 from datetime import datetime
 from collections import defaultdict
 import attrs
-from pydra.workers.base import Worker, WORKERS
+import logging
 from pydra.engine.graph import DiGraph
 from pydra.utils.general import (
     task_fields,
@@ -25,7 +25,8 @@
 from pydra.utils.general import default_run_cache_dir
 from pydra.compose import workflow
 from pydra.engine.state import State
-import logging
+from pydra.workers.base import Worker
+
 
 logger = logging.getLogger("pydra.submitter")
 
@@ -143,24 +144,16 @@ def __init__(
         self._own_loop = not self.loop.is_running()
         if isinstance(worker, Worker):
             self._worker = worker
-            self.worker_name = worker.plugin_name
         else:
-            if isinstance(worker, str):
-                self.worker_name = worker
-                try:
-                    worker_cls = WORKERS[self.worker_name]
-                except KeyError:
-                    raise NotImplementedError(
-                        f"No worker for '{self.worker_name}' worker"
-                    )
-            else:
-                try:
-                    self.worker_name = worker.plugin_name
-                except AttributeError:
-                    raise ValueError(
-                        "Worker class must have a 'plugin_name' str attribute"
-                    )
+            if issubclass(worker, Worker):
                 worker_cls = worker
+            elif isinstance(worker, str):
+                worker_cls = Worker.plugin(worker)
+            else:
+                raise TypeError(
+                    "Worker must be a Worker object, name of a worker or a Worker "
+                    f"class, not {type(worker)}"
+                )
             try:
                 self._worker = worker_cls(**kwargs)
             except TypeError as e:
@@ -170,7 +163,7 @@ def __init__(
         self.clean_stale_locks = (
             clean_stale_locks
             if clean_stale_locks is not None
-            else (self.worker_name == "debug")
+            else (self.worker.plugin_name() == "debug")
         )
         self.worker_kwargs = kwargs
         self._worker.loop = self.loop
@@ -216,7 +209,7 @@ def __call__(
         from pydra.environments.base import Environment
 
         if raise_errors is None:
-            raise_errors = self.worker_name == "debug"
+            raise_errors = self.worker.plugin_name() == "debug"
         if not isinstance(raise_errors, bool):
             raise TypeError(
                 f"'raise_errors' must be a boolean or None, not {type(raise_errors)}"
@@ -317,14 +310,12 @@ def __getstate__(self):
         # Remove the unpicklable entries or those that should not be pickled
         # When unpickled (in another process) the submitter can't be called
         state["loop"] = None
-        state["_worker"] = None
         return state
 
     def __setstate__(self, state):
         self.__dict__.update(state)
         # Restore the loop and worker
         self.loop = get_open_loop()
-        self._worker = WORKERS[self.worker_name](**self.worker_kwargs)
         self.worker.loop = self.loop
 
     def expand_workflow(self, workflow_task: "Job[workflow.Task]", rerun: bool) -> None:

pydra/engine/tests/test_dockertask.py

@@ -2,7 +2,7 @@
 import pytest
 from pydra.engine.submitter import Submitter
 from fileformats.generic import File
-from pydra.environments.docker import Docker
+from pydra.environments import docker
 from pydra.compose import shell, workflow
 from pydra.engine.job import Job
 from .utils import no_win, need_docker, run_submitter, run_no_submitter
@@ -20,11 +20,13 @@ def test_docker_1_nosubm(tmp_path):
     docky_task = Job(
         task=docky,
         name="docky",
-        submitter=Submitter(environment=Docker(image="busybox"), cache_dir=tmp_path),
+        submitter=Submitter(
+            environment=docker.Environment(image="busybox"), cache_dir=tmp_path
+        ),
     )
     assert docky_task.environment.image == "busybox"
     assert docky_task.environment.tag == "latest"
-    assert isinstance(docky_task.environment, Docker)
+    assert isinstance(docky_task.environment, docker.Environment)
     assert docky.cmdline == cmd
 
     res = docky_task.run()
@@ -42,7 +44,9 @@ def test_docker_1(worker, tmp_path):
     Docky = shell.define(cmd)
     docky = Docky()
 
-    with Submitter(cache_dir=tmp_path, environment=Docker(image="busybox")) as sub:
+    with Submitter(
+        cache_dir=tmp_path, environment=docker.Environment(image="busybox")
+    ) as sub:
         res = sub(docky)
 
     assert res.outputs.stdout == "root\n"
@@ -61,7 +65,9 @@ def test_docker_2(run_function, worker, tmp_path):
     docky = Docky()
     # cmdline doesn't know anything about docker
     assert docky.cmdline == cmdline
-    outputs = run_function(docky, tmp_path, worker, environment=Docker(image="busybox"))
+    outputs = run_function(
+        docky, tmp_path, worker, environment=docker.Environment(image="busybox")
+    )
     assert outputs.stdout.strip() == " ".join(cmdline.split()[1:])
     assert outputs.return_code == 0
 
@@ -80,7 +86,9 @@ def test_docker_2a(run_function, worker, tmp_path):
     assert docky.executable == cmd
     assert docky.cmdline == " ".join(cmd)
 
-    outputs = run_function(docky, tmp_path, worker, environment=Docker(image="busybox"))
+    outputs = run_function(
+        docky, tmp_path, worker, environment=docker.Environment(image="busybox")
+    )
     assert outputs.stdout.strip() == " ".join(cmd[1:])
     assert outputs.return_code == 0
 
@@ -99,7 +107,9 @@ def test_docker_st_1(run_function, worker, tmp_path):
     Docky = shell.define("docky")  # cmd is just a placeholder
     docky = Docky().split(executable=cmd)
 
-    outputs = run_function(docky, tmp_path, worker, environment=Docker(image="busybox"))
+    outputs = run_function(
+        docky, tmp_path, worker, environment=docker.Environment(image="busybox")
+    )
     assert (
         outputs.stdout[0]
         == f"/mnt/pydra{tmp_path}/{attrs.evolve(docky, executable=cmd[0])._checksum}\n"
@@ -121,7 +131,7 @@ def test_docker_outputspec_1(worker, tmp_path):
     Docky = shell.define("touch <out|newfile$newfile_tmp.txt>")
     docky = Docky()
 
-    outputs = docky(worker=worker, environment=Docker(image="ubuntu"))
+    outputs = docky(worker=worker, environment=docker.Environment(image="ubuntu"))
     assert outputs.stdout == ""
 
 
@@ -154,7 +164,9 @@ def test_docker_inputspec_1(tmp_path, worker):
     docky = Docky(file=filename)
 
     outputs = docky(
-        cache_dir=tmp_path, worker=worker, environment=Docker(image="busybox")
+        cache_dir=tmp_path,
+        worker=worker,
+        environment=docker.Environment(image="busybox"),
     )
     assert outputs.stdout.strip() == "hello from pydra"
 
@@ -187,7 +199,7 @@ def test_docker_inputspec_1a(tmp_path):
 
     docky = Docky()
 
-    outputs = docky(cache_dir=tmp_path, environment=Docker(image="busybox"))
+    outputs = docky(cache_dir=tmp_path, environment=docker.Environment(image="busybox"))
     assert outputs.stdout.strip() == "hello from pydra"
 
 
@@ -231,7 +243,7 @@ def test_docker_inputspec_2(worker, tmp_path):
 
     outputs = docky(
         name="docky",
-        environment=Docker(image="busybox"),
+        environment=docker.Environment(image="busybox"),
     )
     assert outputs.stdout.strip() == "hello from pydra\nhave a nice one"
 
@@ -278,7 +290,9 @@ def test_docker_inputspec_2a_except(worker, tmp_path):
     assert docky.file2.fspath == filename_2
 
     outputs = docky(
-        cache_dir=tmp_path, worker=worker, environment=Docker(image="busybox")
+        cache_dir=tmp_path,
+        worker=worker,
+        environment=docker.Environment(image="busybox"),
     )
     assert outputs.stdout.strip() == "hello from pydra\nhave a nice one"
 
@@ -323,7 +337,9 @@ def test_docker_inputspec_2a(worker, tmp_path):
     docky = Docky(file2=filename_2)
 
     outputs = docky(
-        cache_dir=tmp_path, worker=worker, environment=Docker(image="busybox")
+        cache_dir=tmp_path,
+        worker=worker,
+        environment=docker.Environment(image="busybox"),
     )
     assert outputs.stdout.strip() == "hello from pydra\nhave a nice one"
 
@@ -350,7 +366,7 @@ def test_docker_inputspec_3(worker, tmp_path):
     ]
 
     docky = shell.define(cmd, inputs=inputs)(
-        environment=Docker(image="busybox"),
+        environment=docker.Environment(image="busybox"),
         file=filename,
         strip=True,
     )
@@ -391,7 +407,9 @@ class Outputs(shell.Outputs):
     docky = Docky(orig_file=str(file))
 
     outputs = docky(
-        cache_dir=tmp_path, worker=worker, environment=Docker(image="busybox")
+        cache_dir=tmp_path,
+        worker=worker,
+        environment=docker.Environment(image="busybox"),
     )
     assert outputs.stdout == ""
     out_file = outputs.out_file.fspath
@@ -436,7 +454,9 @@ def test_docker_inputspec_state_1(worker, tmp_path):
     docky = Docky().split(file=[str(filename_1), str(filename_2)])
 
     outputs = docky(
-        worker=worker, cache_dir=tmp_path, environment=Docker(image="busybox")
+        worker=worker,
+        cache_dir=tmp_path,
+        environment=docker.Environment(image="busybox"),
     )
     assert outputs.stdout[0].strip() == "hello from pydra"
     assert outputs.stdout[1].strip() == "have a nice one"
@@ -473,7 +493,9 @@ def test_docker_inputspec_state_1b(worker, tmp_path):
     docky = Docky().split(file=[str(file_1), str(file_2)])
 
     outputs = docky(
-        cache_dir=tmp_path, worker=worker, environment=Docker(image="busybox")
+        cache_dir=tmp_path,
+        worker=worker,
+        environment=docker.Environment(image="busybox"),
     )
     assert outputs.stdout[0].strip() == "hello from pydra"
     assert outputs.stdout[1].strip() == "have a nice one"
@@ -507,7 +529,7 @@ def Workflow(file):
 
         docky = workflow.add(
             Docky(file=file),
-            environment=Docker(image="busybox"),
+            environment=docker.Environment(image="busybox"),
         )
 
         return docky.stdout
@@ -549,7 +571,7 @@ def Workflow(file):
 
         docky = workflow.add(
             Docky(file=file),
-            environment=Docker(image="busybox"),
+            environment=docker.Environment(image="busybox"),
         )
 
         return docky.stdout
@@ -593,7 +615,7 @@ def Workflow(file):
 
         docky = workflow.add(
             Docky(file=file),
-            environment=Docker(image="busybox"),
+            environment=docker.Environment(image="busybox"),
         )
 
         return docky.stdout