def test_docker_compose_gateway(deployments_addresses):
args = set_gateway_parser().parse_args([
'--env', 'ENV_VAR:ENV_VALUE', '--port-expose', '32465', '--port-in',
'33455'
]) # envs are
# ignored for gateway
deployment_config = DockerComposeConfig(
args, deployments_addresses=deployments_addresses)
name, gateway_config = deployment_config.to_docker_compose_config()[0]
assert name == 'gateway'
assert gateway_config['image'] == 'jinaai/jina:test-pip'
assert gateway_config['entrypoint'] == ['jina']
assert gateway_config['ports'] == [
f'{args.port_expose}:{args.port_expose}',
f'{args.port_in}:{args.port_in}',
]
assert gateway_config['expose'] == [
f'{args.port_expose}', f'{args.port_in}'
]
args = gateway_config['command']
assert args[0] == 'gateway'
assert '--port-in' in args
assert args[args.index('--port-in') + 1] == '33455'
assert '--port-expose' in args
assert args[args.index('--port-expose') + 1] == '32465'
assert '--env' not in args
assert '--pod-role' in args
assert args[args.index('--pod-role') + 1] == 'GATEWAY'
if deployments_addresses is not None:
assert '--deployments-addresses' in args
assert args[args.index('--deployments-addresses') +
1] == json.dumps(deployments_addresses)
def test_docker_compose_gateway(deployments_addresses, custom_gateway):
if custom_gateway:
os.environ['JINA_GATEWAY_IMAGE'] = custom_gateway
elif 'JINA_GATEWAY_IMAGE' in os.environ:
del os.environ['JINA_GATEWAY_IMAGE']
args = set_gateway_parser().parse_args(
['--env', 'ENV_VAR:ENV_VALUE', '--port', '32465']) # envs are
# ignored for gateway
deployment_config = DockerComposeConfig(
args, deployments_addresses=deployments_addresses)
name, gateway_config = deployment_config.to_docker_compose_config()[0]
assert name == 'gateway'
assert (
gateway_config['image'] == custom_gateway if custom_gateway else
f'jinaai/jina:{deployment_config.worker_services[0].version}-py38-standard'
)
assert gateway_config['entrypoint'] == ['jina']
assert gateway_config['ports'] == [f'{args.port}:{args.port}']
assert gateway_config['expose'] == [f'{args.port}']
args = gateway_config['command']
assert args[0] == 'gateway'
assert '--port' in args
assert args[args.index('--port') + 1] == '32465'
assert '--env' not in args
assert '--pod-role' in args
assert args[args.index('--pod-role') + 1] == 'GATEWAY'
if deployments_addresses is not None:
assert '--deployments-addresses' in args
assert args[args.index('--deployments-addresses') +
1] == json.dumps(deployments_addresses)
def test_worker_services(name: str, shards: str):
args = set_deployment_parser().parse_args(
['--name', name, '--shards', shards])
deployment_config = DockerComposeConfig(args)
actual_services = deployment_config.worker_services
assert len(actual_services) == int(shards)
for i, deploy in enumerate(actual_services):
if int(shards) > 1:
assert deploy.name == f'{name}-{i}'
else:
assert deploy.name == name
assert deploy.jina_deployment_name == name
assert deploy.shard_id == i
def test_parse_args(
shards: int,
replicas: int,
uses_with,
uses_metas,
uses_before,
uses_after,
):
args_list = [
'--shards',
str(shards),
'--replicas',
str(replicas),
'--name',
'executor',
]
if uses_before is not None:
args_list.extend(['--uses-before', uses_before])
if uses_after is not None:
args_list.extend(['--uses-after', uses_after])
if uses_with is not None:
args_list.extend(['--uses-with', uses_with])
if uses_metas is not None:
args_list.extend(['--uses-metas', uses_metas])
args = set_deployment_parser().parse_args(args_list)
deployment_config = DockerComposeConfig(args)
assert namespace_equal(
deployment_config.services_args['head_service'],
args,
skip_attr=(
'runtime_cls',
'pod_role',
'port_in',
'k8s_namespace',
'k8s_connection_pool',
'name',
'uses',
'connection_list',
'uses_with',
'uses_metas',
'uses_before',
'uses_after',
'uses_before_address',
'uses_after_address',
'replicas',
),
)
assert deployment_config.services_args[
'head_service'].name == 'executor/head'
assert deployment_config.services_args[
'head_service'].runtime_cls == 'HeadRuntime'
assert deployment_config.services_args['head_service'].uses is None
assert deployment_config.services_args['head_service'].uses_before is None
assert deployment_config.services_args['head_service'].uses_after is None
assert deployment_config.services_args['head_service'].uses_metas is None
assert deployment_config.services_args['head_service'].uses_with is None
if uses_before is None:
assert (
deployment_config.services_args['head_service'].uses_before_address
is None)
else:
assert (deployment_config.services_args['head_service'].
uses_before_address == 'executor-uses-before:8081')
if uses_after is None:
assert (
deployment_config.services_args['head_service'].uses_after_address
is None)
else:
assert (deployment_config.services_args['head_service'].
uses_after_address == 'executor-uses-after:8081')
if shards > 1:
if replicas == 1:
candidate_connection_list = {
str(i): [f'executor-{i}:8081']
for i in range(shards)
}
else:
candidate_connection_list = {}
for shard_id in range(shards):
candidate_connection_list[str(shard_id)] = []
for replica in range(replicas): # TODO
candidate_connection_list[str(shard_id)].append(
f'executor-{shard_id}-rep-{replica}:8081')
else:
if replicas == 1:
candidate_connection_list = {'0': [f'executor:8081']}
else:
candidate_connection_list = {'0': []}
for replica in range(replicas): # TODO
candidate_connection_list['0'].append(
f'executor-rep-{replica}:8081')
assert deployment_config.services_args[
'head_service'].connection_list == json.dumps(
candidate_connection_list)
if uses_before is not None:
assert (deployment_config.services_args['uses_before_service'].name ==
'executor/uses-before')
assert (deployment_config.services_args['uses_before_service'].
runtime_cls == 'WorkerRuntime')
assert (deployment_config.services_args['uses_before_service'].uses ==
uses_before)
assert (
deployment_config.services_args['uses_before_service'].uses_before
is None)
assert deployment_config.services_args[
'uses_before_service'].uses_after is None
assert deployment_config.services_args[
'uses_before_service'].uses_metas is None
assert deployment_config.services_args[
'uses_before_service'].uses_with is None
assert deployment_config.services_args[
'uses_before_service'].shard_id == 0
assert deployment_config.services_args[
'uses_before_service'].replicas == 1
if uses_after is not None:
assert (deployment_config.services_args['uses_after_service'].name ==
'executor/uses-after')
assert (deployment_config.services_args['uses_after_service'].
runtime_cls == 'WorkerRuntime')
assert deployment_config.services_args[
'uses_after_service'].uses == uses_after
assert deployment_config.services_args[
'uses_after_service'].uses_before is None
assert deployment_config.services_args[
'uses_after_service'].uses_after is None
assert deployment_config.services_args[
'uses_after_service'].uses_metas is None
assert deployment_config.services_args[
'uses_after_service'].uses_with is None
assert deployment_config.services_args[
'uses_after_service'].shard_id == 0
assert deployment_config.services_args[
'uses_after_service'].replicas == 1
for i, depl_arg in enumerate(deployment_config.services_args['services']):
import copy
assert (depl_arg.name == f'executor-{i}'
if len(deployment_config.services_args['services']) > 1 else
'executor')
cargs = copy.deepcopy(args)
cargs.shard_id = i
assert depl_arg.k8s_connection_pool is False
assert namespace_equal(
depl_arg,
cargs,
skip_attr=(
'runtime_cls',
'pod_role',
'port_in',
'k8s_namespace',
'k8s_connection_pool',
'uses_before', # the uses_before and after is head business
'uses_after',
'name',
),
)
def test_docker_compose_yaml_regular_deployment(
uses_before,
uses_after,
uses,
shards,
replicas,
uses_with,
uses_metas,
polling,
monkeypatch,
):
def _mock_fetch(name, tag=None, secret=None, force=False):
return (
HubExecutor(
uuid='hello',
name='alias_dummy',
tag='v0',
image_name=f'jinahub/{name}',
md5sum=None,
visibility=True,
archive_url=None,
),
False,
)
monkeypatch.setattr(HubIO, 'fetch_meta', _mock_fetch)
args_list = [
'--name',
'executor',
'--uses',
uses,
'--env',
'ENV_VAR:ENV_VALUE',
'--replicas',
str(replicas),
'--shards',
str(shards),
'--polling',
str(polling),
]
if uses_before is not None:
args_list.extend(['--uses-before', uses_before])
if uses_after is not None:
args_list.extend(['--uses-after', uses_after])
if uses_with is not None:
args_list.extend(['--uses-with', uses_with])
if uses_metas is not None:
args_list.extend(['--uses-metas', uses_metas])
args = set_deployment_parser().parse_args(args_list)
# ignored for gateway
deployment_config = DockerComposeConfig(args)
yaml_configs = deployment_config.to_docker_compose_config()
assert len(yaml_configs) == 1 + shards * replicas + (
1 if uses_before else 0) + (1 if uses_after else 0)
head_name, head_config = yaml_configs[0]
assert head_name == 'executor-head'
assert head_config['image'] == 'jinaai/jina:test-pip'
assert head_config['entrypoint'] == ['jina']
head_args = head_config['command']
assert head_args[0] == 'executor'
assert '--native' in head_args
assert '--runtime-cls' in head_args
assert head_args[head_args.index('--runtime-cls') + 1] == 'HeadRuntime'
assert '--name' in head_args
assert head_args[head_args.index('--name') + 1] == 'executor/head'
assert '--port-in' in head_args
assert head_args[head_args.index('--port-in') + 1] == '8081'
assert '--env' not in head_args
assert '--pod-role' in head_args
assert head_args[head_args.index('--pod-role') + 1] == 'HEAD'
assert '--connection-list' in head_args
connection_list_string = head_args[head_args.index('--connection-list') +
1]
candidate_connection_list = {}
if shards > 1:
if replicas == 1:
candidate_connection_list = {
str(shard_id): [f'executor-{shard_id}:8081']
for shard_id in range(shards)
}
else:
candidate_connection_list = {}
for shard_id in range(shards):
candidate_connection_list[str(shard_id)] = []
for replica in range(replicas):
candidate_connection_list[str(shard_id)].append(
f'executor-{shard_id}-rep-{replica}:8081')
else:
if shards == 1:
if replicas == 1:
candidate_connection_list = {'0': [f'executor:8081']}
else:
candidate_connection_list = {'0': []}
for replica in range(replicas):
candidate_connection_list['0'].append(
f'executor-rep-{replica}:8081')
assert connection_list_string == json.dumps(candidate_connection_list)
if polling == 'ANY':
assert '--polling' not in head_args
else:
assert '--polling' in head_args
assert head_args[head_args.index('--polling') + 1] == 'ALL'
if uses_before is not None:
uses_before_name, uses_before_config = yaml_configs[1]
assert uses_before_name == 'executor-uses-before'
assert uses_before_config['image'] == 'jinahub/HubBeforeExecutor'
assert uses_before_config['entrypoint'] == ['jina']
uses_before_args = uses_before_config['command']
assert uses_before_args[0] == 'executor'
assert '--native' in uses_before_args
assert '--name' in uses_before_args
assert '--env' not in uses_before_args
assert (uses_before_args[uses_before_args.index('--name') +
1] == 'executor/uses-before')
assert '--port-in' in uses_before_args
assert uses_before_args[uses_before_args.index('--port-in') +
1] == '8081'
assert '--connection-list' not in uses_before_args
assert '--polling' not in uses_before_args
if uses_after is not None:
uses_after_name, uses_after_config = (
yaml_configs[1] if uses_before is None else yaml_configs[2])
assert uses_after_name == 'executor-uses-after'
assert uses_after_config['image'] == 'jinahub/HubAfterExecutor'
assert uses_after_config['entrypoint'] == ['jina']
uses_after_args = uses_after_config['command']
assert uses_after_args[0] == 'executor'
assert '--native' in uses_after_args
assert '--name' in uses_after_args
assert '--env' not in uses_after_args
assert (uses_after_args[uses_after_args.index('--name') +
1] == 'executor/uses-after')
assert '--port-in' in uses_after_args
assert uses_after_args[uses_after_args.index('--port-in') +
1] == '8081'
assert '--connection-list' not in uses_after_args
assert '--polling' not in uses_after_args
num_shards_replicas_configs = shards * replicas
shards_replicas_configs = yaml_configs[-num_shards_replicas_configs:]
for shard_id in range(shards):
replicas_configs = shards_replicas_configs[shard_id *
replicas:shard_id *
replicas + replicas]
for i_replica, (replica_name,
replica_config) in enumerate(replicas_configs):
expected_name = 'executor'
expected_arg_name = 'executor'
if shards > 1:
expected_name += f'-{shard_id}'
expected_arg_name += f'-{shard_id}'
if replicas > 1:
expected_name += f'-rep-{i_replica}'
expected_arg_name += f'/rep-{i_replica}'
assert replica_name == expected_name
assert (replica_config['image'] == 'docker_image:latest' if uses
== 'docker_image:latest' else 'jinahub/HubExecutor')
assert replica_config['entrypoint'] == ['jina']
replica_args = replica_config['command']
assert replica_args[0] == 'executor'
assert '--native' in replica_args
assert '--name' in replica_args
assert replica_args[replica_args.index('--name') +
1] == expected_arg_name
assert '--port-in' in replica_args
assert replica_args[replica_args.index('--port-in') + 1] == '8081'
assert '--env' in replica_args
assert (replica_args[replica_args.index('--env') +
1] == '{"ENV_VAR": "ENV_VALUE"}')
assert '--connection-list' not in replica_args
if uses_with is not None:
assert '--uses-with' in replica_args
assert replica_args[replica_args.index('--uses-with') +
1] == uses_with
else:
assert '--uses-with' not in replica_args
expected_uses_metas = {}
if uses_metas is not None:
expected_uses_metas = json.loads(uses_metas)
assert '--uses-metas' in replica_args
assert replica_args[replica_args.index('--uses-metas') +
1] == json.dumps(expected_uses_metas)
def test_parse_args_custom_executor(shards: int, replicas: int):
uses_before = 'custom-executor-before'
uses_after = 'custom-executor-after'
args = set_deployment_parser().parse_args([
'--shards',
str(shards),
'--replicas',
str(replicas),
'--uses-before',
uses_before,
'--uses-after',
uses_after,
'--name',
'executor',
])
deployment_config = DockerComposeConfig(args)
assert deployment_config.services_args[
'head_service'].runtime_cls == 'HeadRuntime'
assert deployment_config.services_args['head_service'].uses_before is None
assert (deployment_config.services_args['head_service'].uses_before_address
== 'executor-uses-before:8081')
assert deployment_config.services_args['head_service'].uses is None
assert deployment_config.services_args['head_service'].uses_after is None
assert (deployment_config.services_args['head_service'].uses_after_address
== f'executor-uses-after:8081')
assert deployment_config.services_args['head_service'].uses_before is None
assert deployment_config.services_args['head_service'].uses_after is None
assert (deployment_config.services_args['uses_before_service'].name ==
'executor/uses-before')
assert (deployment_config.services_args['uses_before_service'].runtime_cls
== 'WorkerRuntime')
assert deployment_config.services_args[
'uses_before_service'].uses == uses_before
assert deployment_config.services_args[
'uses_before_service'].uses_before is None
assert deployment_config.services_args[
'uses_before_service'].uses_after is None
assert deployment_config.services_args[
'uses_before_service'].uses_metas is None
assert deployment_config.services_args[
'uses_before_service'].uses_with is None
assert deployment_config.services_args['uses_before_service'].shard_id == 0
assert deployment_config.services_args['uses_before_service'].replicas == 1
assert (deployment_config.services_args['uses_after_service'].name ==
'executor/uses-after')
assert (deployment_config.services_args['uses_after_service'].runtime_cls
== 'WorkerRuntime')
assert deployment_config.services_args[
'uses_after_service'].uses == uses_after
assert deployment_config.services_args[
'uses_after_service'].uses_before is None
assert deployment_config.services_args[
'uses_after_service'].uses_after is None
assert deployment_config.services_args[
'uses_after_service'].uses_metas is None
assert deployment_config.services_args[
'uses_after_service'].uses_with is None
assert deployment_config.services_args['uses_after_service'].shard_id == 0
assert deployment_config.services_args['uses_after_service'].replicas == 1
for i, depl_arg in enumerate(deployment_config.services_args['services']):
import copy
assert (depl_arg.name == f'executor-{i}'
if len(deployment_config.services_args['services']) > 1 else
'executor')
cargs = copy.deepcopy(args)
cargs.shard_id = i
assert depl_arg.k8s_connection_pool is False
assert namespace_equal(
depl_arg,
cargs,
skip_attr=(
'uses_before',
'uses_after',
'port_in',
'k8s_namespace',
'k8s_connection_pool',
'name',
),
)