Skip to content

Commit

Permalink
Test rework
Browse files Browse the repository at this point in the history
  • Loading branch information
@shouhanzen
shouhanzen committed Feb 29, 2024
1 parent 14e5473 commit 9bb95c2
Show file tree
Hide file tree
Showing 6 changed files with 282 additions and 841 deletions.
1 change: 1 addition & 0 deletions requirements.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,4 +7,5 @@ PyHamcrest
requests_mock
dataclasses-json
python-dotenv
pytest
git+https://github.com/ucsd-ets/[email protected]
235 changes: 30 additions & 205 deletions tests/app/test_gpu_validator.py
View file
Original file line number Diff line number Diff line change
@@ -1,246 +1,71 @@
import inspect
from operator import contains
from dsmlp.app.types import ValidationFailure
from dsmlp.app.validator import Validator
from dsmlp.plugin.awsed import ListTeamsResponse, TeamJson, UserResponse
from dsmlp.plugin.kube import Namespace
from hamcrest import assert_that, contains_inanyorder, equal_to, has_item
from tests.fakes import FakeAwsedClient, FakeLogger, FakeKubeClient
from dsmlp.ext.kube import DefaultKubeClient
from dsmlp.app.gpu_validator import GPUValidator
from tests.app.utils import gen_request, try_val_with_component

class TestValidator:

class TestGPUValidator:
def setup_method(self) -> None:
self.logger = FakeLogger()
self.awsed_client = FakeAwsedClient()
self.kube_client = FakeKubeClient()

self.awsed_client.add_user('user10', UserResponse(uid=10, enrollments=[]))
self.awsed_client.add_user(
'user10', UserResponse(uid=10, enrollments=[]))
self.awsed_client.add_teams('user10', ListTeamsResponse(
teams=[TeamJson(gid=1000)]
))

self.kube_client.add_namespace('user10', Namespace(name='user10', labels={'k8s-sync': 'true'}, gpu_quota=10))

self.kube_client.add_namespace('user10', Namespace(
name='user10', labels={'k8s-sync': 'true'}, gpu_quota=10))
self.kube_client.set_existing_gpus('user10', 0)

def test_no_gpus_requested(self):
response = self.when_validate(
{
"request": {
"uid": "705ab4f5-6393-11e8-b7cc-42010a800002",
"userInfo": {
"username": "user10"
},
"namespace": "user10",
"object": {
"metadata": {
"labels": {}
},
"kind": "Pod",
"spec": {
"containers": [{}]
}
}
}}
self.try_validate(
gen_request(), expected=True, message="Allowed"
)

assert_that(response, equal_to({
"apiVersion": "admission.k8s.io/v1",
"kind": "AdmissionReview",
"response": {
"uid": "705ab4f5-6393-11e8-b7cc-42010a800002",
"allowed": True, "status": {
"message": "Allowed"
}}}))

def test_quota_not_reached(self):
response = self.when_validate(
{
"request": {
"uid": "705ab4f5-6393-11e8-b7cc-42010a800002",
"userInfo": {
"username": "user10"
},
"namespace": "user10",
"object": {
"metadata": {
"labels": {}
},
"kind": "Pod",
"spec": {
"containers": [{
"resources": {
"requests": {
"nvidia.com/gpu": 10
}
}
}]
}
}
}}
)

assert_that(response, equal_to({
"apiVersion": "admission.k8s.io/v1",
"kind": "AdmissionReview",
"response": {
"uid": "705ab4f5-6393-11e8-b7cc-42010a800002",
"allowed": True, "status": {
"message": "Allowed"
}}}))
self.try_validate(
gen_request(gpu_req=10), expected=True, message="Allowed"
)

def test_quota_exceeded(self):
response = self.when_validate(
{
"request": {
"uid": "705ab4f5-6393-11e8-b7cc-42010a800002",
"userInfo": {
"username": "user10"
},
"namespace": "user10",
"object": {
"metadata": {
"labels": {}
},
"kind": "Pod",
"spec": {
"containers": [{
"resources": {
"requests": {
"nvidia.com/gpu": 11
}
}
}]
}
}
}}

self.try_validate(
gen_request(gpu_req=11), expected=False, message="GPU quota exceeded. Wanted 11 but with 0 already in use, the quota of 10 would be exceeded."
)

assert_that(response, equal_to({
"apiVersion": "admission.k8s.io/v1",
"kind": "AdmissionReview",
"response": {
"uid": "705ab4f5-6393-11e8-b7cc-42010a800002",
"allowed": False, "status": {
"message": "GPU quota exceeded. Wanted 11 but with 0 already in use, the quota of 10 would be exceeded."
}}}))

def test_sum_exceeded(self):
self.kube_client.set_existing_gpus('user10', 5)

response = self.when_validate(
{
"request": {
"uid": "705ab4f5-6393-11e8-b7cc-42010a800002",
"userInfo": {
"username": "user10"
},
"namespace": "user10",
"object": {
"metadata": {
"labels": {}
},
"kind": "Pod",
"spec": {
"containers": [{
"resources": {
"requests": {
"nvidia.com/gpu": 6
}
}
}]
}
}
}}

self.try_validate(
gen_request(gpu_req=6), expected=False, message="GPU quota exceeded. Wanted 6 but with 5 already in use, the quota of 10 would be exceeded."
)

assert_that(response, equal_to({
"apiVersion": "admission.k8s.io/v1",
"kind": "AdmissionReview",
"response": {
"uid": "705ab4f5-6393-11e8-b7cc-42010a800002",
"allowed": False, "status": {
"message": "GPU quota exceeded. Wanted 6 but with 5 already in use, the quota of 10 would be exceeded."
}}}))

def test_low_priority(self):
self.kube_client.set_existing_gpus('user10', 5)

response = self.when_validate(
{
"request": {
"uid": "705ab4f5-6393-11e8-b7cc-42010a800002",
"userInfo": {
"username": "user10"
},
"namespace": "user10",
"object": {
"metadata": {
"labels": {}
},
"kind": "Pod",
"spec": {
"containers": [{
"resources": {
"requests": {
"nvidia.com/gpu": 6
}
}
}],
"priorityClassName": "low"
}
}
}}

self.try_validate(
gen_request(gpu_req=6, low_priority=True), expected=True
)

assert_that(response, equal_to({
"apiVersion": "admission.k8s.io/v1",
"kind": "AdmissionReview",
"response": {
"uid": "705ab4f5-6393-11e8-b7cc-42010a800002",
"allowed": True, "status": {
"message": "Allowed"
}}}))

# Should respond to limit as well as request
def test_limit_exceeded(self):
self.kube_client.set_existing_gpus('user10', 5)

response = self.when_validate(
{
"request": {
"uid": "705ab4f5-6393-11e8-b7cc-42010a800002",
"userInfo": {
"username": "user10"
},
"namespace": "user10",
"object": {
"metadata": {
"labels": {}
},
"kind": "Pod",
"spec": {
"containers": [{
"resources": {
"limits": {
"nvidia.com/gpu": 6
}
}
}]
}
}
}}
)

assert_that(response, equal_to({
"apiVersion": "admission.k8s.io/v1",
"kind": "AdmissionReview",
"response": {
"uid": "705ab4f5-6393-11e8-b7cc-42010a800002",
"allowed": False, "status": {
"message": "GPU quota exceeded. Wanted 6 but with 5 already in use, the quota of 10 would be exceeded."
}}}))

def when_validate(self, json):
validator = Validator(self.awsed_client, self.kube_client, self.logger)
response = validator.validate_request(json)
self.try_validate(
gen_request(gpu_lim=6), expected=False, message="GPU quota exceeded. Wanted 6 but with 5 already in use, the quota of 10 would be exceeded."
)

return response
def try_validate(self, json, expected: bool, message: str = None):
try_val_with_component(GPUValidator(
self.kube_client, self.logger), json, expected, message)
Loading

0 comments on commit 9bb95c2

Please sign in to comment.