diff --git a/src/dsmlp/ext/awsed.py b/src/dsmlp/ext/awsed.py index f835948..33e9819 100644 --- a/src/dsmlp/ext/awsed.py +++ b/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 diff --git a/tests/app/test_gpu_validator.py b/tests/app/test_gpu_validator.py index f5754e8..df09044 100644 --- a/tests/app/test_gpu_validator.py +++ b/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) diff --git a/tests/fakes.py b/tests/fakes.py index d977550..b2f2114 100644 --- a/tests/fakes.py +++ b/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