diff --git a/.github/workflows/artifacts.yml b/.github/workflows/artifacts.yml index 76350cf..3a19044 100644 --- a/.github/workflows/artifacts.yml +++ b/.github/workflows/artifacts.yml @@ -11,7 +11,7 @@ jobs: steps: - name: setup-environment run: | - zypper install -y libmount-devel gcc make rpm-build git sudo python3 + zypper install -y libmount-devel gcc make rpm-build git sudo python3 wget - name: clone uses: actions/checkout@v2 - name: build-srpm diff --git a/.gitignore b/.gitignore index 218fb7b..f7268a9 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,5 @@ test/tmp # site generated by mkdocs site + +__pycache__ diff --git a/VERSION b/VERSION index 13d683c..0c89fc9 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -3.0.1 \ No newline at end of file +4.0.0 \ No newline at end of file diff --git a/activate b/activate index 377c68b..cf03c25 100644 --- a/activate +++ b/activate @@ -1,8 +1,10 @@ # only run if bash or zsh -[ -n "${BASH_VERSION:-}" -o -n "${ZSH_VERSION:-}" ] || return 0 +[ -n "${BASH_VERSION:-}" ] || return 0 export UENV_CMD=@@impl@@ +export UENV_IMG_CMD=@@image_impl@@ export UENV_VERSION=@@version@@ +export UENV_PREFIX=@@prefix@@ function uenv { local _last_exitcode=$? @@ -12,23 +14,46 @@ function uenv { echo "" echo "Usage: uenv [--version] [--help] []" echo "" - echo "the following commands are available" + echo "The following commands are available:" echo " run run a command in an environment" echo " start start a new shell with an environment loaded" echo " stop stop a shell with an environment loaded" echo " status print information about each running environment" echo " modules view module status and activate with --use" echo " view activate a view" + echo " image query and pull uenv images" + echo "" + echo "Type 'uenv command --help' for more information and examples for a specific command, e.g." + echo " uenv start --help" + echo " uenv image find --help" } if [ "$1" = "--version" ]; then echo "@@version@@"; - elif [[ $# -eq 0 || "$1" = "--help" || "$1" = "-h" ]]; then - uenv_usage; - elif [[ " $* " =~ [[:space:]](-h|--help)[[:space:]] ]]; then - echo "$($UENV_CMD "$@")" else - eval "$($UENV_CMD "$@")" + flags="" + while [[ $# -gt 0 ]]; do + case $1 in + -?*|--*) + flags="$flags $1" + shift + ;; + *) + break + ;; + esac + done + + if [[ $# -eq 0 ]]; then + uenv_usage; + elif [ "$1" = "image" ]; then + shift; + $UENV_IMG_CMD $flags "$@" + elif [[ " $* " =~ [[:space:]](-h|--help)[[:space:]] ]]; then + echo "$($UENV_CMD $flags "$@")" + else + eval "$($UENV_CMD $flags "$@")" + fi fi unset -f uenv_usage diff --git a/docs/images.md b/docs/images.md new file mode 100644 index 0000000..458cbf9 --- /dev/null +++ b/docs/images.md @@ -0,0 +1,24 @@ + +the `image` +``` +# list images available locally +uenv image pull gromacs/2023 +# pull directly from: requires credentials +uenv image pull --full --name test/2023 + +# list images available locally +uenv image find +# list images available on JFrog +uenv image avail +uenv image avail --cluster=eiger +``` + +updates to other commands +``` +# this should return information about the current cluster? +> uenv status +no uenv loaded +on cluster eiger +``` + + diff --git a/install b/install index 8a6a146..b776901 100755 --- a/install +++ b/install @@ -5,6 +5,20 @@ set -euo pipefail script_path=$(dirname "$(readlink -f "$BASH_SOURCE")") version=$(cat "$script_path/VERSION") +# default installation destination, only required for building rpm packages +destdir="" + +get_oras_arch() { + case "$(uname -m)" in + aarch64) + echo "arm64" + ;; + *) + echo "amd64" + ;; + esac +} + function usage { echo "uenv installer" echo "Usage: install [--prefix=] [--yes] [--help]" @@ -117,29 +131,57 @@ run () { fi } -echo "installing uenv version $version in $prefix" +echo "installing uenv version $version in $prefix from $script_path" echo "local install: $local_install" +libexec=${prefix}/libexec +impl_path="${libexec}/uenv-impl" +img_impl_path="${libexec}/uenv-image" +lib_path="${libexec}/lib" if [ "$local_install" == "no" ] then init_path="${destdir}/etc/profile.d/uenv.sh" - impl_path="${prefix}/libexec/uenv-impl" - impl_path_install="${destdir}/${impl_path}" - else init_path="${prefix}/bin/activate-uenv" - impl_path="${prefix}/libexec/uenv-impl" - impl_path_install="${impl_path}" fi +echo "installing $lib_path" +run install -v -m755 -d "$lib_path" +run install -v -m644 -D -t "${destdir}/${lib_path}" ${script_path}/lib/alps.py +run install -v -m644 -D -t "${destdir}/${lib_path}" ${script_path}/lib/datastore.py +run install -v -m644 -D -t "${destdir}/${lib_path}" ${script_path}/lib/flock.py +run install -v -m644 -D -t "${destdir}/${lib_path}" ${script_path}/lib/jfrog.py +run install -v -m644 -D -t "${destdir}/${lib_path}" ${script_path}/lib/names.py +run install -v -m644 -D -t "${destdir}/${lib_path}" ${script_path}/lib/oras.py +run install -v -m644 -D -t "${destdir}/${lib_path}" ${script_path}/lib/record.py +run install -v -m644 -D -t "${destdir}/${lib_path}" ${script_path}/lib/terminal.py + +echo "installing ${img_impl_path}" +run install -v -m755 -D "$script_path/uenv-image" "${destdir}/${img_impl_path}" + +## ALL GOOD BELOW HERE + echo "installing $impl_path" -run install -D -m755 <(sed "s|@@version@@|$version|g" $script_path/uenv-impl) $impl_path_install +run install -D -m755 <(sed "s|@@version@@|$version|g" $script_path/uenv-impl) ${destdir}/$impl_path + +# install the oras client +echo "installing oras client" +oras_arch=$(get_oras_arch) +oras_version=1.1.0 +oras_file=oras_${oras_version}_linux_${oras_arch}.tar.gz +oras_url=https://github.com/oras-project/oras/releases/download/v${oras_version}/${oras_file} +oras_path=`mktemp -d` +(cd "$oras_path"; run wget --quiet "$oras_url"; run tar -xzf "$oras_file"; run rm *.tar.gz) +oras_exe_path="${libexec}/uenv-oras" +run install -v -m755 -D "${oras_path}/oras" "${destdir}/${oras_exe_path}" echo "installing $init_path" -run install -D -m644 <( +run install -v -m644 -D <( sed \ - -e "s|@@impl@@|$impl_path|g" \ + -e "s|@@impl@@|$impl_path|g" \ -e "s|@@version@@|$version|g" \ + -e "s|@@prefix@@|$prefix|g" \ + -e "s|@@image_impl@@|$img_impl_path|g" \ $script_path/activate) \ $init_path diff --git a/lib/alps.py b/lib/alps.py new file mode 100644 index 0000000..99fa28b --- /dev/null +++ b/lib/alps.py @@ -0,0 +1,22 @@ +# arbitrary alps-specific gubbins + +import os + +# the path used to store a users cached images and meta data +def uenv_repo_path(path: str=None) -> str: + if path is not None: + return path + + # check whether the image path has been explicitly set: + path = os.environ.get("UENV_REPO_PATH") + if path is not None: + return path + + # if not, try to use the path $SCRATCH/.uenv-images/, if SCRATCH exists + path = os.environ.get("SCRATCH") + if path is not None: + return path + "/.uenv-images" + + terminal.error("No repository path available: set UENV_REPO_PATH or use the --repo flag") + + diff --git a/lib/datastore.py b/lib/datastore.py new file mode 100644 index 0000000..a44d358 --- /dev/null +++ b/lib/datastore.py @@ -0,0 +1,159 @@ +import json +import os + +from record import Record +import terminal +import names + +UENV_CLI_API_VERSION=1 + +class DataStore: + + def __init__(self): + # all images store with (key,value) = (sha256,Record) + self._images = {} + self._short_sha = {} + + self._store = {"system": {}, "uarch": {}, "name": {}, "version": {}, "tag": {}} + + def add_record(self, r: Record): + sha = r.sha256 + short_sha = r.sha256[:16] + self._images.setdefault(sha, []).append(r) + # check for (exceedingly unlikely) collision + if short_sha in self._short_sha: + old_sha = self._short_sha[short_sha] + if sha != old_sha: + terminal.error('image hash collision:\n {sha}\n {old_sha}') + self._short_sha[sha[:16]] = sha + self._store["system"] .setdefault(r.system, []).append(sha) + self._store["uarch"] .setdefault(r.uarch, []).append(sha) + self._store["name"] .setdefault(r.name, []).append(sha) + self._store["version"].setdefault(r.version, []).append(sha) + self._store["tag"] .setdefault(r.tag, []).append(sha) + + def find_records(self, **constraints): + if not constraints: + raise ValueError("At least one constraint must be provided") + + for field in constraints: + if (field != "sha") and (field not in self._store): + raise ValueError(f"Invalid field: {field}. Must be one of 'system', 'uarch', 'name', 'version', 'tag', 'sha'") + + if "sha" in constraints: + sha = constraints["sha"] + matching_records_sets = [set()] + if len(sha)<64: + if sha in self._short_sha: + matching_records_sets = [set([self._short_sha[sha]])] + else: + if sha in self._images: + matching_records_sets = [set([sha])] + else: + # Find matching records for each constraint + matching_records_sets = [ + set(self._store[field].get(value, [])) for field, value in constraints.items() + ] + + # Intersect all sets of matching records + if matching_records_sets: + unique = set.intersection(*matching_records_sets) + else: + unique = set() + + results = [] + for sha in unique: + results += (self._images[sha]) + results.sort(reverse=True) + return results + + @property + def images(self): + return self._images + + # return a list of records that match a sha + def get_record(self, sha: str) -> Record: + if names.is_full_sha256(sha): + return self._images.get(sha, []) + elif names.is_short_sha256(sha): + return self._images.get(self._short_sha[sha], []) + raise ValueError(f"{sha} is not a valid sha256 or short (16 character) sha") + + # Convert to a dictionary that can be written to file as JSON + # The serialisation and deserialisation are central: able to represent + # uenv that are available in both JFrog and filesystem directory tree. + def serialise(self, version: int=UENV_CLI_API_VERSION): + image_list = [] + for x in self._images.values(): + image_list += x + terminal.info(f"serialized image list in datastore: {image_list}") + return { + "API_VERSION": version, + "images": [img.dictionary for img in image_list] + } + + # Convert to a dictionary that can be written to file as JSON + # The serialisation and deserialisation are central: able to represent + # uenv that are available in both JFrog and filesystem directory tree. + @classmethod + def deserialise(cls, datastore): + result = cls() + for img in datastore["images"]: + result.add_record(Record.from_dictionary(img)) + +class FileSystemCache(): + def __init__(self, path: str): + self._path = path + self._index = path + "/index.json" + + if not os.path.exists(self._index): + # error: cache does not exists + raise FileNotFoundError(f"filesystem cache not found {self._path}") + + with open(self._index, "r") as fid: + raw = json.loads(fid.read()) + self._database = DataStore() + for img in raw["images"]: + self._database.add_record(Record.fromjson(img)) + + @staticmethod + def create(path: str, exists_ok: bool=False): + if not os.path.exists(path): + terminal.info(f"FileSyStemCache: creating path {path}") + os.makedirs(path) + index_file = f"{path}/index.json" + if not os.path.exists(index_file): + terminal.info(f"FileSyStemCache: creating empty index {index_file}") + empty_config = { "API_VERSION": UENV_CLI_API_VERSION, "images": [] } + with open(index_file, "w") as f: + # default serialisation is str to serialise the pathlib.PosixPath + f.write(json.dumps(empty_config, sort_keys=True, indent=2, default=str)) + f.write("\n") + + terminal.info(f"FileSyStemCache: available {index_file}") + + @property + def database(self): + return self._database + + def add_record(self, record: Record): + self._database.add_record(record) + + # The path where an image would be stored + # will return a path even for images that are not stored + def image_path(self, r: Record) -> str: + return self._path + "/images/" + r.sha256 + + # Return the full record for a given hash + # Returns None if no image with that hash is stored in the repo. + def get_record(self, sha256: str): + if not names.is_valid_sha(sha256): + raise ValueError(f"{sha256} is not a valid image sha256 (neither full 64 or short 16 character form)") + return self._database.get_record(sha256) + + def publish(self): + with open(self._index, "w") as f: + # default serialisation is str to serialise the pathlib.PosixPath + f.write(json.dumps(self._database.serialise(), sort_keys=True, indent=2, default=str)) + f.write("\n") + diff --git a/lib/flock.py b/lib/flock.py new file mode 100644 index 0000000..7827999 --- /dev/null +++ b/lib/flock.py @@ -0,0 +1,36 @@ +import fcntl +import time + +import terminal + +class Lock(): + READ = 1 + WRITE = 2 + def __init__(self, path: str, type: int): + self._lockfile = f"{path}.lock" + + # open the file + self._lock = open(self._lockfile, "a") + + self._time = time.time() + + # acquire lock + self._type = type + if self._type==Lock.READ: + # acquire shared lock + fcntl.flock(self._lock, fcntl.LOCK_SH) + else: + # acquire exclusive lock + fcntl.flock(self._lock, fcntl.LOCK_EX) + + terminal.info(f"aquired lock {self._lockfile} at {self._time}") + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, exc_traceback): + fcntl.flock(self._lock, fcntl.LOCK_UN) # Release the lock + self._lock.close() + endtime = time.time() + terminal.info(f"released lock {self._lockfile} at {endtime}, held for {(endtime - self._time)*1000:.2f} ms") + diff --git a/lib/jfrog.py b/lib/jfrog.py new file mode 100644 index 0000000..a755a0a --- /dev/null +++ b/lib/jfrog.py @@ -0,0 +1,74 @@ +import requests + +from datastore import DataStore +from record import Record +import terminal + + +# The https://cicd-ext-mw.cscs.ch/uenv/list API endpoint returns +# a list of images in the jfrog uenv. +# +#{ +# "results": +# [ +# { +# "repo" : "uenv", +# "path" : "build/clariden/zen3/prgenv-gnu/23.11/1094139948", +# "name" : "manifest.json", +# "created" : "2023-12-04T09:05:44.034Z", +# "size" : "123683707", +# "sha256" : "134c04d01bb3583726804a094b144d3637997877ef6162d1fe19eabff3c72c3a", +# "stats" : [{ +# "downloaded" : "2023-12-11T17:56:59.052Z", +# "downloads" : 11 +# }] +# }, +# ... +# ], +# "range" : +# { +# "start_pos" : 0, +# "end_pos" : 22, +# "total" : 22 +# } +#} +# + +def query() -> tuple: + try: + # GET request to the middleware + url = "https://cicd-ext-mw.cscs.ch/uenv/list" + terminal.info(f"querying jfrog at {url}") + response = requests.get(url) + response.raise_for_status() + + raw_records = response.json() + + deploy_database = DataStore() + build_database = DataStore() + + for record in raw_records["results"]: + path = record["path"] + + date = Record.to_datetime(record["created"]) + sha256 = record["sha256"] + size = record["size"] + if path.startswith("build/"): + r = Record.frompath(path[len("build/"):], date, size, sha256) + build_database.add_record(r) + if path.startswith("deploy/"): + r = Record.frompath(path[len("deploy/"):], date, size, sha256) + deploy_database.add_record(r) + + return (deploy_database, build_database) + + except Exception as error: + raise RuntimeError(f"downloading image data from jfrog.svs.cscs.ch ({str(error)})") + +def relative_from_record(record): + return f"{record.system}/{record.uarch}/{record.name}/{record.version}:{record.tag}" + +def address(record: Record, namespace: str): + if namespace!='deploy' and namespace!='build': + raise RuntimeError("namespace must be one of build or deploy") + return f"jfrog.svc.cscs.ch/uenv/{namespace}/{relative_from_record(record)}" diff --git a/lib/names.py b/lib/names.py new file mode 100644 index 0000000..03a46ba --- /dev/null +++ b/lib/names.py @@ -0,0 +1,77 @@ +import re + +def is_full_sha256(s: str): + pattern = re.compile(r'^[a-fA-F-0-9]{64}$') + return True if pattern.match(s) else False + +def is_short_sha256(s: str): + pattern = re.compile(r'^[a-fA-F-0-9]{16}$') + return True if pattern.match(s) else False + +def is_valid_sha(sha:str) -> bool: + if is_full_sha256(sha): + return True + if is_short_sha256(sha): + return True + return False + +# return dictionary {"name", "version", "tag", "sha"} from a uenv description string +# "prgenv_gnu" -> ("prgenv_gnu", None, None, None) +# "prgenv_gnu/23.11" -> ("prgenv_gnu", "23.11", None, None) +# "prgenv_gnu/23.11:latest" -> ("prgenv_gnu", "23.11", "latest", None) +# "3313739553fe6553" -> (None, None, None, "3313739553fe6553") +def parse_uenv_string(uenv: str) -> dict: + name = version = tag = sha = None + + if is_valid_sha(uenv): + sha = uenv + else: + # todo: assert no more than 1 '/' + # todo: assert no more than 1 ':' + # todo: assert that '/' is before ':' + splits = uenv.split("/",1) + name = splits[0] + if len(splits)>1: + splits = splits[1].split(":",1) + version = splits[0] + if len(splits)>1: + tag = splits[1] + + return {"name": name, "version": version, "tag": tag, "sha": sha} + +def is_complete_description(uenv: dict) -> bool: + if uenv["sha"]: + return True + + if uenv["name"] and uenv["version"] and uenv["tag"]: + return True + + return False + +class IncompleteUenvName(Exception): + """Exception raised for errors related to invalid uenv names.""" + + def __init__(self, message): + self.message = message + super().__init__(self.message) + + def __str__(self): + return f'IncompleteUenvName: {self.message}' + + +def create_filter(uenv, require_complete=False) -> dict: + if uenv is None: + if require_complete: + raise IncompleteUenvName(f"{uenv}") + return {} + + components = parse_uenv_string(uenv) + if require_complete and not is_complete_description(components): + raise IncompleteUenvName(f"{uenv}") + + img_filter = {} + for key, value in components.items(): + if value is not None: + img_filter[key] = value + + return img_filter diff --git a/lib/oras.py b/lib/oras.py new file mode 100644 index 0000000..04fac70 --- /dev/null +++ b/lib/oras.py @@ -0,0 +1,47 @@ +import pathlib +import shutil +import subprocess + +import terminal + +def find_oras() -> str: + # - the oras executable is installed in the libexec path + # - this file is installed in the libexec/lib path + # - the executable is named uenv-oras + # search here for the executable + oras_dir = pathlib.Path(__file__).parent.parent.resolve() + oras_file = oras_dir / "uenv-oras" + terminal.info(f"searching for oras executable: {oras_file}") + if oras_file.is_file(): + terminal.info(f"using {oras_file}") + return oras_file.as_posix() + + # fall back to finding the oras executable + terminal.info(f"{oras_file} does not exist - searching for oras") + oras_file = shutil.which("oras") + if not oras_file: + terminal.error(f"no oras executable found") + terminal.info(f"using {oras_file}") + return oras_file + + +def run_command(args): + try: + command = [find_oras()] + args + + terminal.info(f"calling oras: {' '.join(command)}") + result = subprocess.run( + command, + stdout=subprocess.PIPE, # Capture standard output + stderr=subprocess.PIPE, # Capture standard error + check=True, # Raise exception if command fails + encoding='utf-8' # Decode output from bytes to string + ) + + # Print standard output + terminal.info("Output:\n{result.stdout}") + + except subprocess.CalledProcessError as e: + # Print error message along with captured standard error + terminal.error("An error occurred:\n", e.stderr) + diff --git a/lib/record.py b/lib/record.py new file mode 100644 index 0000000..8833a19 --- /dev/null +++ b/lib/record.py @@ -0,0 +1,127 @@ +from datetime import datetime, timezone + +class Record: + @staticmethod + def to_datetime(date: str): + # In Python 3.6, datetime.fromisoformat is not available. + # Manually parsing the string. + dt_format = '%Y-%m-%dT%H:%M:%S.%fZ' + return datetime.strptime(date, dt_format).replace(tzinfo=timezone.utc) + + + def __init__(self, system: str, uarch: str, name: str, version: str, tag: str, date: str, size_bytes: int, sha256: str): + self._system = system + self._uarch = uarch + self._name = name + self._version = version + self._tag = tag + self._date = date + self._bytes = size_bytes + self._sha256 = sha256 + + # build/eiger/zen2/cp2k/2023/1133706947 + @classmethod + def frompath(cls, path: str, date: str, size_bytes: int, sha256: str): + fields = path.split("/") + if len(fields) != 5: + raise ValueError("Record must have exactly 5 fields") + + system, uarch, name, version, tag = fields + return cls(system, uarch, name, version, tag, date, size_bytes, sha256) + + @classmethod + def fromjson(cls, raw: dict): + system = raw["system"] + uarch = raw["uarch"] + name = raw["name"] + version = raw["version"] + tag = raw["tag"] + date = Record.to_datetime(raw["date"]) + size_bytes = raw["size"] + sha256 = raw["sha256"] + + return cls(system, uarch, name, version, tag, date, size_bytes, sha256) + + def __eq__(self, other): + if not isinstance(other, Record): + return False + return self.sha256==other.sha256 + + def __lt__(self, other): + if self.system < other.system: return True + if other.system < self.system: return False + if self.uarch < other.uarch: return True + if other.uarch < self.uarch: return False + if self.name < other.name: return True + if other.name < self.name: return False + if self.version < other.version: return True + if other.version< self.version: return False + if self.tag=="latest" and other.tag!="latest": return False + if other.tag=="latest" and self.tag!="latest": return True + if self.tag < other.tag: return True + #if other.tag < self.tag: return False + return False + + def __str__(self): + return f"{self.system}/{self.uarch}/{self.name}/{self.version}:{self.tag}/{self.sha256}" + + def __repr__(self): + return f"Record({self.system}, {self.uarch}, {self.name}, {self.version}, {self.tag})" + + @property + def system(self): + return self._system + + @property + def uarch(self): + return self._uarch + + @property + def name(self): + return self._name + + @property + def date(self): + return self._date + + @property + def version(self): + return self._version + + @property + def tag(self): + return self._tag + + @tag.setter + def tag(self, newtag): + self._tag = newtag + + @property + def sha256(self): + return self._sha256 + + @property + def size(self): + return self._bytes + + @property + def datestring(self): + return self.date.strftime('%Y-%m-%dT%H:%M:%S.%f')[:-3] + 'Z' + + @property + def path(self): + return f"{self.system}/{self.uarch}/{self.name}{self.version}/{self.tag}" + + @property + def dictionary(self): + return { + "system": self.system, + "uarch": self.uarch, + "name": self.name, + "date": self.datestring, + "version": self.version, + "tag": self.tag, + "sha256": self.sha256, + "size": self.size + } + diff --git a/lib/terminal.py b/lib/terminal.py new file mode 100644 index 0000000..ff4620c --- /dev/null +++ b/lib/terminal.py @@ -0,0 +1,70 @@ +import os +import sys + +colored_output = True +debug_level = 1 + +# Choose whether to use colored output. +# - by default colored output is ON +# - if the flag --no-color is passed it is OFF +# - if the environment variable NO_COLOR is set it is OFF +def use_colored_output(cli_arg: bool): + global colored_output + + # The --no-color argument overrides all environment variables if passed. + if cli_arg: + colored_output = False + return + + # Check the env. var NO_COLOR and disable color if set. + # See https://no-color.org/ for the informal spec + # + # Command-line software which adds ANSI color to its output by default + # should check for a NO_COLOR environment variable that, when present and + # not an empty string (regardless of its value), prevents the addition of + # ANSI color. + # + if os.environ.get('NO_COLOR') is not None: + color_var = os.environ.get('NO_COLOR') + if len(color_var)>0: + colored_output = False + return + + colored_output = True + +def colorize(string, color): + colors = { + "red": "31", + "green": "32", + "yellow": "33", + "blue": "34", + "magenta": "35", + "cyan": "36", + "white": "37", + "gray": "90", + } + if colored_output: + return f"\033[1;{colors[color]}m{string}\033[0m" + else: + return string + +def set_debug_level(level: int): + global debug_level + debug_level = level + +def error(message, abort=True): + print(f"{colorize('[error]', 'red')} {message}", file=sys.stderr) + if abort: + exit(1) + +def warning(message): + print(f"{colorize('[warning]', 'yellow')} {message}", file=sys.stderr) + +def info(message): + if debug_level>1: + print(f"{colorize('[info]', 'green')} {message}", file=sys.stderr) + +def exit_with_success(): + exit(0) + + diff --git a/rpm/generate-rpm.sh b/rpm/generate-rpm.sh index 280296a..f7f1c7b 100755 --- a/rpm/generate-rpm.sh +++ b/rpm/generate-rpm.sh @@ -27,10 +27,12 @@ mkdir -p "${build_path}" tar_path="${build_path}/${pkg_name}" mkdir -p "${tar_path}" +cp -r "${source_path}/lib" "${tar_path}" cp "${source_path}/install" "${tar_path}" cp "${source_path}/VERSION" "${tar_path}" cp "${source_path}/activate" "${tar_path}" cp "${source_path}/uenv-impl" "${tar_path}" +cp "${source_path}/uenv-image" "${tar_path}" tar_file="${build_path}/SOURCES/${pkg_name}.tar.gz" tar -czf "${tar_file}" --directory "${build_path}" "${pkg_name}" diff --git a/uenv-image b/uenv-image new file mode 100755 index 0000000..b193d53 --- /dev/null +++ b/uenv-image @@ -0,0 +1,510 @@ +#!/usr/bin/python3 + +# --- A note about the shebang --- +# It is hard-coded to /usr/bin/python3 instead of "/usr/bin/env python3" so that the +# same python3 is always used, instead of using a version of python3 that might be +# loaded into the environment. + +import argparse +import copy +import os +import pathlib +import re +import sys +import textwrap + +prefix = pathlib.Path(__file__).parent.resolve() +libpath = prefix / 'lib' +sys.path = [libpath.as_posix()] + sys.path + +import alps +import datastore +import flock +import jfrog +import names +import oras +import record +import terminal +from terminal import colorize + +def make_argparser(): + parser = argparse.ArgumentParser( + prog="uenv image", + formatter_class=argparse.RawDescriptionHelpFormatter, + description=textwrap.dedent( +f"""\ +Manage and query uenv images. + +For more information on how to use individual commands use the --help flag. + +{colorize("Example", "blue")}: get help on the find command + {colorize("uenv image find --help", "white")} +""" + )) + parser.add_argument("--no-color", + action="store_true", + help="Disable color output. By default color output is enabled, unless the NO_COLOR environment variable is set.") + parser.add_argument("-v", "--verbose", + action="store_true", + help="Enable verbose output for debugging.") + parser.add_argument("-r", "--repo", + required=False, default=None, type=str, + help="The path on the local filesystem where uenv are managed. By default the environment variable UENV_REPO_PATH is used, if set, otherwise $SCRATCH/.uenv-images is used if the SCRATCH environment variable is set. This option will override these defaults, and must be set if neither of the defaults is set.") + + subparsers = parser.add_subparsers(dest="command") + + path_parser = subparsers.add_parser("inspect", + formatter_class=argparse.RawDescriptionHelpFormatter, + help="Display detailed information about a uenv.", + epilog=f"""\ +Display detailed information about a uenv. +It is an error if no uenv matching the requested spec is on the local filesystem. + +{colorize("Example", "blue")} - get information about a uenv + {colorize("uenv image inspect prgenv-gnu/24.2:latest", "white")} + +{colorize("Example", "blue")} - it is also possible to get the path of the most relevant version/tag +of a uenv, or use an explicit sha: + {colorize("uenv image inspect prgenv-gnu", "white")} + {colorize("uenv image inspect prgenv-gnu/24.2", "white")} + {colorize("uenv image inspect 3313739553fe6553f789a35325eb6954a37a7b85cdeab943d0878a05edaac998", "white")} + {colorize("uenv image inspect 3313739553fe6553 # the first 16 characters can be used", "white")} + +{colorize("Example", "blue")} - print the path of a uenv + {colorize("uenv image inspect --format '{path}' prgenv-gnu", "white")} + +{colorize("Example", "blue")} - print the name and tag of a uenv un the name:tag format + {colorize("uenv image inspect --format '{name}:{tag}' prgenv-gnu", "white")} + +{colorize("Example", "blue")} - print the location of the squashfs file of an image + {colorize("uenv image inspect --format '{sqfs}' prgenv-gnu", "white")} + +Including name, tag and sqfs, the following variables can be printed in a +format string passed to the --format option: + name: name + version: version + tag: tag + system: the system that the uenv was built for + uarch: the micro-architecture that the uenv was built for + sha256: the unique sha256 hash of the uenv + date: date that the uenv was created + path: absolute path where the uenv is stored + sqfs: absolute path of the squashfs file\ +""") + path_parser.add_argument("-s", "--system", required=False, type=str) + path_parser.add_argument("-a", "--uarch", required=False, type=str) + path_parser.add_argument("--format", type=str, + default= +f"""\ +name: {{name}} +version: {{version}} +tag: {{tag}} +system: {{system}} +uarch: {{uarch}} +sha256: {{sha256}} +date: {{date}} +path: {{path}} +sqfs: {{sqfs}}\ +""" , + help="optional format string") + path_parser.add_argument("uenv", type=str) + + find_parser = subparsers.add_parser("find", + help="Find uenv in the CSCS registry.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=f"""\ +Find uenv in the CSCS registry. + +Uenv can be downloaded from the registry using the pull command, and list the +downloaded uenv with the list command. For more information: + {colorize("uenv image pull --help", "white")} + {colorize("uenv image ls --help", "white")} + +{colorize("Example", "blue")} - find all uenv available (deployed) on this cluster: + {colorize("uenv image find", "white")} + +{colorize("Example", "blue")} - find all uenv with the name prgenv-gnu on this cluster: + {colorize("uenv image find prgenv-gnu", "white")} + +{colorize("Example", "blue")} - find all uenv with name prgenv-gnu and version 24.2 on this cluster: + {colorize("uenv image find prgenv-gnu/24.2", "white")} + +{colorize("Example", "blue")} - find the uenv with name prgenv-gnu, version 24.2 and tag "latest" on this cluster: + {colorize("uenv image find prgenv-gnu/24.2:latest", "white")} + +{colorize("Example", "blue")} - find all uenv with the name prgenv-gnu for uarch target gh200 on this cluster: + {colorize("uenv image find prgenv-gnu --uarch=gh200", "white")} + +{colorize("Example", "blue")} - find all uenv that have a concrete sha256 checksum on this cluster: + {colorize("uenv image find 3313739553fe6553f789a35325eb6954a37a7b85cdeab943d0878a05edaac998", "white")} + {colorize("uenv image find 3313739553fe6553 # the first 16 characters can be used", "white")} + +{colorize("Example", "blue")} - find all uenv that have been generated as build artifacts on this cluster: + {colorize("uenv image find --build", "white")} +""") + find_parser.add_argument("-s", "--system", required=False, type=str) + find_parser.add_argument("-a", "--uarch", required=False, type=str) + find_parser.add_argument("--build", action="store_true", + help="Search undeployed builds.", required=False) + find_parser.add_argument("uenv", nargs="?", default=None, type=str) + + pull_parser = subparsers.add_parser("pull", + help="Pull a uenv from the CSCS registry.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=f"""\ +Pull a uenv from the CSCS registry. + +{colorize("Example", "blue")} - pull the most recent version of uenv with the name prgenv-gnu on this cluster: + {colorize("uenv image pull prgenv-gnu", "white")} + +{colorize("Example", "blue")} - pull the most recent tag of uenv with the name prgenv-gnu/24.2 on this cluster: + {colorize("uenv image pull prgenv-gnu/24.2", "white")} + +{colorize("Example", "blue")} - pull a specific version and tag of prgenv-gnu: + {colorize("uenv image pull prgenv-gnu/24.2:latest", "white")} + +{colorize("Example", "blue")} - pull using the unique sha256 of an uenv. + {colorize("uenv image pull 3313739553fe6553f789a35325eb6954a37a7b85cdeab943d0878a05edaac998", "white")} + {colorize("uenv image pull 3313739553fe6553 # the first 16 characters can be used", "white")} + +{colorize("Example", "blue")} - pull the uenv with the name prgenv-gnu for uarch target gh200 on this cluster: +Note that this is only neccesary when a vCluster has nodes with more than one uarch, and +versions of the same uenv compiled against different uarch has been deployed. + {colorize("uenv image pull prgenv-gnu --uarch=gh200", "white")} + +{colorize("Example", "blue")} - pull a uenv from the build repository. +By default only deployed images are pulled, and this option is only available to users +with appropriate JFrog access and with the JFrog token in their oras keychain. + {colorize("uenv image pull 3313739553fe6553 --build", "white")} +""") + pull_parser.add_argument("-s", "--system", required=False, type=str) + pull_parser.add_argument("-a", "--uarch", required=False, type=str) + pull_parser.add_argument("--build", action="store_true", required=False, + help="enable undeployed builds") + pull_parser.add_argument("uenv", nargs="?", default=None, type=str) + + list_parser = subparsers.add_parser("ls", + help="List available uenv.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=f"""\ +List uenv that are available. + + + +{colorize("Example", "blue")} - list all uenv: + {colorize("uenv image ls", "white")} + +{colorize("Example", "blue")} - list all uenv with the name prgenv-gnu: + {colorize("uenv image ls prgenv-gnu", "white")} + +{colorize("Example", "blue")} - list all uenv with name prgenv-gnu and version 24.2: + {colorize("uenv image ls prgenv-gnu/24.2", "white")} + +{colorize("Example", "blue")} - list the uenv with name prgenv-gnu, version 24.2 and tag "latest": + {colorize("uenv image ls prgenv-gnu/24.2:latest", "white")} + +{colorize("Example", "blue")} - list all uenv with the name prgenv-gnu for uarch target gh200: + {colorize("uenv image ls prgenv-gnu --uarch=gh200", "white")} + +{colorize("Example", "blue")} - list all uenv that have a concrete sha256 checksum: + {colorize("uenv image ls 3313739553fe6553f789a35325eb6954a37a7b85cdeab943d0878a05edaac998", "white")} + {colorize("uenv image ls 3313739553fe6553", "white")} # the first 16 characters can be used +""") + list_parser.add_argument("-s", "--system", required=False, type=str) + list_parser.add_argument("-a", "--uarch", required=False, type=str) + list_parser.add_argument("uenv", nargs="?", default=None, type=str) + + create_parser = subparsers.add_parser("create", + help="Create a local file system repository.") + create_parser.add_argument("--exists-ok", action="store_true", required=False, + help="No error if the local registry exists.") + + deploy_parser = subparsers.add_parser("deploy", + help="Deploy a uenv to the 'deploy' namespace, accessible to all users.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog= +f"""\ +Deploy a uenv from the build repo to the deploy repo, where all users can access it. + +{colorize("Attention", "red")}: this operation can only be performed by CSCS staff with appropriate +access rights and a JFrog token configured for oras. + +{colorize("Attention", "red")}: deployment a new version of a uenv with the :latest tag will modify the version that is +provided by default to all users on the system. + +The recommended method for deploying images is to first find the sha256 of the image +in the build namespace that you wish to deploy + +{colorize("uenv image find --build icon-wcp", "white")} +{colorize("uenv/version:tag uarch date sha256 size", "gray")} +{colorize("icon-wcp/v1:1208323951 gh200 2024-03-11 e901b85c7652802e 7.8 GB", "gray")} +{colorize("icon-wcp/v1:1206570249 gh200 2024-03-08 0cb8e7e23b4fdb5c 7.8 GB", "gray")} + +In this example, we choose the most recent build with sha256 e901b85c7652802e + +{colorize("uenv image deploy e901b85c7652802e --tags=latest,v1", "white")} + +Then check that it has been deployed: + +{colorize("uenv image find e901b85c7652802e", "white")} +{colorize("uenv/version:tag uarch date sha256 size", "gray")} +{colorize("icon-wcp/v1:latest gh200 2024-03-11 e901b85c7652802e 7.8 GB", "gray")} +{colorize("icon-wcp/v1:v1 gh200 2024-03-11 e901b85c7652802e 7.8 GB", "gray")} +""" + ) + deploy_parser.add_argument("--tags", required=False, default="latest", help="Comma separated list of tags.", type=str) + deploy_parser.add_argument("source", nargs=1, type=str, metavar="SOURCE", + help="The full name/version:tag or sha256 of the uenv to deploy.") + + return parser + +def get_options(args): + options = {} + if args.system is None: + sys_name = os.getenv("CLUSTER_NAME") + if sys_name is None: + raise ValueError("No system name was provided, and the CLUSTER_NAME environment variable is not set.") + options["system"] = sys_name + else: + options["system"] = args.system + + options["name"] = args.uenv + options["uarch"] = args.uarch + + return options + +def get_filter(args): + options = get_options(args) + + name = options["name"] + if name is None: + terminal.info(f"get_filter: no search term provided") + img_filter = {} + else: + terminal.info(f"get_filter: parsing {name}") + img_filter = names.create_filter(name, require_complete=False) + + img_filter["system"] = options["system"] + + if options["uarch"] is not None: + img_filter["uarch"] = options["uarch"] + + terminal.info(f"get_filter: filter {img_filter}") + return img_filter + +def inspect_string(record: record.Record, cache: datastore.FileSystemCache, format_string: str) -> str: + try: + return format_string.format( + system = record.system, + uarch = record.uarch, + name = record.name, + date = record.date, + version = record.version, + tag = record.tag, + sha256 = record.sha256, + path = cache.image_path(record), + sqfs = cache.image_path(record) + "/store.squashfs", + ) + except Exception as err: + terminal.error(f"unable to format {str(err)}") + +# pretty print a list of Record +def print_records(records): + if len(records)>0: + print(terminal.colorize(f"{'uenv/version:tag':40}{'uarch':6}{'date':10} {'sha256':16} {'size':<10}", "yellow")) + for r in records: + namestr = f"{r.name}/{r.version}" + tagstr = f"{r.tag}" + label = namestr + ":" + tagstr + datestr = r.date.strftime("%Y-%m-%d") + S = r.size + if S<1024: + size_str = f"{S:<} bytes" + elif S<1024*1024: + size_str = f"{(S/1024):<.0f} kB" + elif S<1024*1024*1024: + size_str = f"{(S/(1024*1024)):<.0f} MB" + else: + size_str = f"{(S/(1024*1024*1024)):<.1f} GB" + print(f"{label:<40}{r.uarch:6}{datestr:10} {r.sha256[:16]:16} {size_str:<10}") + +if __name__ == "__main__": + + parser = make_argparser() + args = parser.parse_args() + + if args.command is None: + parser.print_help() + sys.exit(0) + + terminal.use_colored_output(args.no_color) + if args.verbose: + terminal.set_debug_level(2) + + terminal.info(f"command mode: {args.command}") + + repo_path = alps.uenv_repo_path(args.repo) + + terminal.info(f"local repository: {repo_path}") + + if args.command in ["find", "pull"]: + img_filter = get_filter(args) + terminal.info(f"using {'build' if args.build else 'deploy'} remote repo") + + try: + deploy, build = jfrog.query() + except RuntimeError as err: + terminal.error(f"{str(err)}") + + terminal.info(f"downloaded jfrog meta data: build->{len(build.images)}, deploy->{len(deploy.images)}") + + remote_database = build if args.build else deploy + + records = remote_database.find_records(**img_filter) + + terminal.info(f"The following records matched the query: {records}") + + if args.command == "find": + if len(records)>0: + print_records(records) + else: + print("no images match the query") + + elif args.command == "pull": + # verify that there is at least one image that matches the query + if len(records)==0: + terminal.error(f"no images match the query {args.uenv}") + + # check that there is only one matching uenv + if len(set([r.name for r in records]))>1: + print_records(records) + print() + terminal.error(f"ambiguous uenv {args.uenv}") + + # check that there is only one matching uenv + if len(set([r.uarch for r in records]))>1: + print_records(records) + print() + terminal.error( + "more than one uarch matches the the requested uenv. " + "Specify the desired uarch with the --uarch flag") + + t = records[0] + source_address = jfrog.address(t, 'build' if args.build else 'deploy') + + terminal.info(f"pulling {t} from {source_address} {t.size/(1024*1024):.0f} MB") + + terminal.info(f"repo path: {repo_path}") + + with flock.Lock(f"{repo_path}/index.json", flock.Lock.WRITE) as lk: + cache = datastore.FileSystemCache(repo_path) + + image_path = cache.image_path(t) + + # if the record isn't already in the filesystem repo download it + if not cache.get_record(t.sha256): + terminal.info(f"downloading {t.sha256}") + # download the image using oras + oras.run_command(["pull", "-o", image_path, source_address]) + # add the record to the cache + terminal.info(f"updating file system cache") + cache.add_record(t) + # publish the updated index + terminal.info(f"publishing file system cache") + cache.publish() + else: + terminal.info(f"image {t.sha256} is already in the cache") + terminal.info(f"image downloaded at {image_path}/store.squashfs") + + sys.exit(0) + + elif args.command == "ls": + terminal.info(f"repo path: {repo_path}") + + img_filter = get_filter(args) + + with flock.Lock(f"{repo_path}/index.json", flock.Lock.READ) as lk: + fscache = datastore.FileSystemCache(repo_path) + + records = fscache.database.find_records(**img_filter) + print_records(records) + + sys.exit(0) + + elif args.command == "inspect": + terminal.info(f"repo path: {repo_path}") + + img_filter = get_filter(args) + + with flock.Lock(f"{repo_path}/index.json", flock.Lock.READ) as lk: + fscache = datastore.FileSystemCache(repo_path) + + records = fscache.database.find_records(**img_filter) + + if len(records)==0: + terminal.error(f"no uenv matches the spec: {args.uenv}") + sys.exit(1) + + path = fscache.image_path(records[0]) + formatted_output = inspect_string(records[0], fscache, args.format) + print(formatted_output) + + sys.exit(0) + + elif args.command == "create": + terminal.info(f"repo path: {repo_path}") + + try: + datastore.FileSystemCache.create(repo_path, exists_ok=args.exists_ok) + except Exception as err: + terminal.error(f"unable to find or initialise the local registry: {str(err)}") + + sys.exit(0) + + elif args.command == "deploy": + source = args.source[0] + terminal.info(f"request to deploy '{source}' with tags '{args.tags}'") + + # for deployment, we require a complete description, i.e + # name/version:tag OR sha256 + try: + img_filter = names.create_filter(source, require_complete=True) + except names.IncompleteUenvName as err: + terminal.error(f"source {source} must have be of the form name/version:tag or sha256.") + + # query JFrog for the list of images + try: + _, build_database = jfrog.query() + except RuntimeError as err: + terminal.error(f"{str(err)}") + + terminal.info(f"downloaded jfrog build meta data: {len(build_database.images)} images") + + + # expect that src has [name, version, tag] keys + records = build_database.find_records(**img_filter) + + if not (len(records)==1): + terminal.error(f"source {source} is not an image in the build repository") + + source_record = records[0] + terminal.info(f"the source is {source_record}") + + target_record = copy.deepcopy(source_record) + + # create comma separated list of tags to be attached to the deployed image + tags = [ tag.strip() for tag in args.tags.split(',') ] + target_record.tag = ','.join(tags) + + terminal.info(f"source: {source_record}") + source_address = jfrog.address(source_record, 'build') + target_address = jfrog.address(target_record, 'deploy') + terminal.info(f"source address: {source_address}") + terminal.info(f"target address: {target_address}") + + oras.run_command(["cp", "--concurrency", "10", "--recursive", source_address, target_address]) + + terminal.info(f"successfully deployed {target_address}") + + sys.exit(0) + diff --git a/uenv-impl b/uenv-impl index fa7aa55..9221cb3 100755 --- a/uenv-impl +++ b/uenv-impl @@ -11,6 +11,19 @@ import os import pathlib import sys import subprocess +import textwrap + +prefix = pathlib.Path(__file__).parent.resolve() +libpath = prefix / 'lib' +sys.path = [libpath.as_posix()] + sys.path + +import alps +import datastore +import flock +import names +import record +import terminal +from terminal import colorize VERSION="@@version@@" @@ -18,26 +31,61 @@ shell_noop=" :" shell_error="local _exitcode=1" def make_argparser(): - parser = argparse.ArgumentParser() - parser.add_argument("--no-color", action="store_true") - parser.add_argument("--verbose", action="store_true") + parser = argparse.ArgumentParser( + prog="uenv", + formatter_class=argparse.RawDescriptionHelpFormatter, + description=textwrap.dedent( +f"""\ +Interact, manage and query with CSCS uenv. + +For more information on how to use individual commands use the --help flag. + +{colorize("Example", "blue")}: get help on the start command + {colorize("uenv start --help", "white")} +""" + )) + parser.add_argument("--no-color", action="store_true", help="disable color output") + parser.add_argument("--verbose", action="store_true", help="verbose output for debugging") + parser.add_argument("-r", "--repo", + required=False, default=None, type=str, + help="The path on the local filesystem where uenv are managed. By default the environment variable UENV_REPO_PATH is used, if set, otherwise $SCRATCH/.uenv-images is used if the SCRATCH environment variable is set. This option will override these defaults, and must be set if neither of the defaults is set.") + subparsers = parser.add_subparsers(dest="command") + #### run run_parser = subparsers.add_parser("run", - help="run a command in a user environment") + help="run a command in a user environment", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=textwrap.dedent( +f"""\ +Run a command in an environment. + +{colorize("Example", "blue")}: run the script job.sh in an evironmnent + {colorize("uenv run prgenv-gnu ./job.sh", "white")} +This will mount prgenv-gnu, execute job.sh, then return to the calling shell. + +{colorize("Example", "blue")}: the run command can be used when + {colorize("uenv run gromacs/23.1 ./simulation.sh", "white")} + {colorize("uenv run paraview ./render.sh", "white")} +""" + )) run_parser.add_argument("runline", nargs=argparse.REMAINDER, type=str) + #### start start_parser = subparsers.add_parser("start", help="start a user environment") start_parser.add_argument("image", nargs='+', type=str, help="the uenv to start") + #### stop stop_parser = subparsers.add_parser("stop", help="stop a running user environment") + #### status status_parser = subparsers.add_parser("status", help="print information about running environments") + #### modules modules_parser = subparsers.add_parser("modules", help="use modules if they are available") modules_subparsers = modules_parser.add_subparsers(dest="modules_command") @@ -46,6 +94,7 @@ def make_argparser(): modules_use_parser.add_argument("image", nargs='*', type=str, help="the uenv(s) with the modules to load") + #### view views_parser = subparsers.add_parser("view", help="activate a view") views_parser.add_argument("view_name", @@ -57,48 +106,6 @@ def make_argparser(): # Utility functions for string handling, printing, etc. ############################################################################### -# Choose whether to use colored output. -# - by default colored output is ON -# - if the flag --no-color is passed it is OFF -# - if the environment variable NO_COLOR is set it is OFF -def use_colored_output(cli_arg): - # The --no-color argument overrides all environment variables if passed. - if cli_arg: - return False - - # Check the env. var NO_COLOR and disable color if set. - if os.environ.get('NO_COLOR') is not None: - color_var = os.environ.get('NO_COLOR') - if len(color_var)>0 and color_var != "0": - return False - - return True - -def colorize(string, color): - colors = { - "red": "31", - "green": "32", - "yellow": "33", - "blue": "34", - "magenta": "35", - "cyan": "36", - "white": "37", - } - if colored_output: - return f"\033[1;{colors[color]}m{string}\033[0m" - else: - return string - -def print_error(msg): - print(f"{colorize('[error]', 'red')} {msg}", file=sys.stderr) - -def print_warning(msg): - print(f"{colorize('[warning]', 'yellow')} {msg}", file=sys.stderr) - -def print_debug(msg): - if verbose: - print(f"{colorize('[info]', 'cyan')} {msg}", file=sys.stderr) - def echo_from_lines(lines): return [("echo '" + line + "'") for line in lines] @@ -242,11 +249,11 @@ class uenv: if self.modules in [str(e) for e in env_module_paths]: self.modules_loaded = True if self.modules_loaded: - print_debug(f"uenv {self.name} at {str(self.mount)} has modules loaded") + terminal.info(f"uenv {self.name} at {str(self.mount)} has modules loaded") if not self.is_native_mounted: - print_warning(f"The uenv mounted at {self.mount} should be mounted at {self.native_mount}.") - print_warning(f"Features like modules and views will be be disabled.") + terminal.warning(f"The uenv mounted at {self.mount} should be mounted at {self.native_mount}.") + terminal.warning(f"Features like modules and views will be be disabled.") def load_image(self, path): """return information (if any) about the image mounted at path""" @@ -264,10 +271,10 @@ class uenv: env = json.load(fid) return {"mount": path, "uenv": env} except: - print_error(f"unable to read environment configuration {env_file}") + terminal.error(f"unable to read environment configuration {env_file}", abort=False) return empty_image - print_warning(f"the environment mounted at {meta_path} has no configuration {env_file}") + terminal.warning(f"the environment mounted at {meta_path} has no configuration {env_file}") return empty_image @@ -276,28 +283,28 @@ class uenv: # no uenv - return the real mount point if not self.is_uenv: - print_debug(f"image at {self.mount} is not a uenv: setting native mount to {self.mount}") + terminal.info(f"image at {self.mount} is not a uenv: setting native mount to {self.mount}") return self.mount # check whether the mount point is explicitly in the image meta mnt = self._image["uenv"].get("mount", None) if mnt is not None: - print_debug(f"image at {self.mount} has native mount {mnt} set in meta data") + terminal.info(f"image at {self.mount} has native mount {mnt} set in meta data") return pathlib.Path(mnt) # check whether it can be inferred from a view path for e in self.views: - print_debug(f"image at {self.mount} has native mount {pathlib.Path(e['root']).parents[1]} inferred from views") + terminal.info(f"image at {self.mount} has native mount {pathlib.Path(e['root']).parents[1]} inferred from views") return pathlib.Path(e["root"]).parents[1] # check whether it can be inferred from a module path if self.modules is not None: - print_debug(f"image at {self.mount} has native mount {pathlib.Path(e['root']).parents[1]} inferred from modules") + terminal.info(f"image at {self.mount} has native mount {pathlib.Path(e['root']).parents[1]} inferred from modules") return pathlib.Path(self.modules).parents[0] # no mount information found, so assume that the actual mount point is valid. - print_debug(f"image at {self.mount} has native mount {self.mount} assumed") + terminal.info(f"image at {self.mount} has native mount {self.mount} assumed") return self.mount @property @@ -384,7 +391,7 @@ class uenv: def generate_command(args): env = environment() if env.old_api: - print_error(f"the version of squashfs-mount on this system is too old.") + terminal.error(f"the version of squashfs-mount on this system is too old.", abort=False) return shell_error if args.command == "run": @@ -400,7 +407,7 @@ def generate_command(args): elif args.command == "modules": return generate_modules_command(args, env) - print_error(f"unknown command '{args.command}'") + terminal.error(f"unknown command '{args.command}'", abort=False) return shell_error """ @@ -418,15 +425,30 @@ Specifying what to mount will use some "special characters" gromacs@2023/hohgant/a100:/user-environment |system=hohgant,arch=a100,gromacs:2023|/user-environment $SCRATCH/gromacs.squashfs:/user-environment + + name[/version][:tag][:mount] + + gromacs + gromacs None + gromacs:/uenv + gromacs /uenv + gromacs:uenv + gromacs:uenv None + gromacs/5.0:uenv + gromacs/5.0:uenv None + gromacs/5.0:uenv:/uenv + gromacs/5.0:uenv /uenv + gromacs/5.0:/uenv + gromacs/5.0 /uenv + + mount must always be an absolute path, so must start with '/' """ def parse_image_description(desc): - s = desc.split(':') - if len(s)==1: - return {"image": s[0], "mount": None} - elif len(s)==2: + s = desc.rsplit(':', 1) + if len(s)==2 and s[1].startswith("/"): return {"image": s[0], "mount": s[1]} - return None + return {"image": desc, "mount": None} def generate_image_pairs(mnt_list): """ @@ -444,16 +466,12 @@ def generate_image_pairs(mnt_list): # check for zero length input if num_images==0: - print_error("No environment is provided") + terminal.error("No environment is provided", abort=False) return [] # inspect the first image img = parse_image_description(descriptions[0]) - if img is None: - # handle error - print_error(f"Invalid image description {descriptions[0]}") - return [] - elif img["mount"] is None: + if img["mount"] is None: img["mount"] = "/user-environment" else: implicit_mounts = False @@ -465,29 +483,20 @@ def generate_image_pairs(mnt_list): # inspect the second image img = parse_image_description(descriptions[1]) - - if img is None: - # handle error - print_error(f"Invalid image description {descriptions[1]}") - return [] - elif img["mount"] is None: + if img["mount"] is None: if implicit_mounts == True: img["mount"] = "/user-tools" else: - print_error(f"missing mount point in '{descriptions[1]}'") + terminal.error(f"missing mount point in '{descriptions[1]}'", abort=False) return [] - implicit_mounts = False mounts.append(img) for d in descriptions[2:]: img = parse_image_description(d) - if img is None: - print_error(f"Invalid image description '{d}'") - return [] if img["mount"] is None: - print_error(f"missing mount point in '{d}'") + terminal.error(f"missing mount point in '{d}'", abort=False) return [] mounts.append(img) @@ -499,22 +508,51 @@ def parse_image_descriptions(mnt_list): mount_pairs = [] for m in mounts: - img_path = pathlib.Path(m["image"]) - if not img_path.is_absolute(): - img_path = pathlib.Path.cwd() / img_path - if not img_path.is_file(): - print_error(f"the user environment {colorize(img_path, 'white')} does not exist") - return [] + uenv = m["image"] + terminal.info(f"requested mount {m['image']} at {m['mount']}") + uenv_path = pathlib.PosixPath(uenv) + if not uenv_path.is_file(): + terminal.info(f"no file matching description found") + + # TODO support --repo CLI argument + repo_path = alps.uenv_repo_path(None) + terminal.info(f"looking in the repository {repo_path}") + + img_filter = names.create_filter(uenv, require_complete=False) + if os.environ.get('CLUSTER_NAME') is not None: + img_filter["system"] = os.environ.get('CLUSTER_NAME') + else: + terminal.warning("the CLUSTER_NAME environment variable is not set") + + # TODO: the uarch will have to be set for systems with more than one uarch + #img_filter["uarch"] = "gh200" + + with flock.Lock(f"{repo_path}/index.json", flock.Lock.READ) as lk: + fscache = datastore.FileSystemCache(repo_path) + + records = fscache.database.find_records(**img_filter) + + if len(records)==0: + terminal.error(f"no uenv matches the spec: {uenv}") + sys.exit(1) + + terminal.info(f"selected image {records[0]} from repository {repo_path}") + + uenv_path = fscache.image_path(records[0]) + "/store.squashfs" + terminal.info(f"lookup {uenv} returned {uenv_path}") + else: + uenv_path = uenv_path.resolve() mnt_path = pathlib.Path(m["mount"]) - if not mnt_path.is_absolute(): - mnt_path = pathlib.Path.cwd() / mnt_path if not mnt_path.is_dir(): - print_error(f"the mount point {colorize(mnt_path, 'white')} does not exist") + terminal.error(f"the mount point {colorize(mnt_path, 'white')} does not exist", abort=False) + return [] + if not mnt_path.is_absolute(): + terminal.error(f"the mount point {colorize(mnt_path, 'white')} is a relative path", abort=False) return [] - mount_pairs.append(f"{img_path}:{mnt_path}") - print_debug(f" mounting: {img_path}:{mnt_path}") + mount_pairs.append(f"{uenv_path}:{mnt_path}") + terminal.info(f" mounting: {uenv_path}:{mnt_path}") return mount_pairs @@ -536,16 +574,16 @@ def split_runline(args): return images, cmd, cmdargs def generate_run_command(args, env): - print_debug(f"parsing run command with arguments: {args.command}") + terminal.info(f"parsing run command with arguments: {args.command}") if env.active: - print_error("a uenv is already running") + terminal.error("a uenv is already running", abort=False) return shell_error images, cmd, cmdargs = split_runline(args.runline) - print_debug(f"images: {images}") - print_debug(f"cmd: {cmd}") - print_debug(f"cmdargs: {cmdargs}") + terminal.info(f"images: {images}") + terminal.info(f"cmd: {cmd}") + terminal.info(f"cmdargs: {cmdargs}") mount_pairs=parse_image_descriptions(images) if mount_pairs==[]: @@ -564,9 +602,9 @@ def generate_run_command(args, env): def generate_start_command(args, env): - print_debug(f"parsing start command with arguments: {args.image}") + terminal.info(f"parsing start command with arguments: {args.image}") if env.active: - print_error("a uenv is already running") + terminal.error("a uenv is already running", abort=False) return shell_error mount_pairs=parse_image_descriptions(args.image) @@ -585,10 +623,10 @@ def generate_start_command(args, env): def generate_modules_command(args, env): - print_debug(f"parsing modules command: {args}") + terminal.info(f"parsing modules command: {args}") if not env.active: - print_error(f'there is no environment loaded') + terminal.error(f'there is no environment loaded', abort=False) return shell_error # generate a list of all the mounted environments that provide modules @@ -597,7 +635,7 @@ def generate_modules_command(args, env): for e in env.uenvs if (e.modules is not None) and e.is_native_mounted] - print_debug(f"modules are provided by {module_envs}") + terminal.info(f"modules are provided by {module_envs}") # No use command, i.e. the folloing CLI command was made: # uenv modules @@ -633,14 +671,14 @@ def generate_modules_command(args, env): and (e.modules is not None) and (e.is_native_mounted)] if len(matches)==0: - print_error(f"no uenv matching {i} provides modules") + terminal.error(f"no uenv matching {i} provides modules", abort=False) return shell_error - print_debug(f" uenv {i} mounted at {matches[0]}") + terminal.info(f" uenv {i} mounted at {matches[0]}") mounts.append(matches[0]) modulepaths = [str(p / 'modules') for p in mounts] for p in modulepaths: - print_debug(f" using modules in {p}") + terminal.info(f" using modules in {p}") modulecmds = [f"module use {p}" for p in modulepaths] modulecmds.append(f"export UENV_MODULE_PATH={','.join(modulepaths)}") return modulecmds @@ -651,17 +689,17 @@ def generate_modules_command(args, env): def generate_view_command(args, env): if not env.active: - print_error(f'there is no environment loaded') + terminal.error(f'there is no environment loaded', abort=False) return shell_error if env.loaded_view is not None: - print_error(f'a view is already loaded: {env.loaded_view}') + terminal.error(f'a view is already loaded: {env.loaded_view}', abort=False) return shell_error uenv = env.uenvs[0] name = args.view_name - print_debug(f"trying to load view {name} in env loaded at {uenv.mount}") + terminal.info(f"trying to load view {name} in env loaded at {uenv.mount}") available_views = [v["name"] for v in uenv.views] if name in available_views: @@ -670,7 +708,7 @@ def generate_view_command(args, env): return [f"source '{path}'", f"export UENV_VIEW={uenv.mount}:{name}",] else: - print_error(f'the view "{name}" is not one of the available views: {available_views}') + terminal.error(f'the view "{name}" is not one of the available views: {available_views}', abort=False) return shell_error return shell_noop @@ -684,7 +722,7 @@ def generate_status_command(args, env): first = True loaded_view = env.loaded_view - print_debug(f"loaded view: {loaded_view}") + terminal.info(f"loaded view: {loaded_view}") for uenv in env.uenvs: if not first: @@ -737,7 +775,7 @@ def generate_status_command(args, env): def generate_stop_command(args, env): if not env.active: - print_error(f"there is no running environment to stop") + terminal.error(f"there is no running environment to stop", abort=False) return shell_error return "exit $_last_exitcode" @@ -747,11 +785,10 @@ if __name__ == "__main__": parser = make_argparser() args = parser.parse_args() - global colored_output - colored_output = use_colored_output(args.no_color) + terminal.use_colored_output(args.no_color) - global verbose - verbose = args.verbose + if args.verbose: + terminal.set_debug_level(2) cmd = generate_command(args)