def _build_with_pods(self, flow: Flow, pod_args: List[PodModel]): """ Since we rely on PodModel, this can accept all params that a Pod can accept """ for current_pod_args in pod_args: _current_pod_args = current_pod_args.dict() _current_pod_args.pop('log_config') flow = flow.add(**_current_pod_args) return flow
def query(yaml_path, num_docs=1, index_type='bytes', embed_dim=16): if index_type == 'bytes': input_fn = random_bytes_generator(num_docs=1, num_bytes=10) elif index_type == 'jina_pb2.Document': input_fn = random_docs_generator(num_docs=1, chunks_per_doc=5, embed_dim=embed_dim) elif index_type == 'sentences': pass f = Flow().load_config(filename=yaml_path) f = f.add(name='benchmark_pod', uses='yamls/pods/benchmark_driver.yml') with f: f.search(input_fn=input_fn) return list(f._pod_nodes.keys())
def _build_with_pods(self, pod_args: List[Dict]): """ Since we rely on SinglePodModel, this can accept all params that a Pod can accept """ flow = Flow() for current_pod_args in pod_args: # Hacky code here. We build `SinglePodModel` from `Dict` everytime to reset the default values SinglePodModel = build_pydantic_model(model_name='SinglePodModel', module='pod') _current_pod_args = SinglePodModel(**current_pod_args).dict() if not _current_pod_args.get('pod_role'): _current_pod_args.update(pod_role=PodRoleType.POD) _current_pod_args.pop('log_config') flow = flow.add(**_current_pod_args) return flow
def index(yaml_path, num_docs, input_type='bytes', batch_size=256, embed_dim=16, num_bytes_per_doc=10, num_chunks_per_doc=5, num_sentences_per_doc=5): if input_type == 'bytes': input_fn = random_bytes_generator(num_docs=num_docs, num_bytes=num_bytes_per_doc) elif input_type == 'jina_pb2.Document': input_fn = random_docs_generator(num_docs=num_docs, chunks_per_doc=num_chunks_per_doc, embed_dim=embed_dim) elif input_type == 'sentences': pass f = Flow().load_config(filename=yaml_path) f = f.add(name='benchmark_pod', uses=f'yamls/pods/benchmark_driver.yml') with f: f.index(input_fn=input_fn, batch_size=batch_size) return list(f._pod_nodes.keys())
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.replicas_args['head'] tail_args = pod.replicas_args['tail'] middle_args = pod.replicas_args['replicas'] 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.peas_args['head'].port_in, pod.peas_args[ 'tail'].port_out else: assert pod.args.replicas == len(middle_args) return pod.replicas_args['head'].port_in, pod.replicas_args[ 'tail'].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 # replica_head_out = pod.replicas_args['head'].port_out, # equals # replica_tail_in = pod.replicas_args['tail'].port_in, # equals for pea in pod.peas: if 'head' in pea.name: assert pea.args.port_in == pod.args.port_in assert pea.args.port_out == replica_port_in if 'tail' in pea.name: assert pea.args.port_in == replica_port_out assert pea.args.port_out == pod.args.port_out if pod.args.replicas > 1: for replica in pod.replica_list: validate_ports_replica( replica, replica_port_in, replica_port_out, getattr(pod.args, 'parallel', 1), ) assert pod