feat(resources): support serialization of Resources dataclasses (#2250)

* feat(resources): inherit from from DataClassJSONMixin

- support serialization of Resources and ResourceSpec
- allow passing these as configuration to workflows or between tasks
Signed-off-by: Cameron Smith <[email protected]>

* test(resources): add tests to (de)serialize Resources instances
Signed-off-by: Cameron Smith <[email protected]>
Signed-off-by: Jan Fiedler <[email protected]>

flytekit/core/resources.py

@@ -1,11 +1,13 @@
 from dataclasses import dataclass
 from typing import List, Optional
 
+from mashumaro.mixins.json import DataClassJSONMixin
+
 from flytekit.models import task as task_models
 
 
 @dataclass
-class Resources(object):
+class Resources(DataClassJSONMixin):
     """
     This class is used to specify both resource requests and resource limits.
 
@@ -45,7 +47,7 @@ def _check_none_or_str(value):
 
 
 @dataclass
-class ResourceSpec(object):
+class ResourceSpec(DataClassJSONMixin):
     requests: Resources
     limits: Resources
 

tests/flytekit/unit/core/test_resources.py

@@ -75,3 +75,29 @@ def test_incorrect_type_resources():
         Resources(gpu=1)  # type: ignore
     with pytest.raises(AssertionError):
         Resources(ephemeral_storage=1)  # type: ignore
+
+
+def test_resources_serialization():
+    resources = Resources(cpu="2", mem="1Gi", gpu="1", ephemeral_storage="10Gi")
+    json_str = resources.to_json()
+    assert isinstance(json_str, str)
+    assert '"cpu": "2"' in json_str
+    assert '"mem": "1Gi"' in json_str
+    assert '"gpu": "1"' in json_str
+    assert '"ephemeral_storage": "10Gi"' in json_str
+
+
+def test_resources_deserialization():
+    json_str = '{"cpu": "2", "mem": "1Gi", "gpu": "1", "ephemeral_storage": "10Gi"}'
+    resources = Resources.from_json(json_str)
+    assert resources.cpu == "2"
+    assert resources.mem == "1Gi"
+    assert resources.gpu == "1"
+    assert resources.ephemeral_storage == "10Gi"
+
+
+def test_resources_round_trip():
+    original = Resources(cpu="4", mem="2Gi", gpu="2", ephemeral_storage="20Gi")
+    json_str = original.to_json()
+    result = Resources.from_json(json_str)
+    assert original == result