From fa676e317a234a37913baecd01b0973d65af556a Mon Sep 17 00:00:00 2001 From: Geoffrey Yu Date: Sun, 21 Apr 2024 13:25:37 -0400 Subject: [PATCH] [Envs] Implement remote environment installation (pyenv, python) --- errors/errors.yml | 4 + src/conductor/config.py | 8 ++ src/conductor/envs/install_maestro.py | 83 ++++++++++++++++ src/conductor/envs/resources/__init__.py | 0 src/conductor/envs/resources/install_pyenv.sh | 95 +++++++++++++++++++ src/conductor/errors/generated.py | 14 +++ 6 files changed, 204 insertions(+) create mode 100644 src/conductor/envs/install_maestro.py create mode 100644 src/conductor/envs/resources/__init__.py create mode 100644 src/conductor/envs/resources/install_pyenv.sh diff --git a/errors/errors.yml b/errors/errors.yml index b0b3130..20bade7 100644 --- a/errors/errors.yml +++ b/errors/errors.yml @@ -138,6 +138,10 @@ name: ConductorAbort message: "Conductor's execution has been aborted by the user." +3005: + name: MaestroInstallError + message: "An error occurred when installing Maestro in a remote environment: {error_message}" + # Archive and restore errors (error code 4xxx) 4001: diff --git a/src/conductor/config.py b/src/conductor/config.py index eede451..1d58663 100644 --- a/src/conductor/config.py +++ b/src/conductor/config.py @@ -65,3 +65,11 @@ # COND files can only include files with this extension. COND_INCLUDE_EXTENSION = ".cond" + +### +### Conductor Environment Configs +### + +MAESTRO_ROOT = ".conductor-maestro" + +MAESTRO_PYTHON_VERSION = "3.10.12" diff --git a/src/conductor/envs/install_maestro.py b/src/conductor/envs/install_maestro.py new file mode 100644 index 0000000..3316800 --- /dev/null +++ b/src/conductor/envs/install_maestro.py @@ -0,0 +1,83 @@ +import importlib.resources as pkg_resources + +from fabric.connection import Connection + +import conductor.envs.resources as env_resources +from conductor.config import MAESTRO_ROOT, MAESTRO_PYTHON_VERSION +from conductor.errors import MaestroInstallError + + +def install_maestro(c: Connection) -> None: + """ + Installs Maestro in the remote environment connected to by `c`. + """ + # General strategy: + # - Keep all Maestro-related code under a special root directory, specified + # using `MAESTRO_ROOT`. The assumption is that `MAESTRO_ROOT` is stored + # under `$HOME` (we assume a Unix-like remote environment). + # - Rely on `pyenv` to install a compatible version of Python in the remote + # environment. Note that `pyenv` depends on `git`. + # - Copy the entire Conductor module into the remote environment and install + # its dependencies in the pyenv. + # + # Design goals: + # - We should make this install as seamless as possible, similar to how easy + # it is to install VSCode server. + # - OK to assume a Unix-like remote environment. Should ideally make sure + # this works in a generic EC2 environment and inside a Docker container. + + try: + installer = _MaestroInstaller.create(c) + installer.ensure_root_dir() + installer.ensure_pyenv_installed() + installer.install_python() + except Exception as ex: + raise MaestroInstallError(error_msg=str(ex)) from ex + + +class _MaestroInstaller: + @classmethod + def create(cls, c: Connection) -> "_MaestroInstaller": + result = c.run("echo $HOME") + home_dir = result.stdout.strip() + return cls(c, home_dir) + + def __init__(self, c: Connection, home_dir: str) -> None: + self._c = c + # This is an absolute path in the remote environment. + self._maestro_root = f"{home_dir}/{MAESTRO_ROOT}" + self._pyenv_root = f"{self._maestro_root}/pyenv" + + def ensure_root_dir(self) -> None: + self._c.run(f"mkdir -p {self._maestro_root}") + + def ensure_pyenv_installed(self) -> None: + result = self._c.run(f"ls {self._pyenv_root}", warn=True, hide="both") + if result.ok: + # Assume install succeeded. + return + + try: + # Need to install. + install_script = pkg_resources.files(env_resources).joinpath( + "install_pyenv.sh" + ) + with pkg_resources.as_file(install_script) as path: + self._c.put(path, f"{self._maestro_root}/install_pyenv.sh") + # We want a custom install path to avoid interfering with the existing environment. + self._c.run( + f"bash {self._maestro_root}/install_pyenv.sh", + env={"PYENV_ROOT": self._pyenv_root}, + hide="both", + ) + except Exception: + # If the installation failed, remove the directory so that we will + # retry next time. + self._c.run(f"rm -r {self._pyenv_root}") + raise + + def install_python(self) -> None: + self._c.run( + f"{self._pyenv_root}/bin/pyenv install {MAESTRO_PYTHON_VERSION}", + env={"PYENV_ROOT": self._pyenv_root}, + ) diff --git a/src/conductor/envs/resources/__init__.py b/src/conductor/envs/resources/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/conductor/envs/resources/install_pyenv.sh b/src/conductor/envs/resources/install_pyenv.sh new file mode 100644 index 0000000..f6a0aa2 --- /dev/null +++ b/src/conductor/envs/resources/install_pyenv.sh @@ -0,0 +1,95 @@ +#!/usr/bin/env bash + +# Source: https://raw.githubusercontent.com/pyenv/pyenv-installer/master/bin/pyenv-installer +# Repository: https://github.com/pyenv/pyenv-installer +# Commit hash: 86a08ac9e38ec3a267e4b5c758891caf1233a2e4 + +set -e +[ -n "$PYENV_DEBUG" ] && set -x + +if [ -z "$PYENV_ROOT" ]; then + if [ -z "$HOME" ]; then + printf "$0: %s\n" \ + "Either \$PYENV_ROOT or \$HOME must be set to determine the install location." \ + >&2 + exit 1 + fi + export PYENV_ROOT="${HOME}/.pyenv" +fi + +colorize() { + if [ -t 1 ]; then printf "\e[%sm%s\e[m" "$1" "$2" + else echo -n "$2" + fi +} + +# Checks for `.pyenv` file, and suggests to remove it for installing +if [ -d "${PYENV_ROOT}" ]; then + { echo + colorize 1 "WARNING" + echo ": Can not proceed with installation. Kindly remove the '${PYENV_ROOT}' directory first." + echo + } >&2 + exit 1 +fi + +failed_checkout() { + echo "Failed to git clone $1" + exit -1 +} + +checkout() { + [ -d "$2" ] || git -c advice.detachedHead=0 clone --branch "$3" --depth 1 "$1" "$2" || failed_checkout "$1" +} + +if ! command -v git 1>/dev/null 2>&1; then + echo "pyenv: Git is not installed, can't continue." >&2 + exit 1 +fi + +# Check ssh authentication if USE_SSH is present +if [ -n "${USE_SSH}" ]; then + if ! command -v ssh 1>/dev/null 2>&1; then + echo "pyenv: configuration USE_SSH found but ssh is not installed, can't continue." >&2 + exit 1 + fi + + # `ssh -T git@github.com' returns 1 on success + # See https://docs.github.com/en/authentication/connecting-to-github-with-ssh/testing-your-ssh-connection + ssh -T git@github.com 1>/dev/null 2>&1 || EXIT_CODE=$? + if [[ ${EXIT_CODE} != 1 ]]; then + echo "pyenv: github ssh authentication failed." + echo + echo "In order to use the ssh connection option, you need to have an ssh key set up." + echo "Please generate an ssh key by using ssh-keygen, or follow the instructions at the following URL for more information:" + echo + echo "> https://docs.github.com/en/repositories/creating-and-managing-repositories/troubleshooting-cloning-errors#check-your-ssh-access" + echo + echo "Once you have an ssh key set up, try running the command again." + exit 1 + fi +fi + +if [ -n "${USE_SSH}" ]; then + GITHUB="git@github.com:" +else + GITHUB="https://github.com/" +fi + +checkout "${GITHUB}pyenv/pyenv.git" "${PYENV_ROOT}" "${PYENV_GIT_TAG:-master}" +checkout "${GITHUB}pyenv/pyenv-doctor.git" "${PYENV_ROOT}/plugins/pyenv-doctor" "master" +checkout "${GITHUB}pyenv/pyenv-update.git" "${PYENV_ROOT}/plugins/pyenv-update" "master" +checkout "${GITHUB}pyenv/pyenv-virtualenv.git" "${PYENV_ROOT}/plugins/pyenv-virtualenv" "master" + +if ! command -v pyenv 1>/dev/null; then + { echo + colorize 1 "WARNING" + echo ": seems you still have not added 'pyenv' to the load path." + echo + } >&2 + + { # Without args, `init` commands print installation help + "${PYENV_ROOT}/bin/pyenv" init || true + "${PYENV_ROOT}/bin/pyenv" virtualenv-init || true + } >&2 +fi diff --git a/src/conductor/errors/generated.py b/src/conductor/errors/generated.py index 8eb8363..a780384 100644 --- a/src/conductor/errors/generated.py +++ b/src/conductor/errors/generated.py @@ -386,6 +386,19 @@ def _message(self): ) +class MaestroInstallError(ConductorError): + error_code = 3005 + + def __init__(self, **kwargs): + super().__init__() + self.error_message = kwargs["error_message"] + + def _message(self): + return "An error occurred when installing Maestro in a remote environment: {error_message}".format( + error_message=self.error_message, + ) + + class OutputFileExists(ConductorError): error_code = 4001 @@ -636,6 +649,7 @@ def _message(self): "TaskFailed", "OutputDirTaken", "ConductorAbort", + "MaestroInstallError", "OutputFileExists", "OutputPathDoesNotExist", "NoTaskOutputsToArchive",