def wait_for(client: ArmadaClient, queue, job_set_id=None):
"""
Waits for a queue and optionally the job_set_id to be active.
Ensures that following steps will not fail.
"""
timeout = 20
while True:
try:
# queue active test
client.get_queue(name=queue)
if job_set_id:
events = client.get_job_events_stream(queue=queue,
job_set_id=job_set_id)
for _ in events:
break
return True
except grpc.RpcError as e:
code = e.code()
if code != grpc.StatusCode.NOT_FOUND:
raise e
timeout -= 1
time.sleep(1)
if timeout <= 0:
raise Exception("Timeout")
def no_auth_client() -> ArmadaClient:
server_name = os.environ.get("ARMADA_SERVER", "localhost")
server_port = os.environ.get("ARMADA_PORT", "50051")
server_ssl = os.environ.get("ARMADA_SSL", "false")
if server_ssl.lower() == "true":
channel_credentials = grpc.ssl_channel_credentials()
return ArmadaClient(channel=grpc.secure_channel(
f"{server_name}:{server_port}", channel_credentials))
else:
return ArmadaClient(
channel=grpc.insecure_channel(f"{server_name}:{server_port}"))
def test_submit_job_and_cancel_by_id(client: ArmadaClient, queue_name):
job_set_name = f"set-{uuid.uuid1()}"
jobs = client.submit_jobs(
queue=queue_name,
job_set_id=job_set_name,
job_request_items=submit_sleep_job(client),
)
wait_for(client, queue=queue_name, job_set_id=job_set_name)
cancelled_message = client.cancel_jobs(
job_id=jobs.job_response_items[0].job_id)
assert cancelled_message.cancelled_ids[0] == jobs.job_response_items[
0].job_id
def test_submit_job_and_cancel_by_queue_job_set(client: ArmadaClient,
queue_name):
job_set_name = f"set-{uuid.uuid1()}"
client.submit_jobs(
queue=queue_name,
job_set_id=job_set_name,
job_request_items=submit_sleep_job(client),
)
wait_for(client, queue=queue_name, job_set_id=job_set_name)
cancelled_message = client.cancel_jobs(queue=queue_name,
job_set_id=job_set_name)
expected = f"all jobs in job set {job_set_name}"
assert expected == cancelled_message.cancelled_ids[0]
def __init__(self, host, port, username, password, disable_ssl=True):
# TODO: generalize this so tests can be run with a variety of auth schemas
if disable_ssl:
channel_credentials = grpc.local_channel_credentials()
else:
# TODO pass root certs, private key, cert chain if this is needed
channel_credentials = grpc.ssl_channel_credentials()
channel = grpc.secure_channel(
f"{host}:{port}",
grpc.composite_channel_credentials(
channel_credentials,
grpc.metadata_call_credentials(
GrpcBasicAuth(username, password)),
),
)
self.client = ArmadaClient(channel)
def test_get_job_events_stream(client: ArmadaClient, queue_name):
job_set_name = f"set-{uuid.uuid1()}"
client.submit_jobs(
queue=queue_name,
job_set_id=job_set_name,
job_request_items=submit_sleep_job(client),
)
wait_for(client, queue=queue_name, job_set_id=job_set_name)
event_stream = client.get_job_events_stream(queue=queue_name,
job_set_id=job_set_name)
# Avoiding fickle tests, we will just pass this test as long as we find an event.
found_event = False
for _ in event_stream:
found_event = True
break
assert found_event
def server_mock():
server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
submit_pb2_grpc.add_SubmitServicer_to_server(SubmitService(), server)
event_pb2_grpc.add_EventServicer_to_server(EventService(), server)
server.add_insecure_port("[::]:50051")
server.start()
yield
server.stop(False)
channel = grpc.insecure_channel(target="127.0.0.1:50051")
tester = ArmadaClient(
grpc.insecure_channel(
target="127.0.0.1:50051",
options={
"grpc.keepalive_time_ms": 30000,
}.items(),
))
def test_submit_job():
pod = core_v1.PodSpec(containers=[
core_v1.Container(
name="Container1",
image="index.docker.io/library/ubuntu:latest",
args=["sleep", "10s"],
securityContext=core_v1.SecurityContext(runAsUser=1000),
resources=core_v1.ResourceRequirements(
requests={
"cpu": api_resource.Quantity(string="120m"),
def test_get_queue(client: ArmadaClient, queue_name):
queue = client.get_queue(name=queue_name)
assert queue.name == queue_name
def create_queue(client: ArmadaClient, queue_name):
client.create_queue(name=queue_name, priority_factor=1)
wait_for(client, queue=queue_name)
def test_get_queue_info(client: ArmadaClient, queue_name):
queue = client.get_queue_info(name=queue_name)
assert queue.name == queue_name
assert not queue.active_job_sets
class BasicAuthTest:
def __init__(self, host, port, username, password, disable_ssl=True):
# TODO: generalize this so tests can be run with a variety of auth schemas
if disable_ssl:
channel_credentials = grpc.local_channel_credentials()
else:
# TODO pass root certs, private key, cert chain if this is needed
channel_credentials = grpc.ssl_channel_credentials()
channel = grpc.secure_channel(
f"{host}:{port}",
grpc.composite_channel_credentials(
channel_credentials,
grpc.metadata_call_credentials(
GrpcBasicAuth(username, password)),
),
)
self.client = ArmadaClient(channel)
def job_submit_request_items_for_test(self):
pod = core_v1.PodSpec(containers=[
core_v1.Container(
name="Container1",
image="index.docker.io/library/ubuntu:latest",
args=["sleep", "10s"],
securityContext=core_v1.SecurityContext(runAsUser=1000),
resources=core_v1.ResourceRequirements(
requests={
"cpu": api_resource.Quantity(string="120m"),
"memory": api_resource.Quantity(string="510Mi"),
},
limits={
"cpu": api_resource.Quantity(string="120m"),
"memory": api_resource.Quantity(string="510Mi"),
},
),
)
], )
return [self.client.create_job_request_item(priority=1, pod_spec=pod)]
def submit_test_job(self, queue, job_set_id):
jsr_items = self.job_submit_request_items_for_test()
submit_pb2.JobSubmitRequest(queue="test",
job_set_id=job_set_id,
job_request_items=jsr_items)
self.client.submit_jobs(queue, job_set_id, jsr_items)
def test_watch_events(self):
queue_name = "test"
job_set_id = f"set-{uuid.uuid1()}"
self.client.delete_queue(name=queue_name)
self.client.create_queue(name=queue_name, priority_factor=200)
self.submit_test_job(queue=queue_name, job_set_id=job_set_id)
self.client.cancel_jobs(queue=queue_name, job_set_id=job_set_id)
count = 0
def event_counter():
nonlocal count
count += 1
event_stream = self.client.watch_events(on_event=event_counter,
queue=queue_name,
job_set_id=job_set_id)
time.sleep(1)
print(count)
self.client.unwatch_events(event_stream)
"""
This is an example of a Airflow dag that uses a BashOperator and an ArmadaOperator
"""
with DAG(
dag_id="hello_armada",
start_date=pendulum.datetime(2016, 1, 1, tz="UTC"),
schedule_interval="@daily",
catchup=False,
default_args={"retries": 2},
) as dag:
"""
The ArmadaOperator requires a python client and a JobServiceClient
so we initialize them and set up their channel arguments.
"""
no_auth_client = ArmadaClient(channel=grpc.insecure_channel(
target="127.0.0.1:50051"))
job_service_client = JobServiceClient(channel=grpc.insecure_channel(
target="127.0.0.1:60003"))
"""
This defines an Airflow task that runs Hello World and it gives the airflow
task name of dummy.
"""
op = BashOperator(task_id="dummy", bash_command="echo Hello World!")
"""
This is creating an Armada task with the task_id of armada and name of armada.
The Airflow operator needs queue and job-set for Armada
You also specify the PythonClient and JobServiceClient for each task.
You should reuse them for all your tasks.
This job will use the podspec defined above.
"""
armada = ArmadaOperator(
yield
server.stop(False)
@pytest.fixture(scope="session", autouse=True)
def job_service_mock():
server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
jobservice_pb2_grpc.add_JobServiceServicer_to_server(JobService(), server)
server.add_insecure_port("[::]:60008")
server.start()
yield
server.stop(False)
tester_client = ArmadaClient(grpc.insecure_channel(target="127.0.0.1:50052", ))
tester_jobservice = JobServiceClient(
grpc.insecure_channel(target="127.0.0.1:60008"))
def sleep_job():
pod = core_v1.PodSpec(containers=[
core_v1.Container(
name="container-1",
image="busybox",
args=["sleep", "10s"],
securityContext=core_v1.SecurityContext(runAsUser=1000),
resources=core_v1.ResourceRequirements(
requests={
"cpu": api_resource.Quantity(string="120m"),
"memory": api_resource.Quantity(string="510Mi"),