From b8efa00a05578bbe4d677eb461b883d5f0a9f07e Mon Sep 17 00:00:00 2001
From: Benoit Seignovert <benoit.seignovert@univ-nantes.fr>
Date: Tue, 20 Feb 2024 11:18:46 +0100
Subject: [PATCH] Reformat spawner resources entities

---
 src/glicid_spawner/resources.py | 64 +++++++++++++++++----------------
 tests/test_resources.py         | 36 +++++++++----------
 2 files changed, 52 insertions(+), 48 deletions(-)

diff --git a/src/glicid_spawner/resources.py b/src/glicid_spawner/resources.py
index e5be588..b4cdc52 100644
--- a/src/glicid_spawner/resources.py
+++ b/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
diff --git a/tests/test_resources.py b/tests/test_resources.py
index f9c695e..482ba1c 100644
--- a/tests/test_resources.py
+++ b/tests/test_resources.py
@@ -1,24 +1,24 @@
 """Test resources module."""
 
-from glicid_spawner.resources import CPU, GPU, RAM
+from glicid_spawner.resources import CPU, GPU_DEFAULTS, MEMORY, gpu_max_duration
 
 
-def test_glicid_default_resources():
-    """Test GLiCID default resources."""
-    assert len(CPU) == 6
-    assert CPU[0].description == '1'
-    assert CPU[0].max_duration == 24
-    assert CPU[-1].description == '24'
-    assert CPU[-1].max_duration == 1
+def test_default_resources():
+    """Test default resources duration default allocations."""
+    assert list(CPU) == [1, 2, 4, 8, 12, 24]
+    assert list(CPU.values()) == [24, 12, 6, 3, 2, 1]
 
-    assert len(RAM) == 6
-    assert RAM[0].description == '4 GB'
-    assert RAM[0].max_duration == 24
-    assert RAM[-1].description == '96 GB'
-    assert RAM[-1].max_duration == 1
+    assert list(MEMORY) == [4, 8, 16, 32, 48, 96]
+    assert list(MEMORY.values()) == [24, 12, 6, 3, 2, 1]
 
-    assert len(GPU) == 2
-    assert GPU[0].description == 'No'
-    assert GPU[0].max_duration == 24
-    assert GPU[-1].description == 'A100'
-    assert GPU[-1].max_duration == 1
+    assert list(GPU_DEFAULTS) == ['None', 'A40', 'A100']
+    assert list(GPU_DEFAULTS.values()) == [24, 2, 1]
+
+
+def test_gpu_max_duration():
+    """Test GPU max duration allocation."""
+    gpu = gpu_max_duration(['T4', 'A100', 'None', 'K80', 'A100', 'T4'], unknown_default=3)
+
+    # Sorted by defaults order, then unknowns
+    assert list(gpu) == ['None', 'A100', 'T4', 'K80']
+    assert list(gpu.values()) == [24, 1, 3, 3]