def register_metadata_store(metadata_store_name: str, metadata_store_type: str,
args: List[str]) -> None:
"""Register a metadata store."""
try:
parsed_args = cli_utils.parse_unknown_options(args)
except AssertionError as e:
cli_utils.error(str(e))
return
repo: Repository = Repository()
try:
# TODO [ENG-187]: Remove when we rework the registry logic
from zenml.core.component_factory import metadata_store_factory
comp = metadata_store_factory.get_single_component(metadata_store_type)
except AssertionError as e:
cli_utils.error(str(e))
return
metadata_store = comp(**parsed_args)
service = repo.get_service()
service.register_metadata_store(metadata_store_name, metadata_store)
cli_utils.declare(
f"Metadata Store `{metadata_store_name}` successfully registered!")
def set_metadata_store(store_type, args):
"""Set metadata store for local config."""
try:
parsed_args = parse_unknown_options(args)
except AssertionError as e:
click.echo(str(e))
return
# TODO: [LOW] Hard-coded
config = {'type': store_type, 'args': parsed_args}
from zenml.core.metadata.metadata_wrapper import ZenMLMetadataStore
store = ZenMLMetadataStore.from_config(config)
repo: Repository = Repository.get_instance()
repo.zenml_config.set_metadata_store(store)
click.echo(f'Metadata store set to: {store.to_config()}')
def register_orchestrator(orchestrator_name: str, orchestrator_type: str,
args: List[str]) -> None:
"""Register an orchestrator."""
try:
parsed_args = cli_utils.parse_unknown_options(args)
except AssertionError as e:
cli_utils.error(str(e))
return
repo: Repository = Repository()
# TODO [ENG-186]: Remove when we rework the registry logic
from zenml.core.component_factory import orchestrator_store_factory
comp = orchestrator_store_factory.get_single_component(orchestrator_type)
orchestrator_ = comp(**parsed_args)
service = repo.get_service()
service.register_orchestrator(orchestrator_name, orchestrator_)
cli_utils.declare(
f"Orchestrator `{orchestrator_name}` successfully registered!")
def test_parse_unknown_options_returns_a_dict_of_known_options() -> None:
"""Check that parse_unknown_options returns a dict of known options"""
parsed_sample_args = parse_unknown_options(SAMPLE_CUSTOM_ARGUMENTS)
assert isinstance(parsed_sample_args, dict)
assert len(parsed_sample_args.values()) == 3
assert parsed_sample_args["best_cat"] == '"aria"'