Parse and serialize Step resources as WorkloadResources

tests/test_serialize_step.py

@@ -1,7 +1,7 @@
 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
+    resources = config.steps["contains kubernetes resources"].serialize()["resources"]
 
     assert isinstance(resources, dict), "Resources should be defined."
     assert "cpu" in resources, "Resources should contain data."
@@ -10,7 +10,9 @@ def test_serialize_workload_resources(step_with_resources):
 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
+    resources = config.steps["contains partial workload resources"].serialize()[
+        "resources"
+    ]
 
     assert "min" in resources["cpu"]
     assert "max" not in resources["cpu"]

tests/test_workload_resources.py

@@ -62,7 +62,7 @@ def test_missing_sub_resources(resource_name, missing_key):
     resources = create_resources(resource_name, missing_key)
 
     for this_resource_name, sub_resources in resources.get_data().items():
-        for name, value in sub_resources.get_data().items():
+        for name, value in sub_resources.items():
             if this_resource_name == resource_name and name == missing_key:
                 assert value is None
             else:

valohai_yaml/objs/step.py

@@ -84,6 +84,8 @@ def parse(cls, data: SerializedDict) -> "Step":
         kwargs["source_path"] = kwargs.pop("source-path", None)
         kwargs["stop_condition"] = kwargs.pop("stop-condition", None)
         kwargs["upload_store"] = kwargs.pop("upload-store", None)
+        if "resources" in kwargs:
+            kwargs["resources"] = WorkloadResources.parse(kwargs["resources"])
         inst = cls(**kwargs)
         inst._original_data = data
         return inst
@@ -121,7 +123,10 @@ def serialize(self) -> OrderedDict:  # type: ignore[type-arg]
             ("icon", self.icon),
             ("category", self.category),
             ("source-path", self.source_path),
-            ("resources", self.resources),
+            (
+                "resources",
+                self.resources and self.resources.get_data(),
+            ),
             ("stop-condition", self.stop_condition),
             ("upload-store", self.upload_store),
         ]:

valohai_yaml/objs/workload_resources.py

@@ -19,6 +19,11 @@ def __repr__(self) -> str:
         """CPU data."""
         return f'ResourceCPU("max": {self.max}, "min": {self.min})'
 
+    def get_data(self) -> SerializedDict:
+        return {
+            key: value for key, value in super().get_data().items() if value is not None
+        }
+
 
 class ResourceMemory(Item):
     """Memory configuration."""
@@ -35,6 +40,11 @@ def __repr__(self) -> str:
         """Memory data."""
         return f'ResourceMemory("max": {self.max}, "min": {self.min})'
 
+    def get_data(self) -> SerializedDict:
+        return {
+            key: value for key, value in super().get_data().items() if value is not None
+        }
+
 
 class ResourceDevices(Item):
     """Devices configuration."""
@@ -98,6 +108,13 @@ def parse(cls, data: SerializedDict) -> "WorkloadResources":
         )
         return super().parse(data_with_resources)
 
+    def get_data(self) -> SerializedDict:
+        return {
+            key: value.get_data()
+            for key, value in super().get_data().items()
+            if value is not None
+        }
+
     def __repr__(self) -> str:
         """Resources contents."""
         return f'WorkloadResources("cpu": {self.cpu}, "memory": {self.memory}, "devices": {self.devices})'