Prevent projects with certain SU types from being given the new PI cr…

…edit
CCI-MOC · May 8, 2024 · 5e40c30 · 5e40c30
1 parent cd231b4
commit 5e40c30
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 10 deletions.
diff --git a/process_report/process_report.py b/process_report/process_report.py
@@ -336,6 +336,7 @@ def validate_pi_names(dataframe):
 def apply_credits_new_pi(dataframe, old_pi_file):
     new_pi_credit_code = "0002"
     new_pi_credit_amount = 1000
+    EXCLUDE_SU_TYPES = ["OpenShift GPUA100SXM4", "OpenStack GPUA100SXM4"]
 
     dataframe[CREDIT_FIELD] = None
     dataframe[CREDIT_CODE_FIELD] = None
@@ -357,16 +358,16 @@ def apply_credits_new_pi(dataframe, old_pi_file):
             print(f"Found new PI {pi}")
             remaining_credit = new_pi_credit_amount
             for i, row in pi_projects.iterrows():
-                project_cost = row[COST_FIELD]
-                applied_credit = min(project_cost, remaining_credit)
-
-                dataframe.at[i, CREDIT_FIELD] = applied_credit
-                dataframe.at[i, CREDIT_CODE_FIELD] = new_pi_credit_code
-                dataframe.at[i, BALANCE_FIELD] = row[COST_FIELD] - applied_credit
-                remaining_credit -= applied_credit
-
-                if remaining_credit == 0:
-                    break
+                if remaining_credit == 0 or row[SU_TYPE_FIELD] in EXCLUDE_SU_TYPES:
+                    dataframe.at[i, BALANCE_FIELD] = row[COST_FIELD]
+                else:
+                    project_cost = row[COST_FIELD]
+                    applied_credit = min(project_cost, remaining_credit)
+
+                    dataframe.at[i, CREDIT_FIELD] = applied_credit
+                    dataframe.at[i, CREDIT_CODE_FIELD] = new_pi_credit_code
+                    dataframe.at[i, BALANCE_FIELD] = row[COST_FIELD] - applied_credit
+                    remaining_credit -= applied_credit
 
     dump_old_pis(old_pi_file, old_pi_dict)
 

diff --git a/process_report/tests/unit_tests.py b/process_report/tests/unit_tests.py
@@ -278,9 +278,31 @@ def setUp(self):
                 "ProjectE",
                 "ProjectF",
             ],
+            "SU Type": ["CPU", "CPU", "CPU", "GPU", "GPU", "GPU"],
             "Cost": [10, 100, 10000, 5000, 800, 1000],
         }
         self.dataframe = pandas.DataFrame(data)
+
+        data_no_gpu = {
+            "Invoice Month": [
+                "2024-03",
+                "2024-03",
+                "2024-03",
+                "2024-03",
+                "2024-03",
+            ],
+            "Manager (PI)": ["PI1", "PI1", "PI1", "PI2", "PI2"],
+            "SU Type": [
+                "GPU",
+                "OpenShift GPUA100SXM4",
+                "OpenStack GPUA100SXM4",
+                "OpenShift GPUA100SXM4",
+                "OpenStack GPUA100SXM4",
+            ],
+            "Cost": [500, 100, 100, 500, 500],
+        }
+        self.dataframe_no_gpu = pandas.DataFrame(data_no_gpu)
+
         old_pi = [
             "PI2,2023-09",
             "PI3,2024-02",
@@ -291,8 +313,42 @@ def setUp(self):
             old_pi_file.write(pi + "\n")
         self.old_pi_file = old_pi_file.name
 
+        old_pi_no_gpu = [
+            "OldPI,2024-03",
+        ]
+        old_pi_no_gpu_file = tempfile.NamedTemporaryFile(
+            delete=False, mode="w", suffix=".csv"
+        )
+        for pi in old_pi_no_gpu:
+            old_pi_no_gpu_file.write(pi + "\n")
+        self.old_pi_no_gpu_file = old_pi_no_gpu_file.name
+        self.no_gpu_df_answer = pandas.DataFrame(
+            {
+                "Invoice Month": [
+                    "2024-03",
+                    "2024-03",
+                    "2024-03",
+                    "2024-03",
+                    "2024-03",
+                ],
+                "Manager (PI)": ["PI1", "PI1", "PI1", "PI2", "PI2"],
+                "SU Type": [
+                    "GPU",
+                    "OpenShift GPUA100SXM4",
+                    "OpenStack GPUA100SXM4",
+                    "OpenShift GPUA100SXM4",
+                    "OpenStack GPUA100SXM4",
+                ],
+                "Cost": [500, 100, 100, 500, 500],
+                "Credit": [500, None, None, None, None],
+                "Credit Code": ["0002", None, None, None, None],
+                "Balance": [0.0, 100.0, 100.0, 500.0, 500.0],
+            }
+        )
+
     def tearDown(self):
         os.remove(self.old_pi_file)
+        os.remove(self.old_pi_no_gpu_file)
 
     def test_apply_credit_0002(self):
         dataframe = process_report.apply_credits_new_pi(
@@ -332,6 +388,13 @@ def test_apply_credit_0002(self):
         with open(self.old_pi_file, "r") as f:
             self.assertEqual(updated_old_pi_answer, f.read())
 
+    def test_no_gpu(self):
+        dataframe = process_report.apply_credits_new_pi(
+            self.dataframe_no_gpu, self.old_pi_no_gpu_file
+        )
+        dataframe = dataframe.astype({"Credit": "float64", "Balance": "float64"})
+        self.assertTrue(self.no_gpu_df_answer.equals(dataframe))
+
     def test_apply_credit_error(self):
         old_pi_dict = {"PI1": "2024-12"}
         invoice_month = "2024-03"