From 071a762b3660402cf792ffecc95b9634acdb2b7c Mon Sep 17 00:00:00 2001 From: trn024 Date: Thu, 29 Aug 2024 15:24:56 -0700 Subject: [PATCH] Change return object of the API to integer instead of UserQuotaResponse object --- src/dsmlp/ext/awsed.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/dsmlp/ext/awsed.py b/src/dsmlp/ext/awsed.py index e1f52ef..b01d608 100644 --- a/src/dsmlp/ext/awsed.py +++ b/src/dsmlp/ext/awsed.py @@ -33,15 +33,16 @@ def list_user_teams(self, username: str) -> ListTeamsResponse: return ListTeamsResponse(teams=teams) - # Fetch user's GPU quota with AWSED Api and assign to UserQuotaResponse object - def get_user_gpu_quota(self, username: str) -> UserQuotaResponse: + # Fetch user's GPU quota with AWSED Api + def get_user_gpu_quota(self, username: str) -> int: try: usrGpuQuota = self.client.get_user_quota(username) if not usrGpuQuota: return None 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) + UserQuotaResponse(quota=quota) + return gpu_quota except Exception as e: self.logger.error(f"Failed to fetch GPU quota for user {username}: {e}") return None