Add resourceattrtypes to defauls command. Fix parsing of gpu count. Prevent bulk update error

Author: antonio-rodriguez-tam

coldfront/core/resource/management/commands/add_resource_defaults.py

@@ -43,7 +43,7 @@ def handle(self, *args, **options):
             ('xdmod_resource', 'Text'),
             # ('eula', 'Text'),
             # ('OnDemand','Yes/No'),
-            # ('ServiceEnd', 'Date'),
+            ('ServiceEnd', 'Date'),
             # ('ServiceStart', 'Date'),
             ('slurm_cluster', 'Text'),
             ('slurm_specs', 'Attribute Expanded Text'),

coldfront/plugins/slurm/management/commands/slurm_manage_resources.py

@@ -31,7 +31,6 @@ def get_output_from_file(self, file_path):
         except IOError as e:
             print(f"An error occurred: {e}")
 
-
     def add_arguments(self, parser):
         parser.add_argument("-e", "--environment", help="Environment, use dev to simulate output")
         parser.add_argument('--profile', action='store_true', default=False)
@@ -48,7 +47,7 @@ def _handle(self, *args, **options):
         def calculate_gpu_count(gres_value):
             if 'null' in gres_value:
                 return 0
-            gpu_list = gres_value.split(',')
+            gpu_list = gres_value.split('),')
             return reduce(lambda x, y: x + y,[int(gpu_info.split(':')[2].replace('(S','')) for gpu_info in gpu_list])
 
         def calculate_cpu_count(row):
@@ -77,13 +76,11 @@ def calculate_owner_value(project_list, row):
         project_list = Project.objects.all()
         compute_node, compute_node_created = ResourceType.objects.get_or_create(name='Compute Node', description='Compute Node')
         partition_resource_type, partition_created = ResourceType.objects.get_or_create(name='Cluster Partition', description='Cluster Partition')
-        int_attribute_type = AttributeType.objects.get(name='Int')
-        text_attribute_type = AttributeType.objects.get(name='Text')
-        gpu_count_attribute_type, gpu_count_created = ResourceAttributeType.objects.get_or_create(name='GPU Count', defaults={'attribute_type': int_attribute_type})
-        core_count_attribute_type, core_count_created = ResourceAttributeType.objects.get_or_create(name='Core Count', defaults={'attribute_type': int_attribute_type})
-        features_attribute_type, features_created = ResourceAttributeType.objects.get_or_create(name='Features', defaults={'attribute_type': text_attribute_type})
-        owner_attribute_type, owner_created = ResourceAttributeType.objects.get_or_create(name='Owner', defaults={'attribute_type': text_attribute_type})
-        service_end_attribute_type, service_end_created = ResourceAttributeType.objects.get_or_create(name='ServiceEnd', defaults={'attribute_type': text_attribute_type})
+        gpu_count_attribute_type = ResourceAttributeType.objects.get(name='GPU Count')
+        core_count_attribute_type = ResourceAttributeType.objects.get(name='Core Count')
+        features_attribute_type = ResourceAttributeType.objects.get(name='Features')
+        owner_attribute_type = ResourceAttributeType.objects.get(name='Owner')
+        service_end_attribute_type = ResourceAttributeType.objects.get(name='ServiceEnd')
         processed_resources = set()
         bulk_process_resource_attribute = []
         bulk_update_resource = []
@@ -95,15 +92,24 @@ def calculate_owner_value(project_list, row):
             bulk_process_resource_attribute.append(ResourceAttribute(resource_attribute_type=features_attribute_type, resource=new_resource, value=row.get('avail_features', '(null)')))
             bulk_process_resource_attribute.append(ResourceAttribute(resource_attribute_type=owner_attribute_type, resource=new_resource, value=calculate_owner_value(project_list, row)))
             if new_resource.is_available is False:
-                bulk_update_resource.append(Resource(name=row['nodelist'], is_available=True))
-                bulk_process_resource_attribute.append(ResourceAttribute(resource=new_resource, value=' ', resource_attribute_type=service_end_attribute_type))
+                bulk_update_resource.append(Resource(name=row['nodelist'], is_available=True, resource_type=compute_node))
+                bulk_process_resource_attribute.append(ResourceAttribute(resource=new_resource, value=None, resource_attribute_type=service_end_attribute_type))
             processed_resources.add(new_resource.name)
-        ResourceAttribute.objects.bulk_create(bulk_process_resource_attribute, update_conflicts=True, unique_fields=[], update_fields=['value'])
-        Resource.objects.bulk_create(bulk_update_resource, update_conflicts=True, unique_fields=[],  update_fields=['is_available'])
+        try:
+            ResourceAttribute.objects.bulk_create(bulk_process_resource_attribute, update_conflicts=True, unique_fields=[], update_fields=['value'])
+            Resource.objects.bulk_create(bulk_update_resource, update_conflicts=True, unique_fields=[],  update_fields=['is_available'])
+        except Exception as e:
+            logger.error(f'Error processing resources info: {str(e)}')
+            raise
         bulk_process_resource_attribute = []
         bulk_update_resource = []
         for resource_to_delete in Resource.objects.exclude(name__in=list(processed_resources)).filter(is_available=True, resource_type=compute_node):
-            bulk_update_resource.append(Resource(name=resource_to_delete.name, is_available=False))
+            resource_to_delete.is_available = False
+            bulk_update_resource.append(resource_to_delete)
             bulk_process_resource_attribute.append(ResourceAttribute(resource=resource_to_delete, value=str(datetime.now()), resource_attribute_type=service_end_attribute_type))
-        ResourceAttribute.objects.bulk_create(bulk_process_resource_attribute, update_conflicts=True, unique_fields=[], update_fields=['value'])
-        Resource.objects.bulk_create(bulk_update_resource, update_conflicts=True, unique_fields=[], update_fields=['is_available'])
+        try:
+            ResourceAttribute.objects.bulk_create(bulk_process_resource_attribute, update_conflicts=True, unique_fields=[], update_fields=['value'])
+            Resource.objects.bulk_create(bulk_update_resource, update_conflicts=True, unique_fields=[], update_fields=['is_available'])
+        except Exception as e:
+            logger.error(f'Error cleaning up resources: {str(e)}')
+            raise