def test_port_configuration(replicas_and_shards):
def validate_ports_replica(shard, replica_port_in, replica_port_out, shards):
assert replica_port_in == shard.args.port_in
assert shard.args.port_out == replica_port_out
peas_args = shard.peas_args
peas = peas_args['peas']
assert len(peas) == shards
if shards == 1:
assert peas_args['head'] is None
assert peas_args['tail'] is None
assert peas[0].port_in == replica_port_in
assert peas[0].port_out == replica_port_out
else:
shard_head = peas_args['head']
shard_tail = peas_args['tail']
assert shard.args.port_in == shard_head.port_in
assert shard.args.port_out == shard_tail.port_out
for pea in peas:
assert shard_head.port_out == pea.port_in
assert pea.port_out == shard_tail.port_in
flow = Flow()
for i, (replicas, shards) in enumerate(replicas_and_shards):
flow.add(
name=f'pod{i}',
replicas=replicas,
shards=shards,
copy_flow=False,
)
with flow:
pods = flow._pod_nodes
for pod_name, pod in pods.items():
if pod_name == 'gateway':
continue
if pod.args.replicas == 1:
if int(pod.args.shards) == 1:
assert len(pod.peas_args['peas']) == 1
else:
assert len(pod.peas_args) == 3
shard_port_in = pod.args.port_in
shard_port_out = pod.args.port_out
else:
shard_port_in = pod.head_args.port_out
shard_port_out = pod.tail_args.port_in
assert pod.head_args.port_in == pod.args.port_in
assert pod.head_args.port_out == shard_port_in
assert pod.tail_args.port_in == shard_port_out
assert pod.tail_args.port_out == pod.args.port_out
if pod.args.shards > 1:
for shard in pod.shards:
validate_ports_replica(
shard,
shard_port_in,
shard_port_out,
getattr(pod.args, 'replicas', 1),
)
assert pod
def test_flow_to_k8s_yaml_external_pod(tmpdir, has_external):
flow = Flow(name='test-flow', port=8080).add(name='executor0', )
if has_external:
flow = flow.add(name='external_executor',
external=True,
host='1.2.3.4',
port=9090)
else:
flow = flow.add(name='external_executor')
dump_path = os.path.join(str(tmpdir), 'test_flow_k8s')
namespace = 'test-flow-ns'
flow.to_kubernetes_yaml(
output_base_path=dump_path,
k8s_namespace=namespace,
)
yaml_dicts_per_deployment = {
'gateway': [],
'executor0': [],
}
assert len(set(os.listdir(dump_path))) == 2 if has_external else 3
for pod_name in set(os.listdir(dump_path)):
file_set = set(os.listdir(os.path.join(dump_path, pod_name)))
for file in file_set:
with open(os.path.join(dump_path, pod_name, file)) as f:
yml_document_all = list(yaml.safe_load_all(f))
yaml_dicts_per_deployment[file[:-4]] = yml_document_all
gateway_objects = yaml_dicts_per_deployment['gateway']
gateway_args = gateway_objects[2]['spec']['template']['spec'][
'containers'][0]['args']
assert (
gateway_args[gateway_args.index('--graph-description') + 1] ==
'{"executor0": ["external_executor"], "start-gateway": ["executor0"], "external_executor": ["end-gateway"]}'
)
if has_external:
assert '--deployments-addresses' in gateway_args
assert (
gateway_args[gateway_args.index('--deployments-addresses') + 1] ==
'{"executor0": ["grpc://executor0.test-flow-ns.svc:8080"], "external_executor": ["grpc://1.2.3.4:9090"]}'
)
else:
assert '--deployments-addresses' in gateway_args
assert (
gateway_args[gateway_args.index('--deployments-addresses') + 1] ==
'{"executor0": ["grpc://executor0.test-flow-ns.svc:8080"], "external_executor": ["grpc://external-executor.test-flow-ns.svc:8080"]}'
)
def test_flow_equalities():
f1 = Flow().add().add(needs='gateway').needs_all(name='joiner')
f2 = Flow().add(name='pod0').add(name='pod1', needs='gateway').add(name='joiner', needs=['pod0', 'pod1'])
assert f1 == f2
f2 = f2.add()
assert f1 != f2
def test_port_configuration(replicas_and_parallel):
def extract_pod_args(pod):
if 'replicas' not in pod.args or int(pod.args.replicas) == 1:
head_args = pod.peas_args['head']
tail_args = pod.peas_args['tail']
middle_args = pod.peas_args['peas']
else:
head_args = pod.head_args
tail_args = pod.tail_args
middle_args = pod.replicas_args
return pod, head_args, tail_args, middle_args
def get_outer_ports(pod, head_args, tail_args, middle_args):
if not 'replicas' in pod.args or int(pod.args.replicas) == 1:
if not 'parallel' in pod.args or int(pod.args.parallel) == 1:
assert tail_args is None
assert head_args is None
replica = middle_args[0] # there is only one
return replica.port_in, replica.port_out
else:
return pod.head_args.port_in, pod.tail_args.port_out
else:
assert pod.args.replicas == len(middle_args)
return pod.head_args.port_in, pod.tail_args.port_out
def validate_ports_pods(pods):
for i in range(len(pods) - 1):
_, port_out = get_outer_ports(*extract_pod_args(pods[i]))
port_in_next, _ = get_outer_ports(*extract_pod_args(pods[i + 1]))
assert port_out == port_in_next
def validate_ports_replica(replica, replica_port_in, replica_port_out,
parallel):
assert replica_port_in == replica.args.port_in
assert replica.args.port_out == replica_port_out
peas_args = replica.peas_args
peas = peas_args['peas']
assert len(peas) == parallel
if parallel == 1:
assert peas_args['head'] is None
assert peas_args['tail'] is None
assert peas[0].port_in == replica_port_in
assert peas[0].port_out == replica_port_out
else:
shard_head = peas_args['head']
shard_tail = peas_args['tail']
assert replica.args.port_in == shard_head.port_in
assert replica.args.port_out == shard_tail.port_out
for pea in peas:
assert shard_head.port_out == pea.port_in
assert pea.port_out == shard_tail.port_in
flow = Flow()
for i, (replicas, parallel) in enumerate(replicas_and_parallel):
flow.add(
name=f'pod{i}',
replicas=replicas,
parallel=parallel,
port_in=f'51{i}00',
# info: needs to be set in this test since the test is asserting pod args with pod tail args
port_out=
f'51{i + 1}00', # outside this test, it don't have to be set
copy_flow=False,
)
with flow:
pods = flow._pod_nodes
validate_ports_pods(
[pods['gateway']] +
[pods[f'pod{i}']
for i in range(len(replicas_and_parallel))] + [pods['gateway']])
for pod_name, pod in pods.items():
if pod_name == 'gateway':
continue
if pod.args.replicas == 1:
if int(pod.args.parallel) == 1:
assert len(pod.peas_args['peas']) == 1
else:
assert len(pod.peas_args) == 3
replica_port_in = pod.args.port_in
replica_port_out = pod.args.port_out
else:
replica_port_in = pod.head_args.port_out
replica_port_out = pod.tail_args.port_in
assert pod.head_pea.args.port_in == pod.args.port_in
assert pod.head_pea.args.port_out == replica_port_in
assert pod.tail_pea.args.port_in == replica_port_out
assert pod.tail_pea.args.port_out == pod.args.port_out
if pod.args.replicas > 1:
for replica in pod.replicas:
validate_ports_replica(
replica,
replica_port_in,
replica_port_out,
getattr(pod.args, 'parallel', 1),
)
assert pod
