#!/usr/bin/env python3
import os
import subprocess
import time
import datetime
import argparse
import sys
from pathlib import Path
about_tool = """
Build toolchains using Spack.\n
This function builds toolchains for MPSD-HPC at the appropriate directory, \n
for given system architecture and MPSD software stack version.\n
The toolchains are built using the bash script spack_setup.sh, and the results are logged.
"""
config_vars = {
"cmd_log_file": "logs/install-software-environment.log",
"build_log_file": f"build_toolchains_mpsd_spack_ver_{time.strftime('%Y%m%d-%H%M%S')}.log", # TODO: modify toolchains,mpsd_spack_ver when the variable is available
"spack_environments_repo": "https://gitlab.gwdg.de/mpsd-cs/spack-environments.git",
}
shared_var = {
"script_branch": subprocess.run(
["git", "rev-parse", "--abbrev-ref", "HEAD"], stdout=subprocess.PIPE
)
.stdout.decode()
.strip(),
"script_commit_hash": subprocess.run(
["git", "rev-parse", "--short", "HEAD"], stdout=subprocess.PIPE
)
.stdout.decode()
.strip(),
"spe_branch": None,
"spe_commit_hash": None,
"available_toolchains": None,
"toolchain_base_dir": None,
}
# Helper class to change directory via context manager
class os_chdir:
def __init__(self, new_dir):
self.new_dir = new_dir
self.saved_dir = os.getcwd()
def __enter__(self):
os.chdir(self.new_dir)
def __exit__(self, exc_type, exc_val, exc_tb):
os.chdir(self.saved_dir)
def setup_log_cmd(msg=None,*args, **kwargs):
# Create log directory if it doesn't exist
if not os.path.exists("logs"):
os.makedirs("logs")
# Write to the log file with the following format
# --------------------------------------------------
# 2023-02-29T23:32:01, install-software-environment.py --release dev-23a --install ALL
# Software environment installer branch: script_branch (commit hash: script_commit_hash)
# Spack environments branch: dev-23a (commit hash: spe_commit_hash)
# MSGs
with open(config_vars["cmd_log_file"], "a") as f:
if msg:
# Write the message to the log file
f.write(msg + "\n")
else:
# Write the header
f.write("-" * 50 + "\n")
# Gather data to log
# call statement:
cmd_line = " ".join(sys.argv)
# script branch and commit hash
script_branch= subprocess.run(["git", "rev-parse", "--abbrev-ref", "HEAD"], stdout=subprocess.PIPE).stdout.decode().strip()
script_commit_hash = subprocess.run(["git", "rev-parse", "--short", "HEAD"], stdout=subprocess.PIPE ).stdout.decode().strip()
# spack-environments branch and commit hash from kwargs
spe_branch = kwargs.get("spe_branch", None)
spe_commit_hash = kwargs.get("spe_commit_hash", None)
# Write to log file
f.write(f"{datetime.datetime.now().isoformat()}, {cmd_line}\n")
f.write(
f"Software environment installer branch: {script_branch} (commit hash: {script_commit_hash})\n"
)
f.write(
f"Spack environments branch: {spe_branch} (commit hash: {spe_commit_hash})\n"
)
def create_dir_structure(mpsd_release, script_dir):
# Create the directory structure for the release
release_base_dir = os.path.join(script_dir, mpsd_release)
if not os.path.exists(release_base_dir):
os.makedirs(release_base_dir)
else:
# Warn that the target directory already exists.
print(">Target directory already exists. Continuing...")
with os_chdir(release_base_dir):
# Clone the spack-environments repo if it doesn't exist
if not os.path.exists("spack-environments"):
subprocess.run(
[
"git",
"clone",
config_vars["spack_environments_repo"],
]
)
with os_chdir("spack-environments"):
# Git fetch and checkout the release branch and git pull to be sure that the resulting repo is up to date
subprocess.run(["git", "fetch", "--all"])
subprocess.run(["git", "checkout", mpsd_release])
subprocess.run(["git", "pull"])
def get_release_info(mpsd_release, script_dir):
# Get the info for release
release_base_dir = os.path.join(script_dir, mpsd_release)
if not os.path.exists(release_base_dir):
raise Exception("Release directory does not exist. Run create_dir_structure() first.")
with os_chdir(release_base_dir):
with os_chdir("spack-environments"):
# Get the branch and commit hash of the spack-environments repo and store them in shared_var
spe_commit_hash = (
subprocess.run(["git", "rev-parse", "HEAD"], stdout=subprocess.PIPE)
.stdout.decode()
.strip()
)
spe_branch = (
subprocess.run(
["git", "rev-parse", "--abbrev-ref", "HEAD"], stdout=subprocess.PIPE
)
.stdout.decode()
.strip()
)
available_toolchains = os.listdir("toolchains")
return spe_branch,spe_commit_hash,available_toolchains
def prepare_environment(mpsd_release, script_dir):
# create the release folder and clone the spack-environments repo
# get the branch and commit hash of the spack-environments repo and available toolchains
# log the command usage.
create_dir_structure(mpsd_release, script_dir)
spe_branch, spe_commit_hash, available_toolchains = get_release_info(mpsd_release, script_dir)
setup_log_cmd(spe_branch=spe_branch, spe_commit_hash=spe_commit_hash)
return available_toolchains
def install_environment(
release, toolchains, target_dir, force_reinstall, skip_build_cache
):
print(f"Installing release {release} with toolchains {toolchains} to {target_dir}")
shared_var = prepare_environment(target_dir, release, force_reinstall, shared_var)
def remove_environment(release, toolchains, target_dir):
print(f"Removing release {release} with toolchains {toolchains} from {target_dir}")
def start_new_environment(release, from_release, target_dir):
print(f"Starting new release {release} from {from_release} to {target_dir}")
def main():
parser = argparse.ArgumentParser(description=about_tool)
subparsers = parser.add_subparsers(
dest="action", title="actions", description="valid actions", required=True
)
subparsers.required = True
list_of_cmds = [
("prepare", "Prepare the environment for installation on the disk"),
("install", "Install a software environment"),
("reinstall", "Reinstall a software environment"),
("remove", "Remove a software environment or toolchains from an environment"),
("start-new", "Start a new software environment version"),
]
for cmd, help_text in list_of_cmds:
subp = subparsers.add_parser(cmd, help=help_text)
if cmd == "start-new":
subp.add_argument(
"--from-release",
dest="from_release",
type=str,
required=True,
help="Release version to start from",
)
subp.add_argument(
"--to-release",
dest="to_release",
type=str,
required=True,
help="Release version to create",
)
else:
subp.add_argument(
"release", type=str, help="Release version to install or remove"
)
if cmd in ["install", "reinstall"]:
subp.add_argument(
"--toolchains",
type=str,
nargs="*",
default="ALL",
help="List of toolchains to install (use ALL to install all toolchains)",
)
subp.add_argument(
"--enable-build-cache",
action="store_true",
help="Enable Spack build cache. Useful for reinstallation but consumes time and disk space",
)
# Carry out the action
args = parser.parse_args()
# target dir is the place where this script exists. the release `dev` in script_dir/dev-23a
target_dir = os.path.dirname(os.path.realpath(__file__))
# Check the command and run related function
if args.action == "remove":
remove_environment(args.release, args.toolchains, target_dir)
elif args.action == "start-new":
start_new_environment(args.from_release, args.to_release, target_dir)
elif args.action == "install":
install_environment(
args.release, args.toolchains, target_dir, False, args.skip_build_cache
)
elif args.action == "prepare":
prepare_environment(args.release, target_dir)
if __name__ == "__main__":
main()