Simply data

aws-deadline · Dec 20, 2024 · 1ff0567 · 1ff0567
1 parent ee1805a
commit 1ff0567
Showing 1 changed file with 52 additions and 82 deletions.
diff --git a/src/deadline/client/ui/dataclasses/__init__.py b/src/deadline/client/ui/dataclasses/__init__.py
@@ -133,66 +133,26 @@ class HardwareRequirements:
     scratch_space_min: int = field(default=DEFAULT_VALUE)
     scratch_space_max: int = field(default=DEFAULT_VALUE)
 
-    def serialize(self) -> list:
+    def _serialize(self, minimum, maximum, name):
         requirements: List[dict] = []
-        if self.cpu_min != self.DEFAULT_VALUE or self.cpu_max != self.DEFAULT_VALUE:
-            _requirement = {
-                "name": "amount.worker.vcpu",
-            }
-            if self.cpu_min != self.DEFAULT_VALUE:
-                _requirement["min"] = max(self.cpu_min, self.DEFAULT_VALUE)
-            if self.cpu_max != self.DEFAULT_VALUE:
-                _requirement["max"] = min(self.cpu_max, self.MAX_INT_VALUE)
-            requirements.append(_requirement)
-
-        if self.memory_min != self.DEFAULT_VALUE or self.memory_max != self.DEFAULT_VALUE:
-            _requirement = {
-                "name": "amount.worker.memory",
-            }
-            if self.memory_min != self.DEFAULT_VALUE:
-                _requirement["min"] = max(self.memory_min, self.DEFAULT_VALUE)
-            if self.memory_max != 0:
-                _requirement["max"] = min(self.memory_max, self.MAX_INT_VALUE)
-            requirements.append(_requirement)
-
-        if (
-            self.acceleration_min != self.DEFAULT_VALUE
-            or self.acceleration_max != self.DEFAULT_VALUE
-        ):
+        if minimum != self.DEFAULT_VALUE or maximum != self.DEFAULT_VALUE:
             _requirement = {
-                "name": "amount.worker.gpu",
+                "name": name,
             }
-            if self.acceleration_min != self.DEFAULT_VALUE:
-                _requirement["min"] = max(self.acceleration_min, self.DEFAULT_VALUE)
-            if self.acceleration_max != self.DEFAULT_VALUE:
-                _requirement["max"] = min(self.acceleration_max, self.MAX_INT_VALUE)
-            requirements.append(_requirement)
-
-        if (
-            self.acceleration_memory_min != self.DEFAULT_VALUE
-            or self.acceleration_memory_max != self.DEFAULT_VALUE
-        ):
-            _requirement = {
-                "name": "amount.worker.gpu.memory",
-            }
-            if self.acceleration_memory_min != self.DEFAULT_VALUE:
-                _requirement["min"] = max(self.acceleration_memory_min, self.DEFAULT_VALUE)
-            if self.acceleration_memory_max != self.DEFAULT_VALUE:
-                _requirement["max"] = min(self.acceleration_memory_max, self.MAX_INT_VALUE)
+            if minimum != self.DEFAULT_VALUE:
+                _requirement["min"] = max(minimum, self.DEFAULT_VALUE)
+            if maximum != self.DEFAULT_VALUE:
+                _requirement["max"] = min(maximum, self.MAX_INT_VALUE)
             requirements.append(_requirement)
+        return requirements
 
-        if (
-            self.scratch_space_min != self.DEFAULT_VALUE
-            or self.scratch_space_max != self.DEFAULT_VALUE
-        ):
-            _requirement = {
-                "name": "amount.worker.disk.scratch",
-            }
-            if self.scratch_space_min != self.DEFAULT_VALUE:
-                _requirement["min"] = max(self.scratch_space_min, self.DEFAULT_VALUE)
-            if self.scratch_space_max != self.DEFAULT_VALUE:
-                _requirement["max"] = min(self.scratch_space_max, self.MAX_INT_VALUE)
-            requirements.append(_requirement)
+    def serialize(self) -> list:
+        requirements: List[dict] = []
+        requirements.extend(self._serialize(self.cpu_min, self.cpu_max, "amount.worker.vcpu"))
+        requirements.extend(self._serialize(self.memory_min, self.memory_max, "amount.worker.memory"))
+        requirements.extend(self._serialize(self.acceleration_min, self.acceleration_max, "amount.worker.gpu"))
+        requirements.extend(self._serialize(self.acceleration_memory_min, self.acceleration_memory_max, "amount.worker.gpu.memory"))
+        requirements.extend(self._serialize(self.scratch_space_min, self.scratch_space_max, "amount.worker.disk.scratch"))
         return requirements
 
 
@@ -285,39 +245,49 @@ def __iter__(self):
         if self.attributes:
             yield from self.attributes
 
+    def _serialize_amounts(self):
+        requirements = []
+        for amount in self.amounts:
+            if not amount.name:
+                raise ValueError(f"Custom Amount {amount} has no name")
+            _requirement = {
+                "name": "amount.worker." + amount.name,
+            }
+            if amount.min != CustomAmountRequirement.DEFAULT_VALUE:
+                _requirement["min"] = max(amount.min, CustomAmountRequirement.DEFAULT_VALUE)
+            if amount.max != CustomAmountRequirement.DEFAULT_VALUE:
+                _requirement["max"] = min(amount.max, CustomAmountRequirement.MAX_INT_VALUE)
+            requirements.append(_requirement)
+        return requirements
+
+    def _serialize_attributes(self):
+        requirements = []
+        for attribute in self.attributes:
+            if not attribute.name:
+                raise ValueError(f"Custom Attribute {attribute} has no name")
+            _requirement = {
+                "name": "attr.worker." + attribute.name,
+            }
+            if not attribute.option:
+                raise ValueError(f"Custom Attribute {attribute} has no option")
+            elif attribute.option not in ["anyOf", "allOf"]:
+                raise ValueError(
+                    f'Custom Attribute {attribute} has option {attribute.option} which is not in list: ["anyOf", "allOf"]'
+                )
+            if not attribute.values:
+                raise ValueError(f"Custom Attribute {attribute} has no values")
+            _requirement[attribute.option] = attribute.values
+            requirements.append(_requirement)
+        return requirements
+
     def serialize(self) -> Dict[str, List]:
         requirements: Dict[str, List] = {}
 
         if self.amounts:
-            for amount in self.amounts:
-                if not amount.name:
-                    raise ValueError(f"Custom Amount {amount} has no name")
-                _requirement = {
-                    "name": "amount.worker." + amount.name,
-                }
-                if amount.min != CustomAmountRequirement.DEFAULT_VALUE:
-                    _requirement["min"] = max(amount.min, CustomAmountRequirement.DEFAULT_VALUE)
-                if amount.max != CustomAmountRequirement.DEFAULT_VALUE:
-                    _requirement["max"] = min(amount.max, CustomAmountRequirement.MAX_INT_VALUE)
-                requirements.setdefault("amounts", []).append(_requirement)
+            requirements.setdefault("amounts", []).extend(self._serialize_amounts())
 
         if self.attributes:
-            for attribute in self.attributes:
-                if not attribute.name:
-                    raise ValueError(f"Custom Attribute {attribute} has no name")
-                _requirement = {
-                    "name": "attr.worker." + attribute.name,
-                }
-                if not attribute.option:
-                    raise ValueError(f"Custom Attribute {attribute} has no option")
-                elif attribute.option not in ["anyOf", "allOf"]:
-                    raise ValueError(
-                        f'Custom Attribute {attribute} has option {attribute.option} which is not in list: ["anyOf", "allOf"]'
-                    )
-                if not attribute.values:
-                    raise ValueError(f"Custom Attribute {attribute} has no values")
-                _requirement[attribute.option] = attribute.values
-                requirements.setdefault("attributes", []).append(_requirement)
+            requirements.setdefault("attributes", []).extend(self._serialize_attributes())
         return requirements