From 7676aca0f685acb19d9e32ff7853c9d42d3e50ea Mon Sep 17 00:00:00 2001
From: Eva C <29628428+fevac@users.noreply.github.com>
Date: Wed, 13 Dec 2023 16:18:41 +0100
Subject: [PATCH] Set priority for tower compute environment (#2765)(patch)

### Changed
- Set compute environment according the case priority
---
 cg/cli/workflow/rnafusion/base.py   |  2 +-
 cg/cli/workflow/taxprofiler/base.py |  2 +-
 cg/meta/workflow/nf_analysis.py     |  9 ++++++---
 cg/meta/workflow/rnafusion.py       |  2 +-
 cg/meta/workflow/taxprofiler.py     |  1 +
 cg/models/cg_config.py              | 11 ++++++-----
 6 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/cg/cli/workflow/rnafusion/base.py b/cg/cli/workflow/rnafusion/base.py
index 5472e5b8cc..342418a4e6 100644
--- a/cg/cli/workflow/rnafusion/base.py
+++ b/cg/cli/workflow/rnafusion/base.py
@@ -119,7 +119,7 @@ def run(
                 case_id=case_id, params_file=params_file
             ),
             "name": case_id,
-            "compute_env": compute_env or analysis_api.compute_env,
+            "compute_env": compute_env or analysis_api.get_compute_env(case_id=case_id),
             "revision": revision or analysis_api.revision,
             "wait": "SUBMITTED",
             "id": nf_tower_id,
diff --git a/cg/cli/workflow/taxprofiler/base.py b/cg/cli/workflow/taxprofiler/base.py
index c9da3117cd..24beb6941f 100644
--- a/cg/cli/workflow/taxprofiler/base.py
+++ b/cg/cli/workflow/taxprofiler/base.py
@@ -116,7 +116,7 @@ def run(
                 case_id=case_id, params_file=params_file
             ),
             "name": case_id,
-            "compute_env": compute_env or analysis_api.compute_env,
+            "compute_env": compute_env or analysis_api.get_compute_env(case_id=case_id),
             "revision": revision or analysis_api.revision,
             "wait": NfTowerStatus.SUBMITTED,
             "id": nf_tower_id,
diff --git a/cg/meta/workflow/nf_analysis.py b/cg/meta/workflow/nf_analysis.py
index 0e8bba56f2..c0a1640b43 100644
--- a/cg/meta/workflow/nf_analysis.py
+++ b/cg/meta/workflow/nf_analysis.py
@@ -4,8 +4,6 @@
 from pathlib import Path
 from typing import Any
 
-from cg.store.models import Sample
-
 from cg.constants import Pipeline
 from cg.constants.constants import FileExtensions, FileFormat, WorkflowManager
 from cg.constants.nextflow import NFX_WORK_DIR
@@ -17,6 +15,7 @@
 from cg.models.cg_config import CGConfig
 from cg.models.nf_analysis import FileDeliverable, PipelineDeliverables
 from cg.models.rnafusion.rnafusion import CommandArgs
+from cg.store.models import Sample
 from cg.utils import Process
 
 LOG = logging.getLogger(__name__)
@@ -38,7 +37,7 @@ def __init__(self, config: CGConfig, pipeline: Pipeline):
         self.tower_pipeline: str | None = None
         self.account: str | None = None
         self.email: str | None = None
-        self.compute_env: str | None = None
+        self.compute_env_base: str | None = None
         self.revision: str | None = None
         self.nextflow_binary_path: str | None = None
 
@@ -80,6 +79,10 @@ def get_sample_sheet_path(self, case_id: str) -> Path:
             FileExtensions.CSV
         )
 
+    def get_compute_env(self, case_id: str) -> str:
+        """Get the compute environment for the head job based on the case priority."""
+        return f"{self.compute_env_base}-{self.get_slurm_qos_for_case(case_id=case_id)}"
+
     @staticmethod
     def get_nextflow_config_path(nextflow_config: str | None = None) -> Path | None:
         """Path to Nextflow config file."""
diff --git a/cg/meta/workflow/rnafusion.py b/cg/meta/workflow/rnafusion.py
index efcccd8ca3..9611e40a04 100644
--- a/cg/meta/workflow/rnafusion.py
+++ b/cg/meta/workflow/rnafusion.py
@@ -50,7 +50,7 @@ def __init__(
         self.tower_pipeline: str = config.rnafusion.tower_pipeline
         self.account: str = config.rnafusion.slurm.account
         self.email: str = config.rnafusion.slurm.mail_user
-        self.compute_env: str = config.rnafusion.compute_env
+        self.compute_env_base: str = config.rnafusion.compute_env
         self.revision: str = config.rnafusion.revision
         self.nextflow_binary_path: str = config.rnafusion.binary_path
 
diff --git a/cg/meta/workflow/taxprofiler.py b/cg/meta/workflow/taxprofiler.py
index 79d2c53d15..3af5f89ee1 100644
--- a/cg/meta/workflow/taxprofiler.py
+++ b/cg/meta/workflow/taxprofiler.py
@@ -40,6 +40,7 @@ def __init__(
         self.account: str = config.taxprofiler.slurm.account
         self.email: str = config.taxprofiler.slurm.mail_user
         self.nextflow_binary_path: str = config.taxprofiler.binary_path
+        self.compute_env_base: str = config.taxprofiler.compute_env
 
     def get_sample_sheet_content_per_sample(
         self, sample: Sample, instrument_platform: SequencingPlatform.ILLUMINA, fasta: str = ""
diff --git a/cg/models/cg_config.py b/cg/models/cg_config.py
index abeccfd70d..8c3c28be0b 100644
--- a/cg/models/cg_config.py
+++ b/cg/models/cg_config.py
@@ -179,15 +179,16 @@ class RnafusionConfig(CommonAppConfig):
 
 
 class TaxprofilerConfig(CommonAppConfig):
-    root: str
     binary_path: str
+    conda_binary: str | None = None
     conda_env: str
-    profile: str
+    compute_env: str
+    databases: str
+    hostremoval_reference: str
     pipeline_path: str
+    profile: str
     revision: str
-    conda_binary: str | None = None
-    hostremoval_reference: str
-    databases: str
+    root: str
     slurm: SlurmConfig
     tower_binary_path: str
     tower_pipeline: str