.get method is wrong. Fixed according the JSON structure when fetchin…

…g with AWSED
ucsd-ets · Aug 29, 2024 · 61b2aaa · 61b2aaa
1 parent bdd3773
commit 61b2aaa
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/dsmlp/ext/awsed.py b/src/dsmlp/ext/awsed.py
@@ -39,8 +39,8 @@ def get_user_gpu_quota(self, username: str) -> UserQuotaResponse:
             usrGpuQuota = self.client.get_user_quota(username)
             if not usrGpuQuota:
                 return None
-            gpu_quota = usrGpuQuota.get("nvidia.com/gpu", 0)
-            quota = Quota(user=username, resources=gpu_quota)
+            gpu_quota = usrGpuQuota['resources'].get("nvidia.com/gpu", 0)  # Access the correct attribute
+            quota = Quota(user=username, resources={"nvidia.com/gpu": gpu_quota})
             return UserQuotaResponse(quota=quota)
         except Exception as e:
             self.logger.error(f"Failed to fetch GPU quota for user {username}: {e}")