After running log debug, AWSED return "gpu": int(value). Fix and push

src/dsmlp/ext/awsed.py

@@ -40,10 +40,9 @@ def get_user_gpu_quota(self, username: str) -> int:
             self.logger.debug(f"usrGpuQuota: {usrGpuQuota}")  # Log the structure of usrGpuQuota
             if not usrGpuQuota:
                 return None
-            gpu_quota = usrGpuQuota.quota.resources.get("nvidia.com/gpu", "0")  # Access the correct attribute
-            quota = Quota(user=username, resources={"nvidia.com/gpu": gpu_quota})
+            gpu_quota = usrGpuQuota.quota.resources.get("gpu", 0)  # Access the correct attribute
+            quota = Quota(user=username, resources={"gpu": gpu_quota})
             response = UserQuotaResponse(quota=quota)
-            gpu_quota = int(gpu_quota) # convert string to INT according AWSED JSON return
             return gpu_quota
 
         # Debugging

tests/app/test_gpu_validator.py

@@ -92,7 +92,7 @@ def try_validate(self, json, expected: bool, message: str = None):
 
     # Test correct response for get_user_gpu_quota method
     def test_awsed_gpu_quota_correct_response(self):
-        self.awsed_client.assign_user_gpu_quota('user11', {"nvidia.com/gpu": 5})
+        self.awsed_client.assign_user_gpu_quota('user11', {"gpu": 5})
         user_gpu_quota = self.awsed_client.get_user_gpu_quota('user11')
         assert_that(user_gpu_quota, equal_to(5))
 
@@ -123,7 +123,7 @@ def test_gpu_quota_client_priority(self):
 
         self.kube_client.set_existing_gpus('user11', 3)
         # add awsed quota
-        self.awsed_client.assign_user_gpu_quota('user11', {"nvidia.com/gpu": 6})
+        self.awsed_client.assign_user_gpu_quota('user11', {"gpu": 6})
         self.try_validate(
             gen_request(gpu_req=6, username='user11'), expected=False, message="GPU quota exceeded. Wanted 6 but with 3 already in use, the quota of 6 would be exceeded."
         )
@@ -133,7 +133,7 @@ def test_gpu_quota_client_priority2(self):
         self.kube_client.add_namespace('user11', Namespace(
             name='user11', labels={'k8s-sync': 'true'}, gpu_quota=12))
         # add awsed quota
-        self.awsed_client.assign_user_gpu_quota('user11', {"nvidia.com/gpu": 18})
+        self.awsed_client.assign_user_gpu_quota('user11', {"gpu": 18})
 
         # set existing gpu = kube client quota
         self.kube_client.set_existing_gpus('user11', 12)

tests/fakes.py

@@ -34,7 +34,7 @@ def describe_user(self, username: str) -> UserResponse:
     def get_user_gpu_quota(self, username: str) -> int:
         try:
             user_quota_response = self.user_quota[username]
-            return user_quota_response.quota.resources.get("nvidia.com/gpu", 0)
+            return user_quota_response.quota.resources.get("gpu", 0)
         except KeyError:
             return 0