def init( repo_path: Optional[str], ) -> None: """Initialize ZenML on given path. Args: repo_path: Path to the repository. Raises: InvalidGitRepositoryError: If repo is not a git repo. AssertionError """ if sys.version_info.minor == 6: warning( "ZenML support for Python 3.6 will be deprecated soon. Please " "consider upgrading your Python version to ensure ZenML works " "properly in the future." ) if repo_path is None: repo_path = os.getcwd() declare(f"Initializing at {repo_path}") try: Repository.init_repo(repo_path=repo_path) declare(f"ZenML repo initialized at {repo_path}") except git.InvalidGitRepositoryError: # type: ignore[attr-defined] error( f"{repo_path} is not a valid git repository! Please " f"initialize ZenML within a git repository using " f"`git init `" ) except AssertionError as e: error(f"{e}")
def monitor_orchestrator(orchestrator_name: Optional[str] = None) -> None: """Monitor a running orchestrator.""" orchestrator_, orchestrator_name = _get_orchestrator(orchestrator_name) if not orchestrator_.is_running: cli_utils.warning(f"Can't monitor orchestrator '{orchestrator_name}' " f"because it isn't running.") return if not orchestrator_.log_file: cli_utils.warning(f"Can't monitor orchestrator '{orchestrator_name}'.") return cli_utils.declare( f"Monitoring orchestrator '{orchestrator_name}', press CTRL+C to stop." ) try: with open(orchestrator_.log_file, "r") as log_file: # seek to the end of the file log_file.seek(0, 2) while True: line = log_file.readline() if not line: time.sleep(0.1) continue line = line.rstrip("\n") click.echo(line) except KeyboardInterrupt: cli_utils.declare( f"Stopped monitoring orchestrator '{orchestrator_name}'.")
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_logging_verbosity(verbosity: str) -> None: """Set logging level""" # TODO [ENG-150]: Implement this. verbosity = verbosity.upper() if verbosity not in LoggingLevels.__members__: raise KeyError( f"Verbosity must be one of {list(LoggingLevels.__members__.keys())}" ) cli_utils.declare(f"Set verbosity to: {verbosity}")
def up_orchestrator(orchestrator_name: Optional[str] = None) -> None: """Provisions resources for the orchestrator""" orchestrator_, orchestrator_name = _get_orchestrator(orchestrator_name) cli_utils.declare( f"Bootstrapping resources for orchestrator: `{orchestrator_name}`. " f"This might take a few seconds...") orchestrator_.up() cli_utils.declare(f"Orchestrator: `{orchestrator_name}` is up.")
def down_orchestrator(orchestrator_name: Optional[str] = None) -> None: """Tears down resources for the orchestrator""" orchestrator_, orchestrator_name = _get_orchestrator(orchestrator_name) cli_utils.declare( f"Tearing down resources for orchestrator: `{orchestrator_name}`.") orchestrator_.down() cli_utils.declare( f"Orchestrator: `{orchestrator_name}` resources are now torn down.")
def version() -> None: """Version of ZenML""" declare(r""" .-') _ ('-. .-') _ _ .-') ( OO) )_( OO) ( OO ) )( '.( OO )_ ,(_)----.(,------.,--./ ,--,' ,--. ,--.),--. | | | .---'| \ | |\ | `.' | | |.-') '--. / | | | \| | ) | | | | OO ) (_/ / (| '--. | . |/ | |'.'| | | |`-' | / /___ | .--' | |\ | | | | |(| '---.' | || `---.| | \ | | | | | | | `--------'`------'`--' `--' `--' `--' `------' """) click.echo(click.style(f"version: {__version__}", bold=True))
def list(git_examples_handler: GitExamplesHandler) -> None: """List all available examples.""" declare("Listing examples: \n") for example in git_examples_handler.examples: declare(f"{example.name}") declare("\nTo pull the examples, type: ") declare("zenml example pull EXAMPLE_NAME")
def register_stack( stack_name: str, metadata_store: str, artifact_store: str, orchestrator: str, ) -> None: """Register a stack.""" service = Repository().get_service() stack = BaseStack( artifact_store_name=artifact_store, orchestrator_name=orchestrator, metadata_store_name=metadata_store, ) service.register_stack(stack_name, stack) cli_utils.declare(f"Stack `{stack_name}` successfully registered!")
def clean(git_examples_handler: GitExamplesHandler) -> None: """Deletes the ZenML examples directory from your current working directory.""" examples_directory = os.path.join(os.getcwd(), "zenml_examples") if (fileio.file_exists(examples_directory) and fileio.is_dir(examples_directory) and confirmation( "Do you wish to delete the ZenML examples directory? \n" f"{examples_directory}")): git_examples_handler.clean_current_examples() declare( "ZenML examples directory was deleted from your current working " "directory.") elif not fileio.file_exists(examples_directory) and not fileio.is_dir( examples_directory): logger.error(f"Unable to delete the ZenML examples directory - " f"{examples_directory} - " "as it was not found in your current working directory.")
def pull( git_examples_handler: GitExamplesHandler, example_name: str, force: bool, version: str, ) -> None: """Pull examples straight into your current working directory. Add the flag --force or -f to redownload all the examples afresh. Use the flag --version or -v and the version number to specify which version of ZenML you wish to use for the examples.""" git_examples_handler.pull(force=force, version=version) destination_dir = os.path.join(os.getcwd(), "zenml_examples") fileio.create_dir_if_not_exists(destination_dir) examples = (git_examples_handler.examples if not example_name else [ Example( example_name, Path( os.path.join( git_examples_handler.examples_repo.examples_dir, example_name, )), ) ]) for example in examples: if not fileio.file_exists(str(example.path)): error( f"Example {example.name} does not exist! Available examples: " f"{[e.name for e in git_examples_handler.examples]}") return example_destination_dir = os.path.join(destination_dir, example.name) if fileio.file_exists(example_destination_dir): if confirmation(f"Example {example.name} is already pulled. " f"Do you wish to overwrite the directory?"): fileio.rm_dir(example_destination_dir) else: warning(f"Example {example.name} not overwritten.") continue declare(f"Pulling example {example.name}...") git_examples_handler.copy_example(example, example_destination_dir) declare(f"Example pulled in directory: {example_destination_dir}")
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 _get_orchestrator( orchestrator_name: Optional[str] = None, ) -> Tuple["BaseOrchestrator", str]: """Gets an orchestrator for a given name. Args: orchestrator_name: Name of the orchestrator to get. If `None`, the orchestrator of the active stack gets returned. Returns: A tuple containing the orchestrator and its name. Raises: DoesNotExistException: If no orchestrator for the name exists. """ if not orchestrator_name: active_stack = Repository().get_active_stack() orchestrator_name = active_stack.orchestrator_name cli_utils.declare( f"No orchestrator name given, using `{orchestrator_name}` " f"from active stack.") service = Repository().get_service() return service.get_orchestrator(orchestrator_name), orchestrator_name
def get_active_artifact_store() -> None: """Gets the artifact store of the active stack.""" artifact_store_name = Repository().get_active_stack().artifact_store_name cli_utils.declare(f"Active artifact store: {artifact_store_name}")
def delete_stack(stack_name: str) -> None: """Delete a stack.""" service = Repository().get_service() cli_utils.declare(f"Deleting stack: {stack_name}") service.delete_stack(stack_name) cli_utils.declare("Deleted!")
def set_active_stack(stack_name: str) -> None: """Sets a stack active.""" repo = Repository() repo.set_active_stack(stack_name) cli_utils.declare(f"Active stack: {stack_name}")
def get_active_stack() -> None: """Gets the active stack.""" repo = Repository() key = repo.get_active_stack_key() cli_utils.declare(f"Active stack: {key}")
def get_active_orchestrator() -> None: """Gets the orchestrator of the active stack.""" orchestrator_name = Repository().get_active_stack().orchestrator_name cli_utils.declare(f"Active orchestrator: {orchestrator_name}")
def opt_in() -> None: """Opt-in to analytics""" gc = GlobalConfig() gc.analytics_opt_in = True gc.update() cli_utils.declare("Opted in to analytics.")
def is_analytics_opted_in() -> None: """Check whether user is opt-in or opt-out of analytics.""" gc = GlobalConfig() cli_utils.declare(f"Analytics opt-in: {gc.analytics_opt_in}")
def opt_out() -> None: """Opt-out to analytics""" gc = GlobalConfig() gc.analytics_opt_in = False gc.update() cli_utils.declare("Opted out of analytics.")
def delete_orchestrator(orchestrator_name: str) -> None: """Delete an orchestrator.""" service = Repository().get_service() service.delete_orchestrator(orchestrator_name) cli_utils.declare(f"Deleted orchestrator: `{orchestrator_name}`")
def delete_metadata_store(metadata_store_name: str) -> None: """Delete a metadata store.""" service = Repository().get_service() service.delete_metadata_store(metadata_store_name) cli_utils.declare(f"Deleted metadata store: {metadata_store_name}")
def delete_artifact_store(artifact_store_name: str) -> None: """Delete a artifact store.""" service = Repository().get_service() service.delete_artifact_store(artifact_store_name) cli_utils.declare(f"Deleted artifact store: {artifact_store_name}")
def get_active_metadata_store() -> None: """Gets the metadata store of the active stack.""" metadata_store_name = Repository().get_active_stack().metadata_store_name cli_utils.declare(f"Active metadata store: {metadata_store_name}")