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
