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', ), )