Add support for mache from fork/branch

Author: xylar

e3sm_supported_machines/bootstrap.py

@@ -11,7 +11,12 @@
 from mache.spack import make_spack_env, get_spack_script, \
     get_modules_env_vars_and_mpi_compilers
 from mache.permissions import update_permissions
-from shared import parse_args, check_call, install_miniconda, get_conda_base
+from shared import (
+    check_call,
+    get_conda_base,
+    install_miniforge3,
+    parse_args,
+)
 
 
 def get_config(config_file, machine):
@@ -147,16 +152,15 @@ def build_env(is_test, recreate, compiler, mpi, conda_mpi, version,
     packages = f'python={python} pip'
 
     source_activation_scripts = \
-        f'source {conda_base}/etc/profile.d/conda.sh && ' \
-        f'source {conda_base}/etc/profile.d/mamba.sh'
+        f'source {conda_base}/etc/profile.d/conda.sh'
 
     activate_env = f'{source_activation_scripts} && conda activate {env_name}'
 
     if not os.path.exists(env_path) or recreate:
         print(f'creating {env_name}')
         packages = f'{packages} "e3sm-unified={version}={mpi_prefix}_*"'
         commands = f'{activate_base} && ' \
-                   f'mamba create -y -n {env_name} {channels} {packages}'
+                   f'conda create -y -n {env_name} {channels} {packages}'
         check_call(commands)
 
         if conda_mpi == 'hpc':
@@ -174,6 +178,15 @@ def build_env(is_test, recreate, compiler, mpi, conda_mpi, version,
     return env_path, env_name, activate_env, channels, spack_env
 
 
+def install_mache_from_branch(activate_env, fork, branch):
+    print('Clone and install local mache\n')
+    commands = f'{activate_env} && ' \
+                f'cd build_mache/mache && ' \
+                f'python -m pip install --no-deps .'
+
+    check_call(commands)
+
+
 def build_sys_ilamb_esmpy(config, machine, compiler, mpi, template_path,
                           activate_env, channels, spack_base, spack_env):
 
@@ -378,6 +391,8 @@ def main():
 
     config = get_config(args.config_file, machine)
 
+    local_mache = args.mache_fork is not None and args.mache_branch is not None
+
     if args.release:
         is_test = False
     else:
@@ -387,13 +402,12 @@ def main():
     conda_base = os.path.abspath(conda_base)
 
     source_activation_scripts = \
-        f'source {conda_base}/etc/profile.d/conda.sh && ' \
-        f'source {conda_base}/etc/profile.d/mamba.sh'
+        f'source {conda_base}/etc/profile.d/conda.sh'
 
     activate_base = f'{source_activation_scripts} && conda activate'
 
     # install miniconda if needed
-    install_miniconda(conda_base, activate_base)
+    install_miniforge3(conda_base, activate_base)
 
     python, recreate, compiler, mpi, conda_mpi, activ_suffix, env_suffix, \
         activ_path = get_env_setup(args, config, machine)
@@ -405,11 +419,16 @@ def main():
     nompi_suffix = '_login'
     # first, make environment for login nodes.  We're using mpich from
     # conda-forge for now because we haven't had any luck with esmf>8.2.0 nompi
-    env_path, env_nompi, _, _, _ = build_env(
+    env_path, env_nompi, activate_env, _, _ = build_env(
         is_test, recreate, nompi_compiler, mpi, 'mpich', version,
         python, conda_base, nompi_suffix, nompi_suffix, activate_base,
         args.local_conda_build, config)
 
+    if local_mache:
+        install_mache_from_branch(activate_env=activate_env,
+                                  fork=args.mache_fork,
+                                  branch=args.mache_branch)
+
     if not is_test:
         # make a symlink to the environment
         link = os.path.join(conda_base, 'envs', 'e3sm_unified_latest')

e3sm_supported_machines/deploy_e3sm_unified.py

@@ -7,7 +7,12 @@
 
 from configparser import ConfigParser
 
-from shared import parse_args, check_call, install_miniconda, get_conda_base
+from shared import (
+    check_call,
+    get_conda_base,
+    install_miniforge3,
+    parse_args,
+)
 
 
 def get_config(config_file):
@@ -35,18 +40,17 @@ def bootstrap(activate_install_env, source_path, local_conda_build):
     sys.exit(0)
 
 
-def setup_install_env(activate_base, config, use_local):
+def setup_install_env(activate_base, use_local, mache):
     print('Setting up a conda environment for installing E3SM-Unified')
-    mache_version = config.get('e3sm_unified', 'mache')
     channels = []
     if use_local:
         channels.append('--use-local')
-    if 'rc' in mache_version:
+    if 'rc' in mache:
         channels.append('-c conda-forge/label/mache_dev')
     channels = ' '.join(channels)
     commands = f'{activate_base} && ' \
-               f'mamba create -y -n temp_e3sm_unified_install ' \
-               f'{channels} progressbar2 jinja2 mache={mache_version}'
+               f'conda create -y -n temp_e3sm_unified_install ' \
+               f'{channels} progressbar2 jinja2 {mache}'
 
     check_call(commands)
 
@@ -60,6 +64,21 @@ def remove_install_env(activate_base):
     check_call(commands)
 
 
+def install_mache_from_branch(activate_install_env, fork, branch):
+    print('Clone and install local mache\n')
+    commands = f'{activate_install_env} && ' \
+                f'rm -rf build_mache && ' \
+                f'mkdir -p build_mache && ' \
+                f'cd build_mache && ' \
+                f'git clone -b {branch} ' \
+                f'[email protected]:{fork}.git mache && ' \
+                f'cd mache && ' \
+                f'conda install -y --file spec-file.txt && ' \
+                f'python -m pip install --no-deps .'
+
+    check_call(commands)
+
+
 def main():
     args = parse_args(bootstrap=False)
     source_path = os.getcwd()
@@ -70,8 +89,14 @@ def main():
     conda_base = os.path.abspath(conda_base)
 
     source_activation_scripts = \
-        f'source {conda_base}/etc/profile.d/conda.sh && ' \
-        f'source {conda_base}/etc/profile.d/mamba.sh'
+        f'source {conda_base}/etc/profile.d/conda.sh'
+
+    local_mache = args.mache_fork is not None and args.mache_branch is not None
+    if local_mache:
+        mache = ''
+    else:
+        mache_version = config.get('e3sm_unified', 'mache')
+        mache = f'"mache={mache_version}"'
 
     activate_base = f'{source_activation_scripts} && conda activate'
 
@@ -80,9 +105,14 @@ def main():
         f'conda activate temp_e3sm_unified_install'
 
     # install miniconda if needed
-    install_miniconda(conda_base, activate_base)
+    install_miniforge3(conda_base, activate_base)
+
+    setup_install_env(activate_base, args.use_local, mache)
 
-    setup_install_env(activate_base, config, args.use_local)
+    if local_mache:
+        install_mache_from_branch(activate_install_env=activate_install_env,
+                                  fork=args.mache_fork,
+                                  branch=args.mache_branch)
 
     if args.release:
         is_test = False

e3sm_supported_machines/shared.py

@@ -45,9 +45,17 @@ def parse_args(bootstrap):
         parser.add_argument("--local_conda_build", dest="local_conda_build",
                             type=str,
                             help="A path for conda packages (for testing).")
+    parser.add_argument("--mache_fork", dest="mache_fork",
+                        help="Point to a mache fork (and branch) for testing")
+    parser.add_argument("--mache_branch", dest="mache_branch",
+                        help="Point to a mache branch (and fork) for testing")
 
     args = parser.parse_args(sys.argv[1:])
 
+    if (args.mache_fork is None) != (args.mache_branch is None):
+        raise ValueError('You must supply both or neither of '
+                         '--mache_fork and --mache_branch')
+
     return args
 
 
@@ -61,38 +69,35 @@ def check_call(commands, env=None):
         raise subprocess.CalledProcessError(proc.returncode, commands)
 
 
-def install_miniconda(conda_base, activate_base):
+def install_miniforge3(conda_base, activate_base):
     if not os.path.exists(conda_base):
-        print('Installing Miniconda3')
+        print('Installing Miniforge3')
         if platform.system() == 'Linux':
             system = 'Linux'
         elif platform.system() == 'Darwin':
             system = 'MacOSX'
         else:
             system = 'Linux'
-
-        miniconda = f'Mambaforge-{system}-x86_64.sh'
-        url = f'https://github.com/conda-forge/miniforge/releases/latest/download/{miniconda}'
+        miniforge = f'Miniforge3-{system}-x86_64.sh'
+        url = f'https://github.com/conda-forge/miniforge/releases/latest/download/{miniforge}'  # noqa: E501
         print(url)
         req = Request(url, headers={'User-Agent': 'Mozilla/5.0'})
         f = urlopen(req)
         html = f.read()
-        with open(miniconda, 'wb') as outfile:
+        with open(miniforge, 'wb') as outfile:
             outfile.write(html)
         f.close()
 
-        command = f'/bin/bash {miniconda} -b -p {conda_base}'
+        command = f'/bin/bash {miniforge} -b -p {conda_base}'
         check_call(command)
-        os.remove(miniconda)
+        os.remove(miniforge)
 
-    print('Doing initial setup')
+    print('Doing initial setup\n')
     commands = f'{activate_base} && ' \
                f'conda config --add channels conda-forge && ' \
                f'conda config --set channel_priority strict && ' \
-               f'mamba update -y --all && ' \
-               f'cp ~/.bashrc ~/.bashrc.conda_bak && ' \
-               f'mamba init && ' \
-               f'mv ~/.bashrc.conda_bak ~/.bashrc'
+               f'conda update -y --all && ' \
+               f'conda init --no-user'
 
     check_call(commands)