Implement new step.inputs[].download configuration (#157)

Author: hylje

examples/input-extras.yaml

@@ -14,3 +14,7 @@
         keep-directories: suffix
         default:
           - s3://secret-bucket/bars/**.quux
+      - name: larges
+        download: on-demand
+        default:
+          - s3://foo/bigdata.tar

tests/test_step_parsing.py

@@ -1,7 +1,7 @@
 import pytest
 
 from valohai_yaml.objs import Config
-from valohai_yaml.objs.input import Input, KeepDirectories
+from valohai_yaml.objs.input import DownloadIntent, Input, KeepDirectories
 
 
 def test_parse_inputs(example2_config):
@@ -140,8 +140,10 @@ def test_input_extras(input_extras_config):
     step = config.steps["example"]
     assert step.inputs["model"].keep_directories == KeepDirectories.NONE
     assert step.inputs["model"].filename == "model.pb"
+    assert step.inputs["model"].download == DownloadIntent.ALWAYS
     assert step.inputs["foos"].keep_directories == KeepDirectories.FULL
     assert step.inputs["bars"].keep_directories == KeepDirectories.SUFFIX
+    assert step.inputs["larges"].download == DownloadIntent.ON_DEMAND
 
 
 @pytest.mark.parametrize(
@@ -179,3 +181,16 @@ def test_widget(example1_config: Config) -> None:
     assert widget and widget.type == "datumalias"
     assert widget and widget.settings and widget.settings["width"] == 123
     assert parameters["decoder-spec"].widget is None
+
+
+@pytest.mark.parametrize(
+    "value, expected",
+    [
+        *[(dp, dp) for dp in DownloadIntent],
+        ("always", DownloadIntent.ALWAYS),
+        ("on-demand", DownloadIntent.ON_DEMAND),
+        (None, DownloadIntent.ALWAYS),
+    ],
+)
+def test_download_policy(value, expected):
+    assert Input(name="foo", download=value).download == expected

valohai_yaml/objs/input.py

@@ -25,6 +25,21 @@ def cast(cls, value: KeepDirectoriesValue) -> "KeepDirectories":
         return KeepDirectories(str(value).lower())
 
 
+class DownloadIntent(Enum):
+    """Specify input download intention."""
+
+    ALWAYS = "always"
+    ON_DEMAND = "on-demand"
+
+    @classmethod
+    def cast(cls, value: Union[None, str, "DownloadIntent"]) -> "DownloadIntent":
+        if isinstance(value, DownloadIntent):
+            return value
+        if value is None:
+            return DownloadIntent.ALWAYS
+        return DownloadIntent(str(value).lower())
+
+
 class Input(Item):
     """Represents an input definition within a step definition."""
 
@@ -37,18 +52,24 @@ def __init__(
         description: Optional[str] = None,
         keep_directories: KeepDirectoriesValue = False,
         filename: Optional[str] = None,
+        download: Optional[str] = None,
     ) -> None:
         self.name = name
         self.default = default  # may be None, a string or a list of strings
         self.optional = bool(optional)
         self.description = description
         self.keep_directories = KeepDirectories.cast(keep_directories)
         self.filename = filename
+        self.download = DownloadIntent.cast(download)
 
     def get_data(self) -> SerializedDict:
         data = super().get_data()
         if self.keep_directories is not KeepDirectories.NONE:
             data["keep_directories"] = data["keep_directories"].value
         else:
             data.pop("keep_directories", None)
+        if self.download is not DownloadIntent.ALWAYS:
+            data["download"] = data["download"].value
+        else:
+            data.pop("download", None)
         return data

valohai_yaml/schema/input-item.yaml

@@ -38,3 +38,12 @@ properties:
           - "suffix"
     default: false
     description: Whether to retain directories when using wildcards for this input.
+  download:
+    type: string
+    enum:
+      - "always"
+      - "on-demand"
+    default: "always"
+    description: |
+      Select downloading intention for this input. On-demand inputs are not downloaded before
+      the step can run, instead are available to download using an authenticated download URL.