From 18cc5e65ba10b58643312b2ba1fcf7fa463ab37e Mon Sep 17 00:00:00 2001 From: Clif Houck Date: Fri, 23 Aug 2024 12:49:59 -0500 Subject: [PATCH] Add support for binoculars to python client Signed-off-by: Clif Houck --- .../python/armada_client/binoculars_client.py | 45 +++++++++++++++++++ client/python/tests/unit/server_mock.py | 20 ++++++++- .../tests/unit/test_binoculars_client.py | 44 ++++++++++++++++++ scripts/build-python-client.sh | 4 +- 4 files changed, 109 insertions(+), 4 deletions(-) create mode 100644 client/python/armada_client/binoculars_client.py create mode 100644 client/python/tests/unit/test_binoculars_client.py diff --git a/client/python/armada_client/binoculars_client.py b/client/python/armada_client/binoculars_client.py new file mode 100644 index 00000000000..f2b821b1f87 --- /dev/null +++ b/client/python/armada_client/binoculars_client.py @@ -0,0 +1,45 @@ +import datetime +from typing import Optional + + +from armada_client.armada import ( + binoculars_pb2, + binoculars_pb2_grpc, +) + +from armada_client.k8s.io.api.core.v1 import generated_pb2 as core_v1 + + +class BinocularsClient: + """ + Client for accessing Armada's Binoculars service over gRPC. + + :param channel: gRPC channel used for authentication. See + https://grpc.github.io/grpc/python/grpc.html + for more information. + :return: an Binoculars client instance + """ + + def __init__(self, channel): + self.binoculars_stub = binoculars_pb2_grpc.BinocularsStub(channel) + + def logs( + self, + job_id: str, + pod_namespace: str, + since_time: datetime.datetime, + pod_number: Optional[int] = 0, + log_options: Optional[core_v1.PodLogOptions] = None, + ): + log_request = binoculars_pb2.LogRequest( + job_id=job_id, + pod_number=pod_number, + pod_namespace=pod_namespace, + since_time=since_time.isoformat(), + log_options=log_options, + ) + return self.binoculars_stub.Logs(log_request) + + def cordon(self, node_name: str): + cordon_request = binoculars_pb2.CordonRequest(node_name=node_name) + return self.binoculars_stub.Cordon(cordon_request) diff --git a/client/python/tests/unit/server_mock.py b/client/python/tests/unit/server_mock.py index 8d19101203b..1bb94bcf053 100644 --- a/client/python/tests/unit/server_mock.py +++ b/client/python/tests/unit/server_mock.py @@ -1,13 +1,15 @@ from google.protobuf import empty_pb2 from armada_client.armada import ( - submit_pb2_grpc, - submit_pb2, + binoculars_pb2, + binoculars_pb2_grpc, event_pb2, event_pb2_grpc, health_pb2, job_pb2_grpc, job_pb2, + submit_pb2, + submit_pb2_grpc, ) from armada_client.armada.job_pb2 import JobRunState from armada_client.armada.submit_pb2 import JobState @@ -149,3 +151,17 @@ def GetJobRunDetails(self, request, context): for run in request.run_ids } ) + + +class BinocularsService(binoculars_pb2_grpc.BinocularsServicer): + def Logs(self, request, context): + return binoculars_pb2.LogResponse( + log=[ + binoculars_pb2.LogLine(timestamp="now", line="some log contents!"), + binoculars_pb2.LogLine(timestamp="now", line="some more log contents!"), + binoculars_pb2.LogLine(timestamp="now", line="even more log contents!"), + ], + ) + + def Cordon(self, request, context): + return empty_pb2.Empty() diff --git a/client/python/tests/unit/test_binoculars_client.py b/client/python/tests/unit/test_binoculars_client.py new file mode 100644 index 00000000000..cd9a3a5474a --- /dev/null +++ b/client/python/tests/unit/test_binoculars_client.py @@ -0,0 +1,44 @@ +from concurrent import futures +import datetime + +import grpc +import pytest + +from google.protobuf import empty_pb2 + +from server_mock import BinocularsService + +from armada_client.armada import binoculars_pb2_grpc +from armada_client.binoculars_client import BinocularsClient + + +@pytest.fixture(scope="session", autouse=True) +def binoculars_server_mock(): + server = grpc.server(futures.ThreadPoolExecutor(max_workers=10)) + binoculars_pb2_grpc.add_BinocularsServicer_to_server(BinocularsService(), server) + server.add_insecure_port("[::]:4000") + server.start() + + yield + server.stop(False) + + +channel = grpc.insecure_channel(target="127.0.0.1:4000") +tester = BinocularsClient( + grpc.insecure_channel( + target="127.0.0.1:4000", + options={ + "grpc.keepalive_time_ms": 30000, + }.items(), + ) +) + + +def test_logs(): + resp = tester.logs("fake-job-id", "fake-namespace", datetime.datetime.now()) + assert len(resp.log) == 3 + + +def test_cordon(): + result = tester.cordon("fake-node-name") + assert result == empty_pb2.Empty() diff --git a/scripts/build-python-client.sh b/scripts/build-python-client.sh index 9e2bece9a24..5fd23818146 100755 --- a/scripts/build-python-client.sh +++ b/scripts/build-python-client.sh @@ -3,7 +3,7 @@ # make the python package armada.client, not pkg.api mkdir -p proto/armada -cp pkg/api/event.proto pkg/api/submit.proto pkg/api/health.proto pkg/api/job.proto proto/armada +cp pkg/api/event.proto pkg/api/submit.proto pkg/api/health.proto pkg/api/job.proto pkg/api/binoculars/binoculars.proto proto/armada sed -i 's/\([^\/]\)pkg\/api/\1armada/g' proto/armada/*.proto # generate python stubs @@ -11,7 +11,7 @@ cd proto python3 -m grpc_tools.protoc -I. --plugin=protoc-gen-mypy=$(which protoc-gen-mypy) --python_out=../client/python/armada_client --grpc_python_out=../client/python/armada_client --mypy_out=../client/python/armada_client \ google/api/annotations.proto \ google/api/http.proto \ - armada/event.proto armada/submit.proto armada/health.proto armada/job.proto \ + armada/event.proto armada/submit.proto armada/health.proto armada/job.proto armada/binoculars.proto \ github.com/gogo/protobuf/gogoproto/gogo.proto \ k8s.io/api/core/v1/generated.proto \ k8s.io/apimachinery/pkg/api/resource/generated.proto \