Reformat spawner resources entities

src/glicid_spawner/resources.py

@@ -1,35 +1,39 @@
-"""Resources module."""
+"""Default resources module."""
 
-from dataclasses import dataclass
+CPU = {
+    # description: max-duration [h]
+    1: 24,
+    2: 12,
+    4: 6,
+    8: 3,
+    12: 2,
+    24: 1,
+}
+
+MEMORY = {  # [GB]
+    4: 24,
+    8: 12,
+    16: 6,
+    32: 3,
+    48: 2,
+    96: 1,
+}
+
+GPU_DEFAULTS = {
+    'None': 24,
+    'A40': 2,
+    'A100': 1,
+}
 
 
-@dataclass
-class Resource:
-    """Generic cluster resource."""
+def gpu_max_duration(gpus: list, unknown_default=1) -> dict:
+    """GPU max-duration allocation.
 
-    description: str
-    max_duration: int
+    By default 1h is allocated to any unknown GPU resource
+    not listed in the GPU_DEFAULTS.
 
-
-CPU = [
-    Resource('1', 24),
-    Resource('2', 12),
-    Resource('4', 6),
-    Resource('8', 3),
-    Resource('12', 2),
-    Resource('24', 1),
-]
-
-RAM = [
-    Resource('4 GB', 24),
-    Resource('8 GB', 12),
-    Resource('16 GB', 6),
-    Resource('32 GB', 3),
-    Resource('48 GB', 2),
-    Resource('96 GB', 1),
-]
-
-GPU = [
-    Resource('No', 24),
-    Resource('A100', 1),
-]
+    """
+    # Filter defaults value in the same order as GPUS_DEFAULTS
+    defaults = {gpu: duration for gpu, duration in GPU_DEFAULTS.items() if gpu in gpus}
+    unknowns = {gpu: unknown_default for gpu in gpus if gpu not in GPU_DEFAULTS}
+    return defaults | unknowns