[Envs] Implement remote environment installation (pyenv, python)

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:

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"

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},
+        )

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 [email protected]' returns 1 on success
+  # See https://docs.github.com/en/authentication/connecting-to-github-with-ssh/testing-your-ssh-connection
+  ssh -T [email protected] 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="[email protected]:"
+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

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",