Preserve workload resources structure on serialization (#133)

Workload resource keys ("cpu", "memory", etc.) need to be preserved so that the UI can correctly parse the resource values.

tests/conftest.py

@@ -11,6 +11,7 @@
 pipeline_config = config_fixture("pipeline-example.yaml")
 task_config = config_fixture("task-example.yaml")
 step_with_resources = config_fixture("step-with-resources.yaml")
+step_with_partial_resources = config_fixture("step-with-partial-resources.yaml")
 pipeline_overridden_config = config_fixture("pipeline-with-override-example.yaml")
 pipeline_with_tasks_config = config_fixture("pipeline-with-tasks-example.yaml")
 pipeline_with_parameters_config = config_fixture(

tests/test_serialize_step.py

@@ -0,0 +1,16 @@
+def test_serialize_workload_resources(step_with_resources):
+    """Must not flatten workload resource data."""
+    config = step_with_resources
+    resources = config.steps["contains kubernetes resources"].resources
+
+    assert isinstance(resources, dict), "Resources should be defined."
+    assert "cpu" in resources, "Resources should contain data."
+
+
+def test_serialize_partial_resources(step_with_partial_resources):
+    """Serialized data only contains keys found in the config."""
+    config = step_with_partial_resources
+    resources = config.steps["contains partial workload resources"].resources
+
+    assert "min" in resources["cpu"]
+    assert "max" not in resources["cpu"]

valohai_yaml/objs/step.py

@@ -96,6 +96,10 @@ def serialize(self) -> OrderedDict:  # type: ignore[type-arg]
                 ("command", self.command),
             ],
         )
+
+        # Some keys need their value structure preserved
+        keys_not_to_flatten = ["resources"]
+
         for key, source in [
             ("parameters", self.parameters),
             ("inputs", self.inputs),
@@ -125,7 +129,7 @@ def serialize(self) -> OrderedDict:  # type: ignore[type-arg]
                 val,
                 key,
                 source,
-                flatten_dicts=True,
+                flatten_dicts=(key not in keys_not_to_flatten),
                 elide_empty_iterables=True,
             )
         return val