def export_schema(args):
"""Export to JSON Schemas
:param args: args from CLI
"""
if args.yaml_path:
dump_api = api_to_dict()
for yp in args.yaml_path:
f_name = (yp % __version__) if '%s' in yp else yp
with open(f_name, 'w', encoding='utf8') as fp:
JAML.dump(dump_api, fp)
default_logger.info(f'API is exported to {f_name}')
if args.json_path:
dump_api = api_to_dict()
for jp in args.json_path:
f_name = (jp % __version__) if '%s' in jp else jp
with open(f_name, 'w', encoding='utf8') as fp:
json.dump(dump_api, fp, sort_keys=True)
default_logger.info(f'API is exported to {f_name}')
if args.schema_path:
dump_api = get_full_schema()
for jp in args.schema_path:
f_name = (jp % __version__) if '%s' in jp else jp
with open(f_name, 'w', encoding='utf8') as fp:
json.dump(dump_api, fp, sort_keys=True)
default_logger.info(f'API is exported to {f_name}')
def export_api(args: 'Namespace'):
import json
from .export import api_to_dict
from jina.jaml import JAML
from jina import __version__
from jina.logging import default_logger
from jina.schemas import get_full_schema
if args.yaml_path:
dump_api = api_to_dict()
for yp in args.yaml_path:
f_name = (yp % __version__) if '%s' in yp else yp
with open(f_name, 'w', encoding='utf8') as fp:
JAML.dump(dump_api, fp)
default_logger.info(f'API is exported to {f_name}')
if args.json_path:
dump_api = api_to_dict()
for jp in args.json_path:
f_name = (jp % __version__) if '%s' in jp else jp
with open(f_name, 'w', encoding='utf8') as fp:
json.dump(dump_api, fp, sort_keys=True)
default_logger.info(f'API is exported to {f_name}')
if args.schema_path:
dump_api = get_full_schema()
for jp in args.schema_path:
f_name = (jp % __version__) if '%s' in jp else jp
with open(f_name, 'w', encoding='utf8') as fp:
json.dump(dump_api, fp, sort_keys=True)
default_logger.info(f'API is exported to {f_name}')
def test_dump_load_build(monkeypatch):
f: Flow = Flow.load_config('''
jtype: Flow
with:
name: abc
port: 12345
protocol: http
executors:
- name: executor1
port: 45678
shards: 2
- name: executor2
uses: docker://exec
host: 1.2.3.4
- name: executor3
uses: docker://exec
shards: 2
''').build()
f1: Flow = Flow.load_config(JAML.dump(f)).build()
# these were passed by the user
assert f.port == f1.port
assert f.protocol == f1.protocol
assert f['executor1'].args.port == f1['executor1'].args.port
assert f['executor2'].args.host == f1['executor2'].args.host
# this was set during `load_config`
assert f['executor2'].args.port == f1['executor2'].args.port
monkeypatch.setenv('JINA_FULL_CLI', 'true')
f2: Flow = Flow.load_config(JAML.dump(f)).build()
# these were passed by the user
assert f.port == f2.port
# validate gateway args (set during build)
assert f['gateway'].args.port == f2['gateway'].args.port
def test_class_yaml():
class DummyClass:
pass
JAML.register(DummyClass)
a = JAML.load('!DummyClass {}')
assert type(a) == DummyClass
with open(
resource_filename(
'jina', '/'.join(
('resources',
'executors.requests.BaseExecutor.yml')))) as fp:
b = fp.read()
print(b)
c = JAML.load(b)
print(c)
args = set_pea_parser().parse_args([])
with BasePea(args):
pass
from jina.executors.requests import _defaults
assert _defaults is not None
def test_class_yaml2():
with open(
resource_filename(
'jina', '/'.join(
('resources',
'executors.requests.BaseExecutor.yml')))) as fp:
JAML.load(fp)
def test_escape(original, escaped):
assert JAML.escape(original,
include_unknown_tags=False).strip() == escaped.strip()
assert (JAML.unescape(
JAML.escape(original, include_unknown_tags=False),
include_unknown_tags=False,
).strip() == original.strip())
def test_class_yaml():
class DummyClass:
pass
JAML.register(DummyClass)
a = JAML.load('!DummyClass {}')
assert type(a) == DummyClass
def test_dump_load_build(monkeypatch):
f: Flow = Flow.load_config(
'''
jtype: Flow
with:
name: abc
port_expose: 12345
protocol: http
executors:
- name: executor1
port_in: 45678
shards: 2
- name: executor2
uses: docker://exec
host: 1.2.3.4
- name: executor3
uses: docker://exec
shards: 2
'''
).build()
f['gateway'].args.runs_in_docker = True
f['executor1'].args.runs_in_docker = True
f1: Flow = Flow.load_config(JAML.dump(f)).build()
assert not f1[
'gateway'
].args.runs_in_docker # gateway doesn't have custom args set, as env was not set
assert f1['executor1'].args.runs_in_docker
# these were passed by the user
assert f.port_expose == f1.port_expose
assert f.protocol == f1.protocol
assert f['executor1'].args.port_in == f1['executor1'].args.port_in
assert f['executor2'].args.host == f1['executor2'].args.host
# this was set during `load_config`
assert f['executor2'].args.port_in == f1['executor2'].args.port_in
assert f['executor3'].args.port_out == f1['executor3'].args.port_out
# gateway args are not set, if `JINA_FULL_CLI` is not set
assert f['gateway'].args.port_in != f1['gateway'].args.port_in
assert f['gateway'].args.port_out != f1['gateway'].args.port_out
monkeypatch.setenv('JINA_FULL_CLI', 'true')
f2: Flow = Flow.load_config(JAML.dump(f)).build()
assert f2['gateway'].args.runs_in_docker
assert f2['executor1'].args.runs_in_docker
# these were passed by the user
assert f.port_expose == f2.port_expose
# validate gateway args (set during build)
assert f['gateway'].args.port_in == f2['gateway'].args.port_in
assert f['gateway'].args.port_out == f2['gateway'].args.port_out
assert f['gateway'].args.port_ctrl == f2['gateway'].args.port_ctrl
def _write_optimization_parameter(executor_configurations, target_file,
overwrite_parameter_file):
output = [
parameter for config in executor_configurations.values()
for parameter in config
]
if os.path.exists(target_file) and not overwrite_parameter_file:
logger.warning(
f"{target_file} already exists. Skip writing. Please remove it before parameter discovery."
)
else:
with open(target_file, "w") as outfile:
JAML.dump(output, outfile)
def register_class(cls):
"""
Register the class for dumping loading.
:param cls: Target class.
:return: Registered class.
"""
reg_cls_set = getattr(cls, '_registered_class', set())
if cls.__name__ not in reg_cls_set:
reg_cls_set.add(cls.__name__)
setattr(cls, '_registered_class', reg_cls_set)
from jina.jaml import JAML
JAML.register(cls)
return cls
def __init__(self, args: 'argparse.Namespace'):
super().__init__(args)
self.runtime_cls = self._get_runtime_cls()
self.worker = _get_worker(
args=args,
target=run,
kwargs={
'args':
args,
'name':
self.name,
'envs':
self._envs,
'is_started':
self.is_started,
'is_shutdown':
self.is_shutdown,
'is_ready':
self.is_ready,
# the cancel event is only necessary for threads, otherwise runtimes should create and use the asyncio event
'cancel_event':
self.cancel_event
if getattr(args, 'runtime_backend', RuntimeBackendType.THREAD)
== RuntimeBackendType.THREAD else None,
'runtime_cls':
self.runtime_cls,
'jaml_classes':
JAML.registered_classes(),
},
name=self.name,
)
def test_hub_build_push():
args = set_hub_build_parser().parse_args([str(cur_dir / 'hub-mwu'), '--push', '--host-info'])
summary = HubIO(args).build()
with open(cur_dir / 'hub-mwu' / 'manifest.yml') as fp:
manifest = JAML.load(fp)
assert summary['is_build_success']
assert manifest['version'] == summary['version']
assert manifest['description'] == summary['manifest_info']['description']
assert manifest['author'] == summary['manifest_info']['author']
assert manifest['kind'] == summary['manifest_info']['kind']
assert manifest['type'] == summary['manifest_info']['type']
assert manifest['vendor'] == summary['manifest_info']['vendor']
assert manifest['keywords'] == summary['manifest_info']['keywords']
args = set_hub_list_parser().parse_args([
'--name', summary['manifest_info']['name'],
'--keywords', summary['manifest_info']['keywords'][0],
'--type', summary['manifest_info']['type']
])
response = HubIO(args).list()
manifests = response
assert len(manifests) >= 1
assert manifests[0]['name'] == summary['manifest_info']['name']
def assert_bi():
b = BaseIndexer(1)
b.save_config(os.path.join(tmpdir, 'tmp.yml'))
with open(os.path.join(tmpdir, 'tmp.yml')) as fp:
b = JAML.load(fp)
assert b.a == 1
def add_handlers(self, config_path: Optional[str] = None, **kwargs):
"""
Add handlers from config file.
:param config_path: Path of config file.
:param kwargs: Extra parameters.
"""
self.logger.handlers = []
with open(config_path) as fp:
config = JAML.load(fp)
for h in config['handlers']:
cfg = config['configs'].get(h, None)
fmt = getattr(formatter, cfg.get('formatter', 'Formatter'))
if h not in self.supported or not cfg:
raise ValueError(
f'can not find configs for {h}, maybe it is not supported')
handler = None
if h == 'StreamHandler':
handler = logging.StreamHandler(sys.stdout)
handler.setFormatter(fmt(cfg['format'].format_map(kwargs)))
elif h == 'SysLogHandler' and not __windows__:
if cfg['host'] and cfg['port']:
handler = SysLogHandlerWrapper(address=(cfg['host'],
cfg['port']))
else:
# a UNIX socket is used
if platform.system() == 'Darwin':
handler = SysLogHandlerWrapper(
address='/var/run/syslog')
else:
handler = SysLogHandlerWrapper(address='/dev/log')
if handler:
handler.ident = cfg.get('ident', '')
handler.setFormatter(fmt(cfg['format'].format_map(kwargs)))
try:
handler._connect_unixsocket(handler.address)
except OSError:
handler = None
pass
elif h == 'FileHandler':
filename = cfg['output'].format_map(kwargs)
if __windows__:
# colons are not allowed in filenames
filename = filename.replace(':', '.')
handler = logging.FileHandler(filename, delay=True)
handler.setFormatter(fmt(cfg['format'].format_map(kwargs)))
if handler:
self.logger.addHandler(handler)
verbose_level = LogVerbosity.from_string(config['level'])
if 'JINA_LOG_LEVEL' in os.environ:
verbose_level = LogVerbosity.from_string(
os.environ['JINA_LOG_LEVEL'])
self.logger.setLevel(verbose_level.value)
def test_load_yaml1(tmpdir):
with open(os.path.join(cur_dir, 'yaml/test-driver.yml'), encoding='utf8') as fp:
a = JAML.load(fp)
assert isinstance(a[0], KVSearchDriver)
assert isinstance(a[1], ControlReqDriver)
assert isinstance(a[2], BaseDriver)
with open(os.path.join(tmpdir, 'test_driver.yml'), 'w', encoding='utf8') as fp:
JAML.dump(a[0], fp)
with open(os.path.join(tmpdir, 'test_driver.yml'), encoding='utf8') as fp:
b = JAML.load(fp)
assert isinstance(b, KVSearchDriver)
assert b._executor_name == a[0]._executor_name
def test_yaml_single_flow(tmpdir, config):
jsonlines_file = os.path.join(tmpdir, 'docs.jsonlines')
optimizer_yaml = f'''!FlowOptimizer
version: 1
with:
flow_runner: !SingleFlowRunner
with:
flow_yaml: '{os.path.join(cur_dir, 'flow.yml')}'
documents: {jsonlines_file}
request_size: 1
execution_endpoint: 'search'
evaluation_callback: !MeanEvaluationCallback {{}}
parameter_yaml: '{os.path.join(cur_dir, 'parameter.yml')}'
workspace_base_dir: {tmpdir}
n_trials: 5
'''
documents = document_generator(10)
with open(jsonlines_file, 'w') as f:
for document, groundtruth_doc in documents:
document.id = ""
groundtruth_doc.id = ""
json.dump(
{
'document': json.loads(MessageToJson(document).replace('\n', '')),
'groundtruth': json.loads(
MessageToJson(groundtruth_doc).replace('\n', '')
),
},
f,
)
f.write('\n')
optimizer = JAML.load(optimizer_yaml)
result = optimizer.optimize_flow()
validate_result(result, tmpdir)
def test_hub_build_push_push_again():
args = set_hub_build_parser().parse_args([str(cur_dir / 'hub-mwu'), '--push', '--host-info'])
summary = HubIO(args).build()
with open(cur_dir / 'hub-mwu' / 'manifest.yml') as fp:
manifest = JAML.load(fp)
assert summary['is_build_success']
assert manifest['version'] == summary['version']
assert manifest['description'] == summary['manifest_info']['description']
assert manifest['author'] == summary['manifest_info']['author']
assert manifest['kind'] == summary['manifest_info']['kind']
assert manifest['type'] == summary['manifest_info']['type']
assert manifest['vendor'] == summary['manifest_info']['vendor']
assert manifest['keywords'] == summary['manifest_info']['keywords']
args = set_hub_list_parser().parse_args([
'--name', summary['manifest_info']['name'],
'--keywords', summary['manifest_info']['keywords'][0],
'--type', summary['manifest_info']['type']
])
response = HubIO(args).list()
manifests = response
assert len(manifests) >= 1
assert manifests[0]['name'] == summary['manifest_info']['name']
with pytest.raises(ImageAlreadyExists):
# try and push same version again should fail with `--no-overwrite`
args = set_hub_build_parser().parse_args([str(cur_dir / 'hub-mwu'), '--push', '--host-info', '--no-overwrite'])
HubIO(args).build()
def test_exec_type(tmpdir):
from jina.executors.indexers import BaseIndexer
assert 'BaseIndexer' in BaseExecutor._registered_class
# init from YAML should be okay as well
BaseExecutor.load_config('BaseIndexer')
BaseIndexer().save_config(os.path.join(tmpdir, 'tmp.yml'))
with open(os.path.join(tmpdir, 'tmp.yml')) as fp:
_ = JAML.load(fp)
def assert_bi():
b = BaseIndexer(1)
b.save_config(os.path.join(tmpdir, 'tmp.yml'))
with open(os.path.join(tmpdir, 'tmp.yml')) as fp:
b = JAML.load(fp)
assert b.a == 1
# we override BaseIndexer now, without force it shall not store all init values
class BaseIndexer(BaseExecutor):
def __init__(self, a=0):
super().__init__()
self.a = a
with pytest.raises(AssertionError):
assert_bi()
class BaseIndexer(BaseExecutor):
force_register = True
def __init__(self, a=0):
super().__init__()
self.a = a
assert_bi()
def test_yaml_expand4():
os.environ['ENV1'] = 'a'
os.environ['ENV2'] = '{"1": "2"}'
with open(os.path.join(cur_dir, 'yaml/test-expand4.yml')) as fp:
b = JAML.load(
fp,
substitute=True,
context={
'context_var': 3.14,
'context_var2': 'hello-world'
},
)
assert b['components'][0]['metas']['bad_var'] == 'real-compound'
assert b['components'][1]['metas']['bad_var'] == 2
assert b['components'][1]['metas']['float_var'] == 0.232
assert b['components'][1]['metas']['mixed'] == '0.232-2-real-compound'
assert b['components'][1]['metas']['name_shortcut'] == 'test_numpy'
assert b['components'][1]['metas']['mixed_env'] == '0.232-a'
assert b['components'][1]['metas']['random_id'] == 3.14
assert b['components'][1]['metas']['config_str'] == 'hello-world'
assert b['components'][1]['metas']['bracket_env'] == '{"1": "2"}'
assert b['components'][1]['metas']['bracket_env'] == '{"1": "2"}'
assert b['components'][1]['metas']['context_dot'] == 3.14
def test_hub_build_push(monkeypatch, mocker):
monkeypatch.setattr(Path, 'is_file', True)
mock_access_token = mocker.patch.object(hubapi,
'_fetch_access_token',
autospec=True)
mock_access_token.return_value = os.environ.get('GITHUB_TOKEN', None)
args = set_hub_build_parser().parse_args(
[str(cur_dir + '/hub-mwu'), '--push', '--host-info'])
summary = HubIO(args).build()
with open(cur_dir + '/hub-mwu' + '/manifest.yml') as fp:
manifest_jaml = JAML.load(fp, substitute=True)
manifest = expand_dict(manifest_jaml)
assert summary['is_build_success']
assert manifest['version'] == summary['version']
assert manifest['description'] == summary['manifest_info']['description']
assert manifest['author'] == summary['manifest_info']['author']
assert manifest['kind'] == summary['manifest_info']['kind']
assert manifest['type'] == summary['manifest_info']['type']
assert manifest['vendor'] == summary['manifest_info']['vendor']
assert manifest['keywords'] == summary['manifest_info']['keywords']
args = set_hub_list_parser().parse_args([
'--name', summary['manifest_info']['name'], '--keywords',
summary['manifest_info']['keywords'][0], '--type',
summary['manifest_info']['type']
])
response = HubIO(args).list()
manifests = response
assert len(manifests) >= 1
assert manifests[0]['name'] == summary['manifest_info']['name']
def test_include_unknown(include_unk, expected):
y = '''
!BaseExecutor {}
!Blah {}
'''
assert JAML.escape(
y, include_unknown_tags=include_unk).strip() == expected.strip()
def test_yaml_expand3():
with open(os.path.join(cur_dir, 'yaml/test-expand3.yml')) as fp:
a = JAML.load(fp)
b = expand_dict(a)
assert b['max_snapshot'] == 0
assert b['pea_workspace'] != '{root.workspace}/{root.name}-{this.pea_id}'
def test_check_platform():
with resource_stream(
'jina', '/'.join(
('resources', 'hub-builder', 'platforms.yml'))) as fp:
platforms = JAML.load(fp)
check_platform(platforms)
with pytest.raises(ValueError):
check_platform(platforms + ['invalid'])
def test_check_licenses():
with resource_stream(
'jina', '/'.join(
('resources', 'hub-builder', 'osi-approved.yml'))) as fp:
licenses = JAML.load(fp)
for lic in licenses:
check_license(lic)
with pytest.raises(ValueError):
check_license('invalid')
def test_yaml_expand2():
with open(os.path.join(cur_dir, 'yaml/test-expand2.yml')) as fp:
a = JAML.load(fp)
os.environ['ENV1'] = 'a'
b = expand_dict(a)
assert b['components'][0]['metas']['bad_var'] == 'real-compound'
assert b['components'][1]['metas']['bad_var'] == 2
assert b['components'][1]['metas']['float_var'] == 0.232
assert b['components'][1]['metas']['mixed'] == '0.232-2-real-compound'
assert b['components'][1]['metas']['mixed_env'] == '0.232-a'
assert b['components'][1]['metas']['name_shortcut'] == 'test_numpy'
def _add_metas(self, _metas: Optional[Dict]):
from jina.serve.executors.metas import get_default_metas
tmp = get_default_metas()
if _metas:
tmp.update(_metas)
unresolved_attr = False
target = SimpleNamespace()
# set self values filtered by those non-exist, and non-expandable
for k, v in tmp.items():
if k == 'workspace' and not (v is None or v == ''):
warnings.warn(
'Setting `workspace` via `metas.workspace` is deprecated. '
'Instead, use `f.add(..., workspace=...)` when defining a a Flow in Python; '
'the `workspace` parameter when defining a Flow using YAML; '
'or `--workspace` when starting an Executor using the CLI.',
category=DeprecationWarning,
)
if not hasattr(target, k):
if isinstance(v, str):
if not env_var_regex.findall(v):
setattr(target, k, v)
else:
unresolved_attr = True
else:
setattr(target, k, v)
elif type(getattr(target, k)) == type(v):
setattr(target, k, v)
if unresolved_attr:
_tmp = vars(self)
_tmp['metas'] = tmp
new_metas = JAML.expand_dict(_tmp)['metas']
for k, v in new_metas.items():
if not hasattr(target, k):
if isinstance(v, str):
if not (env_var_regex.findall(v)
or internal_var_regex.findall(v)):
setattr(target, k, v)
else:
raise ValueError(
f'{k}={v} is not substitutable or badly referred'
)
else:
setattr(target, k, v)
# `name` is important as it serves as an identifier of the executor
# if not given, then set a name by the rule
if not getattr(target, 'name', None):
setattr(target, 'name', self.__class__.__name__)
self.metas = target
def export_api(args: 'Namespace'):
from .export import api_to_dict
from jina import __version__
from jina.logging import default_logger
if args.yaml_path:
for yp in args.yaml_path:
f_name = (yp % __version__) if '%s' in yp else yp
from jina.jaml import JAML
with open(f_name, 'w', encoding='utf8') as fp:
JAML.dump(api_to_dict(), fp)
default_logger.info(f'API is exported to {f_name}')
if args.json_path:
for jp in args.json_path:
f_name = (jp % __version__) if '%s' in jp else jp
import json
with open(f_name, 'w', encoding='utf8') as fp:
json.dump(api_to_dict(), fp, sort_keys=True)
default_logger.info(f'API is exported to {f_name}')
def test_yaml_expand():
with open(os.path.join(cur_dir, 'yaml/test-expand.yml')) as fp:
a = JAML.load(fp)
b = expand_dict(a)
assert b['quote_dict'] == {}
assert b['quote_string'].startswith('{')
assert b['quote_string'].endswith('}')
assert b['nest']['quote_dict'] == {}
assert b['nest']['quote_string'].startswith('{')
assert b['nest']['quote_string'].endswith('}')
assert b['exist_env'] != '$PATH'
assert b['non_exist_env'] == '$JINA_WHATEVER_ENV'
def parse(self, cls: Type['JAMLCompatible'],
data: Dict) -> 'JAMLCompatible':
"""
:param cls: target class type to parse into, must be a :class:`JAMLCompatible` type
:param data: flow yaml file loaded as python dict
:return: the YAML parser given the syntax version number
"""
expanded_data = JAML.expand_dict(data, None)
if 'with' in data:
obj = cls(**expanded_data.get('with', {}))
else:
obj = cls(**expanded_data)
return obj
def test_sse_client(tmpdir):
class WrapAssert:
def __init__(self):
self.count = 0
conf_path = os.path.join(tmpdir, 'log')
path = os.path.join(conf_path, GROUP_ID)
event = threading.Event()
feed_thread = threading.Thread(name='feed_path_logs',
target=feed_path_logs,
daemon=False,
kwargs={
'path': path,
'threading_event': event
})
with open(os.path.join(cur_dir, 'logserver_config.yml')) as fp:
log_config = JAML.load(fp)
log_config['files']['log'] = conf_path
log_config['port'] = RANDOM_PORT
sse_server_thread = threading.Thread(name='sentinel-sse-logger',
target=start_sse_logger,
daemon=False,
args=(log_config, GROUP_ID, None))
wrap = WrapAssert()
sse_client_thread = threading.Thread(name='sse-client',
target=sse_client,
daemon=False,
kwargs=({
'wrap': wrap
}))
feed_thread.start()
time.sleep(0.5)
sse_server_thread.start()
time.sleep(0.5)
sse_client_thread.start()
time.sleep(0.5)
event.set()
stop_log_server()
feed_thread.join()
sse_server_thread.join()
sse_client_thread.join()
assert wrap.count > 0
