PaddleX/paddlex/repo_manager/core.py

# !/usr/bin/env python3
# -*- coding: UTF-8 -*-
################################################################################
#
# Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
#
################################################################################
"""
Author: PaddlePaddle Authors
"""

import os
import sys
from collections import OrderedDict

from ..utils import logging
from .utils import install_deps_using_pip
from .meta import get_all_repo_names
from .repo import build_repo_instance, build_repo_group_cloner, build_repo_group_installer

__all__ = [
    'set_parent_dirs', 'setup', 'wheel', 'is_initialized', 'initialize',
    'get_versions'
]


class _GlobalContext(object):
    REPO_PARENT_DIR = None
    PDX_COLLECTION_MOD = None
    REPOS = None

    @classmethod
    def set_parent_dirs(cls, repo_parent_dir, pdx_collection_mod):
        """ set_parent_dirs """
        cls.REPO_PARENT_DIR = repo_parent_dir
        cls.PDX_COLLECTION_MOD = pdx_collection_mod

    @classmethod
    def build_repo_instance(cls, repo_name):
        """ build_repo_instance """
        return build_repo_instance(repo_name, cls.REPO_PARENT_DIR,
                                   cls.PDX_COLLECTION_MOD)

    @classmethod
    def is_initialized(cls):
        """ is_initialized """
        return cls.REPOS is not None

    @classmethod
    def initialize(cls):
        """ initialize """
        cls.REPOS = []

    @classmethod
    def add_repo(cls, repo):
        """ add_repo """
        if not cls.is_initialized():
            cls.initialize()
        cls.REPOS.append(repo)

    @classmethod
    def add_repos(cls, repos):
        """ add_repos """
        if len(repos) == 0 and not cls.is_initialized():
            cls.initialize()
        for repo in repos:
            cls.add_repo(repo)


set_parent_dirs = _GlobalContext.set_parent_dirs
is_initialized = _GlobalContext.is_initialized


def setup(repo_names,
          reinstall=None,
          no_deps=False,
          constraints=None,
          platform=None,
          update_repos=False):
    """ setup """
    repos = []
    for repo_name in repo_names:
        repo = _GlobalContext.build_repo_instance(repo_name)
        repos.append(repo)

    repos_to_clone = []
    for repo in repos:
        repo_name = repo.name
        if repo.check_repo_exiting():
            logging.warning(f"Existing of {repo.name} repo.")
            if reinstall is None:
                if sys.stdin.isatty():
                    logging.warning("Should we remove it (y/n)?")
                try:
                    remove_existing = input()
                except EOFError:
                    logging.warning(
                        "Unable to read from stdin. Please set `reinstall` to \
                        True or False to apply a global setting for reclone repos."
                    )
                    raise
                remove_existing = remove_existing.lower() in ('y', 'yes')
            else:
                remove_existing = reinstall
            if remove_existing:
                repo.remove()
                repos_to_clone.append(repo)
            else:
                logging.warning(
                    f"We will use the existing repo of {repo.name}.")
        else:
            repos_to_clone.append(repo)
    repos_to_install = []
    for repo in repos:
        repo_name = repo.name
        if repo.check_installation():
            logging.warning(f"Existing installation of {repo.name} detected.")
            if reinstall is None and not update_repos:
                if sys.stdin.isatty():
                    logging.warning("Should we uninstall it (y/n)?")
                try:
                    uninstall_existing = input()
                except EOFError:
                    logging.warning(
                        "Unable to read from stdin. Please set `reinstall` to \
                        True or False to apply a global setting for reinstalling repos."
                    )
                    raise
                uninstall_existing = uninstall_existing.lower() in ('y', 'yes')
            else:
                if reinstall or update_repos:
                    uninstall_existing = True
            if uninstall_existing:
                repo.uninstall()
                repos_to_install.append(repo)
            else:
                logging.warning(
                    f"We will use the existing installation of {repo.name}.")
        else:
            repos_to_install.append(repo)
    cloner = build_repo_group_cloner(*repos_to_clone)
    installer = build_repo_group_installer(*repos_to_install)

    logging.info("Now cloning the repos...")
    cloner.clone(force_reclone=False, platform=platform)
    logging.info("All repos are existing.")

    if not no_deps:
        logging.info("Dependencies are listed below:")
        logging.info(installer.get_deps())

    logging.info("Now installing the packages...")
    install_deps_using_pip()
    if update_repos:
        installer.update()
        logging.info("All repos are updated.")
    installer.install(
        force_reinstall=False, no_deps=no_deps, constraints=constraints)
    logging.info("All packages are installed.")


def wheel(repo_names, dst_dir='./', fail_fast=False):
    """ wheel """
    for repo_name in repo_names:
        repo = _GlobalContext.build_repo_instance(repo_name)
        logging.info(f"Now building Wheel for {repo_name}...")
        try:
            tgt_dir = os.path.join(dst_dir, repo.pkg_name)
            if os.path.exists(tgt_dir):
                raise FileExistsError(f"{tgt_dir} already exists.")
            repo.wheel(tgt_dir)
        except Exception as e:
            logging.warning(
                f"Failed to build wheel for {repo_name}. We encountered the following error:\n  {str(e)}\n"
            )
            if fail_fast:
                raise
        else:
            logging.info(f"Wheel for {repo_name} is built.\n")


def initialize(repo_names=None):
    """ initialize """
    if _GlobalContext.is_initialized():
        raise RuntimeError(
            "PDX has already been initialized. Reinitialization is not supported."
        )

    if repo_names is None:
        try_all = True
        repo_names = get_all_repo_names()
    else:
        try_all = False

    repos = []
    for repo_name in repo_names:
        logging.debug(f"Now initializing {repo_name}...")
        repo = _GlobalContext.build_repo_instance(repo_name)
        flag = repo.initialize()
        if flag:
            logging.debug(f"{repo_name} is initialized.")
            repos.append(repo)
        else:
            if try_all:
                logging.debug(
                    f"Failed to initialize {repo_name}. Please make sure {repo_name} is properly installed."
                )
            else:
                pass

    _GlobalContext.add_repos(repos)


def get_versions(repo_names=None):
    """ get_versions """
    if repo_names is None:
        repo_names = get_all_repo_names()

    name2versions = OrderedDict()
    for repo_name in repo_names:
        repo = _GlobalContext.build_repo_instance(repo_name)
        versions = repo.get_version()
        name2versions[repo_name] = versions
    return name2versions