feat(config): add partition in application.yaml and use it in get_load

AntaresSimulatorTeam · Jul 25, 2023 · 0c948e8 · 0c948e8
1 parent d549e16
commit 0c948e8
Show file tree

Hide file tree

Showing 7 changed files with 15 additions and 6 deletions.
diff --git a/antarest/core/config.py b/antarest/core/config.py
@@ -185,6 +185,7 @@ class SlurmConfig:
     default_n_cpu: int = 1
     default_json_db_name: str = ""
     slurm_script_path: str = ""
+    partition: str = ""
     max_cores: int = 64
     antares_versions_on_remote_server: List[str] = field(
         default_factory=lambda: []
@@ -205,6 +206,7 @@ def from_dict(data: JSON) -> "SlurmConfig":
             default_n_cpu=data["default_n_cpu"],
             default_json_db_name=data["default_json_db_name"],
             slurm_script_path=data["slurm_script_path"],
+            partition=data["partition"],
             antares_versions_on_remote_server=data[
                 "antares_versions_on_remote_server"
             ],

diff --git a/antarest/launcher/adapters/slurm_launcher/slurm_launcher.py b/antarest/launcher/adapters/slurm_launcher/slurm_launcher.py
@@ -225,6 +225,7 @@ def _init_launcher_parameters(
             json_dir=local_workspace or self.slurm_config.local_workspace,
             default_json_db_name=self.slurm_config.default_json_db_name,
             slurm_script_path=self.slurm_config.slurm_script_path,
+            partition=self.slurm_config.partition,
             antares_versions_on_remote_server=self.slurm_config.antares_versions_on_remote_server,
             default_ssh_dict={
                 "username": self.slurm_config.username,

diff --git a/antarest/launcher/service.py b/antarest/launcher/service.py
@@ -708,7 +708,8 @@ def get_load(self) -> LauncherLoadDTO:
                     key_password=slurm_config.key_password,
                     password=slurm_config.password,
                 )
-                slurm_load = calculates_slurm_load(ssh_config)
+                partition = slurm_config.partition
+                slurm_load = calculates_slurm_load(ssh_config, partition)
                 return LauncherLoadDTO(
                     allocated_cpu_rate=slurm_load[0],
                     cluster_load_rate=slurm_load[1],

diff --git a/antarest/launcher/ssh_client.py b/antarest/launcher/ssh_client.py
@@ -67,12 +67,12 @@ def parse_cpu_load(sinfo_output: str) -> float:
 
 
 def calculates_slurm_load(
-    ssh_config: SSHConfigDTO,
+    ssh_config: SSHConfigDTO, partition: str
 ) -> Tuple[float, float, int]:
     # allocated cpus
     sinfo_cpus_used = execute_command(
         ssh_config,
-        ["sinfo", "--partition", "calin1", "-O", "NodeAIOT", "--noheader"],
+        ["sinfo", "--partition", partition, "-O", "NodeAIOT", "--noheader"],
     )
     allocated_cpus = 100 * parse_cpu_used(sinfo_cpus_used)
     # cluster load
@@ -81,7 +81,7 @@ def calculates_slurm_load(
         [
             "sinfo",
             "--partition",
-            "calin1",
+            partition,
             "-N",
             "-O",
             "CPUsLoad",
@@ -96,7 +96,7 @@ def calculates_slurm_load(
             [
                 "squeue",
                 "--partition",
-                "calin1",
+                partition,
                 "--noheader",
                 "-t",
                 "pending",

diff --git a/resources/application.yaml b/resources/application.yaml
@@ -59,6 +59,7 @@ launcher:
 #    default_n_cpu: 12
 #    default_json_db_name: launcher_db.json
 #    slurm_script_path: /path/to/launchantares_v1.1.3.sh
+#    partition: specific_partition_for_the_resource_allocation
 #    db_primary_key: name
 #    antares_versions_on_remote_server :
 #      - "610"

diff --git a/tests/launcher/test_slurm_launcher.py b/tests/launcher/test_slurm_launcher.py
@@ -35,6 +35,7 @@ def launcher_config(tmp_path: Path) -> Config:
                 local_workspace=tmp_path,
                 default_json_db_name="default_json_db_name",
                 slurm_script_path="slurm_script_path",
+                partition="fake_partition",
                 antares_versions_on_remote_server=["42", "45"],
                 username="username",
                 hostname="hostname",
@@ -111,6 +112,7 @@ def test_init_slurm_launcher_parameters(tmp_path: Path):
                 local_workspace=tmp_path,
                 default_json_db_name="default_json_db_name",
                 slurm_script_path="slurm_script_path",
+                partition="faek_partition",
                 antares_versions_on_remote_server=["42"],
                 username="username",
                 hostname="hostname",
@@ -136,6 +138,7 @@ def test_init_slurm_launcher_parameters(tmp_path: Path):
         main_parameters.slurm_script_path
         == config.launcher.slurm.slurm_script_path
     )
+    assert main_parameters.partition == config.launcher.slurm.partition
     assert (
         main_parameters.antares_versions_on_remote_server
         == config.launcher.slurm.antares_versions_on_remote_server
@@ -516,6 +519,7 @@ def test_kill_job(
         json_dir=Path(tmp_path),
         default_json_db_name="default_json_db_name",
         slurm_script_path="slurm_script_path",
+        partition="fake_partition",
         antares_versions_on_remote_server=["42", "45"],
         default_ssh_dict={
             "username": "username",

diff --git a/tests/launcher/test_ssh_client.py b/tests/launcher/test_ssh_client.py
@@ -28,4 +28,4 @@ def test_parse_cpu_load():
 def test_calculates_slurm_load_whithout_pkey_fails():
     ssh_config = Mock()
     with pytest.raises(SlurmError):
-        calculates_slurm_load(ssh_config)
+        calculates_slurm_load(ssh_config, "fake_partition")