Repository: allenai/tango (language: python)

Issue:
Support Indexing for multiple return values from a Tango Step
### 🚀 The feature, motivation and pitch

If a Tango step returns multiple values, we currently have to pass all the returned values to another Step. The goal is to support indexing so that the user can select which returned values should be passed as an input to another Step. 

### Alternatives

_No response_

### Additional context

_No response_


Discussion / hints:
One alternative worth mentioning (which is currently possible) is to create an intermediate step which takes any kind of `Sequence` and just returns the item corresponding to a given index. So your config could look like:

```jsonnet
sequence: {type: "sequence_step"},  # pretend this step returns a sequence of whatever
second_item: {type: "get_index", index: 1},  # this is the intermediate 'getter' step
final_step: {type: "whatever", input: {type: "ref", ref: second_item}},
```
I would like to second @isunitha98selvan's feature request for being able to index multiple returned values. 

In the following example, `InferenceStep` needs to refer to individual outputs from `ModelCreationStep`. I know there is a workaround here - wrapping model and params into a single return object - but would be nice to have the flexibility!

```python
@Step.register('init_model')
class ModelCreationStep(Step):
    DETERMINISTIC = False
    CACHEABLE = False

    def run(self,image_size):
        model = Model(image_size=image_size)
        params = get_params(model)
        return model, params

@Step.register('infer')
class InferenceStep(Step):
    DETERMINISTIC = False
    CACHEABLE = False

    def run(self,model,params,input):
        output = model.apply(params,input)
        return output
```

@epwalsh 
Gotcha.. Hmm, in the JSON config, I'm thinking we'd support it this way:

```jsonnet
{
  steps: {
    model_and_params: { type: "init_model" },
    predictions: { type: "infer", model: { type: "ref", ref: "model_and_params", key: 0 }, ... },
  }
}
```

The main thing here is the addition of a "key" field in the ref step object. The key could also be a string so we could pull objects out of dictionaries.

What do you think @BigRedT @AkshitaB @dirkgr?
That would work for me!

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3e730403e..44f3fdf6e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- You can now reference into a particular index of the result of another step in a config. For example: `{type: "ref", ref: "some_previous_step", key: 0}`.
+  The key field can be an integer if the result of the referenced step is a list or tuple, or a string if the result of the referenced step is a dictionary.
 - Added `priority` parameter to Beaker executor for setting the default task priority for Beaker jobs.
 - Added `Workspace.step_result()` method for getting a step's result from the latest
   run.
diff --git a/tango/common/from_params.py b/tango/common/from_params.py
index 0e5f6a4c8..583be7d35 100644
--- a/tango/common/from_params.py
+++ b/tango/common/from_params.py
@@ -387,7 +387,7 @@ def construct_arg(
     if popped_params is default:
         return popped_params
 
-    from tango.step import Step, WithUnresolvedSteps
+    from tango.step import Step, StepIndexer, WithUnresolvedSteps
 
     origin = getattr(annotation, "__origin__", None)
     args = getattr(annotation, "__args__", [])
@@ -481,6 +481,18 @@ def construct_arg(
         else:
             return default
 
+    # For StepIndexer, we just return as-is and hope the for the best.
+    # At worst, user will get an error at runtime if they are trying to index a step
+    # result that can't be indexed.
+    # TODO (epwalsh): we could check the return type of the wrapped step here
+    # and make sure that:
+    #  1. It's an index-able object,
+    #  2. The item in the index-able object matches `annotation`.
+    #
+    # But that's complex and might have false negatives.
+    elif type(popped_params) == StepIndexer:
+        return popped_params
+
     # If the parameter type is a Python primitive, just pop it off
     # using the correct casting pop_xyz operation.
     elif annotation in {int, bool}:
diff --git a/tango/common/util.py b/tango/common/util.py
index 3a7e219dc..9bc5836d8 100644
--- a/tango/common/util.py
+++ b/tango/common/util.py
@@ -269,11 +269,13 @@ def local_timezone() -> Optional[tzinfo]:
 
 
 def replace_steps_with_unique_id(o: Any):
-    from tango.step import Step
+    from tango.step import Step, StepIndexer
 
     if isinstance(o, Step):
         return {"type": "ref", "ref": o.unique_id}
-    if isinstance(o, (list, tuple, set)):
+    elif isinstance(o, StepIndexer):
+        return {"type": "ref", "ref": o.step.unique_id, "key": o.key}
+    elif isinstance(o, (list, tuple, set)):
         return o.__class__(replace_steps_with_unique_id(i) for i in o)
     elif isinstance(o, dict):
         return {key: replace_steps_with_unique_id(value) for key, value in o.items()}
diff --git a/tango/step.py b/tango/step.py
index 969646b2c..4d7a371ba 100644
--- a/tango/step.py
+++ b/tango/step.py
@@ -541,8 +541,8 @@ def __eq__(self, other):
             return False
 
     def _replace_steps_with_results(self, o: Any, workspace: "Workspace"):
-        if isinstance(o, Step):
-            return o.result(workspace, self)
+        if isinstance(o, (Step, StepIndexer)):
+            return o.result(workspace=workspace, needed_by=self)
         elif isinstance(o, Lazy):
             return Lazy(
                 o._constructor,
@@ -623,6 +623,8 @@ def dependencies_internal(o: Any) -> Iterable[Step]:
             elif isinstance(o, WithUnresolvedSteps):
                 yield from dependencies_internal(o.args)
                 yield from dependencies_internal(o.kwargs)
+            elif isinstance(o, StepIndexer):
+                yield o.step
             elif isinstance(o, str):
                 return  # Confusingly, str is an Iterable of itself, resulting in infinite recursion.
             elif isinstance(o, (dict, Params)):
@@ -701,6 +703,17 @@ def log_failure(self, exception: BaseException) -> None:
         cli_logger.error('[red]\N{ballot x} Step [bold]"%s"[/] failed[/]', self.name)
 
 
+class StepIndexer:
+    def __init__(self, step: Step, key: Union[str, int]):
+        self.step = step
+        self.key = key
+
+    def result(
+        self, workspace: Optional["Workspace"] = None, needed_by: Optional["Step"] = None
+    ) -> Any:
+        return self.step.result(workspace=workspace, needed_by=needed_by)[self.key]
+
+
 class WithUnresolvedSteps(CustomDetHash):
     """
     This is a helper class for some scenarios where steps depend on other steps.
@@ -794,8 +807,8 @@ def with_resolved_steps(
         :return: A new object that's a copy of the original object, with all instances of :class:`.Step` replaced
                  with the results of the step.
         """
-        if isinstance(o, Step):
-            return o.result(workspace)
+        if isinstance(o, (Step, StepIndexer)):
+            return o.result(workspace=workspace)
         elif isinstance(o, Lazy):
             return Lazy(
                 o._constructor,
diff --git a/tango/step_graph.py b/tango/step_graph.py
index 22b9b2ffc..ab3cceb66 100644
--- a/tango/step_graph.py
+++ b/tango/step_graph.py
@@ -4,7 +4,7 @@
 from tango.common import PathOrStr
 from tango.common.exceptions import ConfigurationError
 from tango.common.params import Params
-from tango.step import Step
+from tango.step import Step, StepIndexer
 
 logger = logging.getLogger(__name__)
 
@@ -142,7 +142,7 @@ def _dict_is_ref(d: Union[dict, Params]) -> bool:
         keys = set(d.keys())
         if keys == {"ref"}:
             return True
-        if keys == {"type", "ref"} and d["type"] == "ref":
+        if keys >= {"type", "ref"} and d["type"] == "ref":
             return True
         return False
 
@@ -168,7 +168,10 @@ def _replace_step_dependencies(cls, o: Any, existing_steps: Mapping[str, Step])
             return o.__class__(cls._replace_step_dependencies(i, existing_steps) for i in o)
         elif isinstance(o, (dict, Params)):
             if cls._dict_is_ref(o):
-                return existing_steps[o["ref"]]
+                if "key" in o:
+                    return StepIndexer(existing_steps[o["ref"]], o["key"])
+                else:
+                    return existing_steps[o["ref"]]
             else:
                 return {
                     key: cls._replace_step_dependencies(value, existing_steps)
@@ -239,6 +242,8 @@ def _to_config(o: Any):
                 return {key: _to_config(value) for key, value in o.items()}
             elif isinstance(o, Step):
                 return {"type": "ref", "ref": o.name}
+            elif isinstance(o, StepIndexer):
+                return {"type": "ref", "ref": o.step.name, "key": o.key}
             elif o is not None and not isinstance(o, (bool, str, int, float)):
                 raise ValueError(o)
             return o


Test patch:
diff --git a/tango/common/testing/steps.py b/tango/common/testing/steps.py
index 5dd8fef1a..03de97f2b 100644
--- a/tango/common/testing/steps.py
+++ b/tango/common/testing/steps.py
@@ -3,6 +3,7 @@
 import random
 import time
 from string import ascii_letters
+from typing import List
 
 import tango.common.logging as common_logging
 from tango import Step
@@ -143,6 +144,12 @@ def run(self, num_proc: int = 2) -> bool:  # type: ignore
         return True
 
 
+@Step.register("range_step")
+class RangeOutput(Step):
+    def run(self, start: int, end: int) -> List[int]:  # type: ignore
+        return list(range(start, end))
+
+
 def _worker_function(worker_id: int):
     common_logging.initialize_worker_logging(worker_id)
     logger = logging.getLogger(MultiprocessingStep.__name__)
diff --git a/tests/end_to_end/test_step_indexing.py b/tests/end_to_end/test_step_indexing.py
new file mode 100644
index 000000000..f1be2b49b
--- /dev/null
+++ b/tests/end_to_end/test_step_indexing.py
@@ -0,0 +1,21 @@
+from tango.common.testing import TangoTestCase
+from tango.workspaces import LocalWorkspace
+
+
+class TestStepIndexing(TangoTestCase):
+    def test_step_indexing(self):
+        run_name = "run1"
+        config = {
+            "steps": {
+                "list": {"type": "range_step", "start": 0, "end": 3},
+                "added": {
+                    "type": "add_numbers",
+                    "a_number": 2,
+                    "b_number": {"type": "ref", "ref": "list", "key": 1},
+                },
+            }
+        }
+        self.run(config, name=run_name)
+        workspace = LocalWorkspace(self.TEST_DIR / "workspace")
+        result = workspace.step_result_for_run(run_name, "added")
+        assert result == 3
diff --git a/tests/step_graph_test.py b/tests/step_graph_test.py
index a5d5c393f..0cf00df96 100644
--- a/tests/step_graph_test.py
+++ b/tests/step_graph_test.py
@@ -1,4 +1,5 @@
 import re
+from copy import deepcopy
 from tempfile import NamedTemporaryFile
 
 import pytest
@@ -92,3 +93,16 @@ def test_to_file_without_config(self):
             step_graph.to_file(file_ref.name)
             new_step_graph = StepGraph.from_file(file_ref.name)
             assert step_graph == new_step_graph
+
+    def test_with_step_indexer(self):
+        config = {
+            "list": {"type": "range_step", "start": 0, "end": 3},
+            "added": {
+                "type": "add_numbers",
+                "a_number": 2,
+                "b_number": {"type": "ref", "ref": "list", "key": 1},
+            },
+        }
+        step_graph = StepGraph.from_params(deepcopy(config))
+        assert [s.name for s in step_graph["added"].dependencies] == ["list"]
+        assert step_graph.to_config() == config