Repository: allenai/tango (language: python)

Issue:
Manual specification of step order
### 🚀 The feature, motivation and pitch

As far as I can tell (and I may be wrong about this), the only thing that dictates parent-child relations in the step graph is direct chaining of outputs to inputs via "ref". But it would be nice to have a way to _manually_ enforce that a step (S2) executes after another (S1), even if S1 is not an ancestor of S2 in the graph determined by "ref" relations. Is there a way to do this?

### Alternatives

A naïve solution to this would be to (e.g.) have the `run` method of S2 take the outputs of S1 as an argument and simply not use that argument. But this seems inelegant.

### Additional context

My personal use case motivating this is a set of LM prompting experiments in which one of the steps in the experiment simply prints a set of randomly sampled, dynamically constructed prompts. I would like to see what these prompts look like before model inference runs (as a separate step) in order to be confident the inputs are correct, but currently the prompts print only after the inference step. And the printing step and the model inference step are siblings in the step graph.

I am using the latest tango release (v0.14.0).


Discussion / hints:
Hey @wgantt, this does seem like a good feature, and you're right the currently the only way to enforce a parent-child relationship is through inputs/outputs. I'm thinking we could add another `Step` parameter like `step_requires: List[Union[Step, str]]` that would take a list of steps / step names / step unique IDs that must be run before the current step.

So the config in your example might look something like this:

```jsonnet
{
  steps: {
    print_prompts: { type: "print_prompts" },
    eval_prompts: { type: "eval_prompts", step_requires: [ "print_prompts" ] },
  },
}
```
Yep! At least from a user's perspective, this seems like the most sensible way to do it. 

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8e3f90e01..c6e32e8f3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Added
+
+- Added `step_extra_dependencies` input field to `Step` class that can be used to force a dependency on another step even if the current step doesn't directly depend on the output of the other step. See [#418](https://github.com/allenai/tango/issues/418) for more context.
+
 ### Changed
 
 - `beaker-py >= 1.10` required.
diff --git a/tango/step.py b/tango/step.py
index 8125920ab..4f042108d 100644
--- a/tango/step.py
+++ b/tango/step.py
@@ -120,6 +120,10 @@ class Step(Registrable, Generic[T]):
     :param step_metadata: use this to specify additional metadata for your step.
       This is added to the :attr:`METADATA` class variable to form the ``self.metadata`` attribute.
       Values in ``step_metadata`` take precedence over ``METADATA``.
+    :param step_extra_dependencies: use this to force a dependency on other steps. Normally dependencies
+      between steps are determined by the inputs and outputs of the steps, but you can use this
+      parameter to force that other steps run before this step even if this step doesn't
+      explicitly depend on the outputs of those steps.
 
     .. important::
         Overriding the unique id means that the step will always map to this value, regardless of the inputs,
@@ -177,6 +181,7 @@ def __init__(
         step_unique_id_override: Optional[str] = None,
         step_resources: Optional[StepResources] = None,
         step_metadata: Optional[Dict[str, Any]] = None,
+        step_extra_dependencies: Optional[Iterable["Step"]] = None,
         **kwargs,
     ):
         if self.VERSION is not None:
@@ -258,6 +263,7 @@ def __init__(
         self.metadata = deepcopy(self.METADATA)
         if step_metadata:
             self.metadata.update(step_metadata)
+        self.extra_dependencies = set(step_extra_dependencies) if step_extra_dependencies else set()
 
     @classmethod
     def massage_kwargs(cls, kwargs: Dict[str, Any]) -> Dict[str, Any]:
@@ -638,7 +644,8 @@ def dependencies_internal(o: Any) -> Iterable[Step]:
             else:
                 return
 
-        return dependencies_internal(self.kwargs.values())
+        yield from self.extra_dependencies
+        yield from dependencies_internal(self.kwargs.values())
 
     @property
     def dependencies(self) -> Set["Step"]:


Test patch:
diff --git a/tests/step_graph_test.py b/tests/step_graph_test.py
index 0cf00df96..7ce1a0487 100644
--- a/tests/step_graph_test.py
+++ b/tests/step_graph_test.py
@@ -106,3 +106,24 @@ def test_with_step_indexer(self):
         step_graph = StepGraph.from_params(deepcopy(config))
         assert [s.name for s in step_graph["added"].dependencies] == ["list"]
         assert step_graph.to_config() == config
+
+    def test_with_forced_dependencies(self):
+        config = {
+            "some_string": {
+                "type": "string",
+                "result": "I should run second",
+                "step_extra_dependencies": [{"type": "ref", "ref": "other_string"}],
+            },
+            "other_string": {"type": "string", "result": "I should run first"},
+            "added": {
+                "type": "concat_strings",
+                "string1": "Some string:",
+                "string2": {"type": "ref", "ref": "some_string"},
+            },
+        }
+        step_graph = StepGraph.from_params(deepcopy(config))
+        assert step_graph["some_string"].dependencies == {step_graph["other_string"]}
+        assert step_graph["added"].recursive_dependencies == {
+            step_graph["other_string"],
+            step_graph["some_string"],
+        }