Enable multiprocessing groups within project config

spacy/cli/_util.py

@@ -1,5 +1,5 @@
 from typing import Dict, Any, Union, List, Optional, Tuple, Iterable
-from typing import TYPE_CHECKING, overload
+from typing import TYPE_CHECKING, overload, cast
 import sys
 import shutil
 from pathlib import Path
@@ -215,30 +215,54 @@ def validate_project_version(config: Dict[str, Any]) -> None:
         msg.fail(err, exits=1)
 
 
+def verify_workflow_step(workflow_name: str, commands: List[str], step: str) -> None:
+    if step not in commands:
+        msg.fail(
+            f"Unknown command specified in workflow '{workflow_name}': {step}",
+            f"Workflows can only refer to commands defined in the 'commands' "
+            f"section of the {PROJECT_FILE}.",
+            exits=1,
+        )
+
+
 def validate_project_commands(config: Dict[str, Any]) -> None:
     """Check that project commands and workflows are valid, don't contain
     duplicates, don't clash  and only refer to commands that exist.
 
     config (Dict[str, Any]): The loaded config.
     """
-    command_names = [cmd["name"] for cmd in config.get("commands", [])]
+
+    commands = [cmd["name"] for cmd in config.get("commands", [])]
     workflows = config.get("workflows", {})
-    duplicates = set([cmd for cmd in command_names if command_names.count(cmd) > 1])
+    duplicates = set([cmd for cmd in commands if commands.count(cmd) > 1])
     if duplicates:
         err = f"Duplicate commands defined in {PROJECT_FILE}: {', '.join(duplicates)}"
         msg.fail(err, exits=1)
-    for workflow_name, workflow_steps in workflows.items():
-        if workflow_name in command_names:
+    for workflow_name, workflow_items in workflows.items():
+        if workflow_name in commands:
             err = f"Can't use workflow name '{workflow_name}': name already exists as a command"
             msg.fail(err, exits=1)
-        for step in workflow_steps:
-            if step not in command_names:
-                msg.fail(
-                    f"Unknown command specified in workflow '{workflow_name}': {step}",
-                    f"Workflows can only refer to commands defined in the 'commands' "
-                    f"section of the {PROJECT_FILE}.",
-                    exits=1,
-                )
+        for workflow_item in workflow_items:
+            if isinstance(workflow_item, str):
+                verify_workflow_step(workflow_name, commands, workflow_item)
+            else:
+                assert isinstance(workflow_item, list)
+                assert isinstance(workflow_item[0], str)
+                steps = cast(List[str], workflow_item)
+                if len(steps) < 2:
+                    msg.fail(
+                        f"Invalid multiprocessing group within '{workflow_name}'.",
+                        f"A multiprocessing group must reference at least two commands.",
+                        exits=1,
+                    )
+                if len(steps) != len(set(steps)):
+                    msg.fail(
+                        f"A multiprocessing group within '{workflow_name}' contains a command more than once.",
+                        f"This is not permitted because it is then not possible to determine when to rerun.",
+                        exits=1,
+                    )
+                for step in steps:
+                    verify_workflow_step(workflow_name, commands, step)
 
 
 def get_hash(data, exclude: Iterable[str] = tuple()) -> str:
@@ -556,3 +580,15 @@ def setup_gpu(use_gpu: int, silent=None) -> None:
         local_msg.info("Using CPU")
         if has_cupy and gpu_is_available():
             local_msg.info("To switch to GPU 0, use the option: --gpu-id 0")
+
+
+def get_workflow_steps(workflow_items: List[Union[str, List[str]]]) -> List[str]:
+    steps: List[str] = []
+    for workflow_item in workflow_items:
+        if isinstance(workflow_item, str):
+            steps.append(workflow_item)
+        else:
+            assert isinstance(workflow_item, list)
+            assert isinstance(workflow_item[0], str)
+            steps.extend(workflow_item)
+    return steps

spacy/cli/project/document.py

@@ -14,8 +14,8 @@
 Commands are only re-run if their inputs have changed."""
 INTRO_WORKFLOWS = f"""The following workflows are defined by the project. They
 can be executed using [`spacy project run [name]`]({DOCS_URL}/api/cli#project-run)
-and will run the specified commands in order. Commands are only re-run if their
-inputs have changed."""
+and will run the specified commands in order. Commands grouped within square brackets
+are run in parallel. Commands are only re-run if their inputs have changed."""
 INTRO_ASSETS = f"""The following assets are defined by the project. They can
 be fetched by running [`spacy project assets`]({DOCS_URL}/api/cli#project-assets)
 in the project directory."""
@@ -69,7 +69,17 @@ def project_document(
         md.add(md.table(data, ["Command", "Description"]))
     # Workflows
     wfs = config.get("workflows", {}).items()
-    data = [(md.code(n), " → ".join(md.code(w) for w in stp)) for n, stp in wfs]
+    data = []
+    for n, steps in wfs:
+        rendered_steps = []
+        for step in steps:
+            if isinstance(step, str):
+                rendered_steps.append(md.code(step))
+            else:
+                rendered_steps.append(
+                    "[" + ", ".join(md.code(p_step) for p_step in step) + "]"
+                )
+        data.append([md.code(n), " → ".join(rendered_steps)])
     if data:
         md.add(md.title(3, "Workflows", "⏭"))
         md.add(INTRO_WORKFLOWS)

spacy/cli/project/dvc.py

@@ -1,12 +1,12 @@
 """This module contains helpers and subcommands for integrating spaCy projects
-with Data Version Controk (DVC). https://dvc.org"""
+with Data Version Control (DVC). https://dvc.org"""
 from typing import Dict, Any, List, Optional, Iterable
 import subprocess
 from pathlib import Path
 from wasabi import msg
 
 from .._util import PROJECT_FILE, load_project_config, get_hash, project_cli
-from .._util import Arg, Opt, NAME, COMMAND
+from .._util import Arg, Opt, NAME, COMMAND, get_workflow_steps
 from ...util import working_dir, split_command, join_command, run_command
 from ...util import SimpleFrozenList
 
@@ -105,13 +105,15 @@ def update_dvc_config(
         dvc_config_path.unlink()
     dvc_commands = []
     config_commands = {cmd["name"]: cmd for cmd in config.get("commands", [])}
-    for name in workflows[workflow]:
+    processed_step = False
+    for name in get_workflow_steps(workflows[workflow]):
         command = config_commands[name]
         deps = command.get("deps", [])
         outputs = command.get("outputs", [])
         outputs_no_cache = command.get("outputs_no_cache", [])
         if not deps and not outputs and not outputs_no_cache:
             continue
+        processed_step = True
         # Default to the working dir as the project path since dvc.yaml is auto-generated
         # and we don't want arbitrary paths in there
         project_cmd = ["python", "-m", NAME, "project", "run", name]
@@ -123,6 +125,11 @@ def update_dvc_config(
             dvc_cmd.append("--always-changed")
         full_cmd = [*dvc_cmd, *deps_cmd, *outputs_cmd, *outputs_nc_cmd, *project_cmd]
         dvc_commands.append(join_command(full_cmd))
+    if not processed_step:
+        msg.fail(
+            f"A DVC workflow must have at least one dependency or output",
+            exits=1,
+        )
     with working_dir(path):
         dvc_flags = {"--verbose": verbose, "--quiet": silent}
         run_dvc_commands(dvc_commands, flags=dvc_flags)

spacy/cli/project/pull.py

@@ -1,3 +1,4 @@
+import multiprocessing
 from pathlib import Path
 from wasabi import msg
 from .remote_storage import RemoteStorage
@@ -37,6 +38,7 @@ def project_pull(project_dir: Path, remote: str, *, verbose: bool = False):
     # We use a while loop here because we don't know how the commands
     # will be ordered. A command might need dependencies from one that's later
     # in the list.
+    mult_group_mutex = multiprocessing.Lock()
     while commands:
         for i, cmd in enumerate(list(commands)):
             logger.debug(f"CMD: {cmd['name']}.")
@@ -52,7 +54,7 @@ def project_pull(project_dir: Path, remote: str, *, verbose: bool = False):
 
                 out_locs = [project_dir / out for out in cmd.get("outputs", [])]
                 if all(loc.exists() for loc in out_locs):
-                    update_lockfile(project_dir, cmd)
+                    update_lockfile(project_dir, cmd, mult_group_mutex=mult_group_mutex)
                 # We remove the command from the list here, and break, so that
                 # we iterate over the loop again.
                 commands.pop(i)