diff --git a/src/dsmlp/ext/awsed.py b/src/dsmlp/ext/awsed.py index b01d608..0079e49 100644 --- a/src/dsmlp/ext/awsed.py +++ b/src/dsmlp/ext/awsed.py @@ -39,9 +39,10 @@ def get_user_gpu_quota(self, username: str) -> int: 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 + gpu_quota = usrGpuQuota['resources'].get("nvidia.com/gpu", "0") # Access the correct attribute + gpu_quota = int(gpu_quota) # convert string to INT according AWSED JSON return quota = Quota(user=username, resources={"nvidia.com/gpu": gpu_quota}) - UserQuotaResponse(quota=quota) + response = UserQuotaResponse(quota=quota) return gpu_quota except Exception as e: self.logger.error(f"Failed to fetch GPU quota for user {username}: {e}")