Merge pull request #20 from NYU-ITS/auto_slurm_account_name_assignment

Auto slurm account name assignment

Author: s-sajid-ali

coldfront/core/allocation/test_views.py

@@ -1,5 +1,7 @@
 import logging
 
+from coldfront.core.allocation.views import GENERAL_RESOURCE_NAME
+from coldfront.core.resource.models import Resource, ResourceType
 from coldfront.core.project.models import Project
 from coldfront.core.school.models import School
 from coldfront.core.user.models import UserProfile, ApproverProfile
@@ -24,7 +26,13 @@
 )
 from coldfront.core.allocation.models import (
     AllocationChangeRequest,
-    AllocationChangeStatusChoice, Allocation, AllocationStatusChoice, AllocationAttributeChangeRequest,
+    AllocationChangeStatusChoice,
+    AllocationStatusChoice,
+    AllocationAttributeChangeRequest,
+    Allocation,
+    AllocationAttribute,
+    AllocationAttributeType,
+    AttributeType,
 )
 
 logging.disable(logging.CRITICAL)
@@ -415,6 +423,59 @@ def setUp(self):
             'resource': f'{self.allocation.resources.first().pk}',
         }
 
+        # Create resources
+        self.resource_standard = Resource.objects.create(name="Tandon",
+                                                         resource_type=ResourceType.objects.create(name="Generic"))
+        self.resource_hpc = Resource.objects.create(name=GENERAL_RESOURCE_NAME,
+                                                    resource_type=ResourceType.objects.create(name="Cluster"))
+
+        # Create AllocationAttributeType for slurm_account_name
+        self.slurm_account_attr_type, _ = AllocationAttributeType.objects.get_or_create(
+            name='slurm_account_name', attribute_type=AttributeType.objects.create(name='Text'),
+            has_usage=False, is_private=False)
+
+    def test_allocationcreateview_post_standard_resource(self):
+        """Test POST to the AllocationCreateView with a standard resource and check slurm_account_name"""
+        self.post_data['resource'] = str(self.resource_standard.pk)
+        initial_count = self.project.allocation_set.count()
+
+        response = self.client.post(self.url, data=self.post_data, follow=True)
+
+        self.assertEqual(response.status_code, 200)
+        self.assertContains(response, "Allocation requested.")
+        self.assertEqual(self.project.allocation_set.count(), initial_count + 1)
+
+        # Verify AllocationAttribute for slurm_account_name
+        allocation = self.project.allocation_set.latest('id')
+        expected_slurm_name = f"pr_{self.project.pk}_{self.resource_standard.name}"
+
+        slurm_attr = AllocationAttribute.objects.get(
+            allocation_attribute_type=self.slurm_account_attr_type,
+            allocation=allocation
+        )
+        self.assertEqual(slurm_attr.value, expected_slurm_name)
+
+    def test_allocationcreateview_post_hpc_resource(self):
+        """Test POST to the AllocationCreateView with University HPC resource and check slurm_account_name"""
+        self.post_data['resource'] = str(self.resource_hpc.pk)
+        initial_count = self.project.allocation_set.count()
+
+        response = self.client.post(self.url, data=self.post_data, follow=True)
+
+        self.assertEqual(response.status_code, 200)
+        self.assertContains(response, "Allocation requested.")
+        self.assertEqual(self.project.allocation_set.count(), initial_count + 1)
+
+        # Verify AllocationAttribute for slurm_account_name
+        allocation = self.project.allocation_set.latest('id')
+        expected_slurm_name = f"pr_{self.project.pk}_general"
+
+        slurm_attr = AllocationAttribute.objects.get(
+            allocation_attribute_type=self.slurm_account_attr_type,
+            allocation=allocation
+        )
+        self.assertEqual(slurm_attr.value, expected_slurm_name)
+
     def test_allocationcreateview_access(self):
         """Test access to the AllocationCreateView"""
         self.allocation_access_tstbase(self.url)

coldfront/core/allocation/views.py

@@ -80,6 +80,7 @@
 ALLOCATION_ACCOUNT_MAPPING = import_from_settings(
     'ALLOCATION_ACCOUNT_MAPPING', {})
 
+GENERAL_RESOURCE_NAME = "University HPC"
 
 logger = logging.getLogger(__name__)
 
@@ -529,6 +530,18 @@ def form_valid(self, form):
 
         allocation_obj.resources.add(resource_obj)
 
+        # Automatically create slurm_account_name AllocationAttribute
+        allocation_attribute_type_obj = AllocationAttributeType.objects.get(name='slurm_account_name')
+        if resource_obj.name == GENERAL_RESOURCE_NAME: # "University HPC"
+            slurm_account_name = f"pr_{allocation_obj.project.pk}_general"
+        else:
+            slurm_account_name = f"pr_{allocation_obj.project.pk}_{resource_obj.name}"
+
+        AllocationAttribute.objects.get_or_create(
+            allocation_attribute_type=allocation_attribute_type_obj,
+            allocation=allocation_obj,
+            value=slurm_account_name)
+
         if ALLOCATION_ACCOUNT_ENABLED and allocation_account and resource_obj.name in ALLOCATION_ACCOUNT_MAPPING:
 
             allocation_attribute_type_obj = AllocationAttributeType.objects.get(

coldfront/core/school/management/commands/data/school_data.csv

@@ -18,4 +18,5 @@
 18	Tisch School of the Arts
 19	NYU Abu Dhabi
 20	NYU Shanghai
-21	NYU IT
+21	NYU IT
+22	Center for Data Science