From b5a8254c6fffb4f96cf02f8d8a83aaeb264ad51a Mon Sep 17 00:00:00 2001 From: bcumming Date: Wed, 26 Jun 2024 16:55:28 +0200 Subject: [PATCH] uenv view now loads using json data if available --- lib/envvars.py | 383 ++++++++++++++++++++++++++++++++++++++++--------- uenv-impl | 38 +++-- 2 files changed, 345 insertions(+), 76 deletions(-) diff --git a/lib/envvars.py b/lib/envvars.py index 68f0074..f73ff89 100644 --- a/lib/envvars.py +++ b/lib/envvars.py @@ -1,30 +1,39 @@ -from enum import Enum +#!/usr/bin/python3 + +import argparse import json import os -from typing import Optional, List +from enum import Enum +from typing import List, Optional + +import yaml -class EnvVarOp (Enum): - PREPEND=1 - APPEND=2 - SET=3 + +class EnvVarOp(Enum): + PREPEND = 1 + APPEND = 2 + SET = 3 def __str__(self): return self.name.lower() -class EnvVarKind (Enum): - SCALAR=2 - LIST=2 + +class EnvVarKind(Enum): + SCALAR = 2 + LIST = 2 + list_variables = { - "ACLOCAL_PATH", - "CMAKE_PREFIX_PATH", - "CPATH", - "LD_LIBRARY_PATH", - "LIBRARY_PATH", - "MANPATH", - "PATH", - "PKG_CONFIG_PATH", - } + "ACLOCAL_PATH", + "CMAKE_PREFIX_PATH", + "CPATH", + "LD_LIBRARY_PATH", + "LIBRARY_PATH", + "MANPATH", + "PATH", + "PKG_CONFIG_PATH", +} + class EnvVarError(Exception): """Exception raised when there is an error with environment variable manipulation.""" @@ -36,13 +45,15 @@ def __init__(self, message): def __str__(self): return self.message + def is_env_value_list(v): return isinstance(v, list) and all(isinstance(item, str) for item in v) -class ListEnvVarUpdate(): + +class ListEnvVarUpdate: def __init__(self, value: List[str], op: EnvVarOp): - # strip white space from each entry - self._value = [v.strip() for v in value] + # clean up paths as they are inserted + self._value = [os.path.normpath(p) for p in value] self._op = op @property @@ -53,12 +64,21 @@ def op(self): def value(self): return self._value + def set_op(self, op: EnvVarOp): + self._op = op + + # remove all paths that have root as common root + def remove_root(self, root: str): + root = os.path.normpath(root) + self._value = [p for p in self._value if root != os.path.commonprefix([root, p])] + def __repr__(self): return f"envvar.ListEnvVarUpdate({self.value}, {self.op})" def __str__(self): return f"({self.value}, {self.op})" + class EnvVar: def __init__(self, name: str): self._name = name @@ -67,60 +87,76 @@ def __init__(self, name: str): def name(self): return self._name + class ListEnvVar(EnvVar): def __init__(self, name: str, value: List[str], op: EnvVarOp): super().__init__(name) self._updates = [ListEnvVarUpdate(value, op)] - def update(self, value: List[str], op:EnvVarOp): + def update(self, value: List[str], op: EnvVarOp): self._updates.append(ListEnvVarUpdate(value, op)) + def remove_root(self, root: str): + for i in range(len(self._updates)): + self._updates[i].remove_root(root) + @property def updates(self): return self._updates - def concat(self, other: 'ListEnvVar'): + def concat(self, other: "ListEnvVar"): self._updates += other.updates + def make_dirty(self): + if len(self._updates) > 0: + self._updates[0].set_op(EnvVarOp.PREPEND) + + @property + def paths(self): + paths = [] + for u in self._updates: + paths += u.value + return paths + # Given the current value, return the value that should be set # current is None implies that the variable is not set # # dirty allows for not overriding the current value of the variable. - def get_value(self, current: Optional[str], dirty: bool=False): + def get_value(self, current: Optional[str], dirty: bool = False): v = current # if the variable is currently not set, first initialise it as empty. if v is None: - if len(self._updates)==0: + if len(self._updates) == 0: return None v = "" first = True for update in self._updates: joined = ":".join(update.value) - if first and dirty and update.op==EnvVarOp.SET: + if first and dirty and update.op == EnvVarOp.SET: op = EnvVarOp.PREPEND else: op = update.op - if v == "" or op==EnvVarOp.SET: + if v == "" or op == EnvVarOp.SET: v = joined - elif op==EnvVarOp.APPEND: + elif op == EnvVarOp.APPEND: v = ":".join([v, joined]) - elif op==EnvVarOp.PREPEND: + elif op == EnvVarOp.PREPEND: v = ":".join([joined, v]) else: - raise EnvVarError(f"Internal error: implement the operation {update.op}"); + raise EnvVarError(f"Internal error: implement the operation {update.op}") first = False # strip any leading/trailing ":" - v = v.strip(':') + v = v.strip(":") return v def __repr__(self): - return f"envvars.ListEnvVar(\"{self.name}\", {self._updates})" + return f'envvars.ListEnvVar("{self.name}", {self._updates})' def __str__(self): return f"(\"{self.name}\": [{','.join([str(u) for u in self._updates])}])" @@ -148,10 +184,11 @@ def get_value(self, value: Optional[str]): return self._value def __repr__(self): - return f"envvars.ScalarEnvVar(\"{self.name}\", \"{self.value}\")" + return f'envvars.ScalarEnvVar("{self.name}", "{self.value}")' def __str__(self): - return f"(\"{self.name}\": \"{self.value}\")" + return f'("{self.name}": "{self.value}")' + class Env: def __init__(self): @@ -160,11 +197,13 @@ def __init__(self): def apply(self, var: EnvVar): self._vars[var.name] = var + # returns true if the environment variable with name is a list variable, # e.g. PATH, LD_LIBRARY_PATH, PKG_CONFIG_PATH, etc. def is_list_var(name: str) -> bool: return name in list_variables + class EnvVarSet: """ A set of environment variable updates. @@ -190,19 +229,26 @@ def clear(self): def scalars(self): return self._scalars + def make_dirty(self): + for name in self._lists: + self._lists[name].make_dirty() + + def remove_root(self, root: str): + for name in self._lists: + self._lists[name].remove_root(root) + def set_scalar(self, name: str, value: str): self._scalars[name] = ScalarEnvVar(name, value) def set_list(self, name: str, value: List[str], op: EnvVarOp): var = ListEnvVar(name, value, op) if var.name in self._lists.keys(): - old = self._lists[var.name] self._lists[var.name].concat(var) else: self._lists[var.name] = var def __repr__(self): - return f"envvars.EnvVarSet(\"{self.lists}\", \"{self.scalars}\")" + return f'envvars.EnvVarSet("{self.lists}", "{self.scalars}")' def __str__(self): s = "EnvVarSet:\n" @@ -219,7 +265,7 @@ def __str__(self): # Update the environment variables using the values in another EnvVarSet. # This operation is used when environment variables are sourced from more # than one location, e.g. multiple activation scripts. - def update(self, other: 'EnvVarSet'): + def update(self, other: "EnvVarSet"): for name, var in other.scalars.items(): self.set_scalar(name, var.value) for name, var in other.lists.items(): @@ -234,7 +280,7 @@ def update(self, other: 'EnvVarSet'): # "post": the list of commands to be executed to revert the environment # # The "post" list is optional, and should not be used for commands that - # update the environment like "uenv view", instead + # update the environment like "uenv view" and "uenv modules use", instead # it should be used for commands that should not alter the calling environment, # like "uenv run" and "uenv start". # @@ -277,9 +323,26 @@ def export(self, dirty=False): return {"pre": pre, "post": post} + def as_dict(self) -> dict: + # create a dictionary with the information formatted for JSON + d = {"list": {}, "scalar": {}} + + for name, var in self.lists.items(): + ops = [] + for u in var.updates: + op = "set" if u.op == EnvVarOp.SET else ("prepend" if u.op == EnvVarOp.PREPEND else "append") + ops.append({"op": op, "value": u.value}) + + d["list"][name] = ops + + for name, var in self.scalars.items(): + d["scalar"][name] = var.value + + return d + # returns a string that represents the environment variable modifications # in json format - #{ + # { # "list": { # "PATH": [ # {"op": "set", "value": "/user-environment/bin"}, @@ -294,62 +357,46 @@ def export(self, dirty=False): # "CUDA_HOME": "/user-environment/env/default", # "MPIF90": "/user-environment/env/default/bin/mpif90" # } - #} - def json(self) -> str: - # create a dictionary with the information formatted for JSON - d = {"list": {}, "scalar": {}} - - for name, var in self.lists.items(): - ops = [] - for u in var.updates: - op = "set" if u.op == EnvVarOp.SET else ("prepend" if u.op==EnvVarOp.PREPEND else "append") - ops.append({"op": op, "value": u.value}) - - d["list"][name] = ops - - for name, var in self.scalars.items(): - d["scalar"][name] = var.value - - return json.dumps(d, separators=(',', ':')) + # } + def as_json(self) -> str: + return json.dumps(self.as_dict(), separators=(",", ":")) def set_post(self, value: bool): self._generate_post = value -def read_activation_script(filename: str, env: Optional[EnvVarSet]=None) -> EnvVarSet: - scalars = {} - lists = {} + +def read_activation_script(filename: str, env: Optional[EnvVarSet] = None) -> EnvVarSet: if env is None: env = EnvVarSet() with open(filename) as fid: for line in fid: - l = line.strip().rstrip(";") + ls = line.strip().rstrip(";") # skip empty lines and comments - if (len(l)==0) or (l[0]=='#'): + if (len(ls) == 0) or (ls[0] == "#"): continue # split on the first whitespace # this splits lines of the form # export Y # where Y is an arbitray string into ['export', 'Y'] - fields = l.split(maxsplit=1) + fields = ls.split(maxsplit=1) # handle lines of the form 'export Y' - if len(fields)>1 and fields[0]=='export': - fields = fields[1].split('=', maxsplit=1) + if len(fields) > 1 and fields[0] == "export": + fields = fields[1].split("=", maxsplit=1) # get the name of the environment variable name = fields[0] # if there was only one field, there was no = sign, so pass - if len(fields)<2: + if len(fields) < 2: continue # rhs the value that is assigned to the environment variable rhs = fields[1] if name in list_variables: - fields = [f for f in rhs.split(":") if len(f.strip())>0] - lists[name] = fields + fields = [f for f in rhs.split(":") if len(f.strip()) > 0] # look for $name as one of the fields (only works for append or prepend) - if len(fields)==0: + if len(fields) == 0: env.set_list(name, fields, EnvVarOp.SET) elif fields[0] == f"${name}": env.set_list(name, fields[1:], EnvVarOp.APPEND) @@ -362,4 +409,204 @@ def read_activation_script(filename: str, env: Optional[EnvVarSet]=None) -> EnvV return env +def read_dictionary(d: dict, env: Optional[EnvVarSet] = None) -> EnvVarSet: + if env is None: + env = EnvVarSet() + + if "scalar" in d: + for name, value in d["scalar"].items(): + env.set_scalar(name, value) + + if "list" in d: + for name, updates in d["list"].items(): + for u in updates: + if u["op"] == "set": + env.set_list(name, u["value"], EnvVarOp.SET) + elif u["op"] == "prepend": + env.set_list(name, u["value"], EnvVarOp.PREPEND) + elif u["op"] == "append": + env.set_list(name, u["value"], EnvVarOp.APPEND) + # just ignore updtes with invalid "op" + + return env + + +def view_impl(args): + print( + f"parsing view {args.root}\n compilers {args.compilers}\n prefix_paths '{args.prefix_paths}'\n \ + build_path '{args.build_path}'" + ) + + if not os.path.isdir(args.root): + print(f"error - environment root path {args.root} does not exist") + exit(1) + + root_path = args.root + activate_path = root_path + "/activate.sh" + if not os.path.isfile(activate_path): + print(f"error - activation script {activate_path} does not exist") + exit(1) + + envvars = read_activation_script(activate_path) + + # force all prefix path style variables (list vars) to use PREPEND the first operation. + envvars.make_dirty() + envvars.remove_root(args.build_path) + + if args.compilers is not None: + if not os.path.isfile(args.compilers): + print(f"error - compiler yaml file {args.compilers} does not exist") + exit(1) + + with open(args.compilers, "r") as file: + data = yaml.safe_load(file) + compilers = [c["compiler"] for c in data["compilers"]] + + compiler_paths = [] + for c in compilers: + local_paths = set([os.path.dirname(v) for _, v in c["paths"].items() if v is not None]) + compiler_paths += local_paths + print(f'adding compiler {c["spec"]} -> {[p for p in local_paths]}') + + envvars.set_list("PATH", compiler_paths, EnvVarOp.PREPEND) + + if args.prefix_paths: + # get the root path of the env + print(f"prefix_paths: searching in {root_path}") + + for p in args.prefix_paths.split(","): + name, value = p.split("=") + paths = [] + for path in [os.path.normpath(p) for p in value.split(":")]: + test_path = f"{root_path}/{path}" + if os.path.isdir(test_path): + paths.append(test_path) + + print(f"{name}:") + for p in paths: + print(f" {p}") + + if len(paths) > 0: + if name in envvars.lists: + ld_paths = envvars.lists[name].paths + final_paths = [p for p in paths if p not in ld_paths] + envvars.set_list(name, final_paths, EnvVarOp.PREPEND) + else: + envvars.set_list(name, paths, EnvVarOp.PREPEND) + + json_path = os.path.join(root_path, "env.json") + print(f"writing JSON data to {json_path}") + envvar_dict = {"version": 1, "values": envvars.as_dict()} + with open(json_path, "w") as fid: + json.dump(envvar_dict, fid) + fid.write("\n") + + +def meta_impl(args): + # verify that the paths exist + if not os.path.exists(args.mount): + print(f"error - uenv mount '{args.mount}' does not exist.") + exit(1) + + # parse the uenv meta data from file + meta_in_path = os.path.normpath(f"{args.mount}/meta/env.json.in") + meta_path = os.path.normpath(f"{args.mount}/meta/env.json") + print(f"loading meta data to update: {meta_in_path}") + with open(meta_in_path) as fid: + meta = json.load(fid) + + for name, data in meta["views"].items(): + env_root = data["root"] + + # read the json view data from file + json_path = os.path.join(env_root, "env.json") + print(f"reading view {name} data rom {json_path}") + + if not os.path.exists(json_path): + print(f"error - meta data file '{json_path}' does not exist.") + exit(1) + + with open(json_path, "r") as fid: + envvar_dict = json.load(fid) + + # update the global meta data to include the environment variable state + meta["views"][name]["env"] = envvar_dict + meta["views"][name]["type"] = "spack-view" + + # process spack and modules + if args.modules: + module_path = f"{args.mount}/modules" + meta["views"]["modules"] = { + "activate": "/dev/null", + "description": "activate modules", + "root": module_path, + "env": { + "version": 1, + "type": "augment", + "values": {"list": {"MODULEPATH": [{"op": "prepend", "value": [module_path]}]}, "scalar": {}}, + }, + } + + if args.spack is not None: + spack_url, spack_version = args.spack.split(",") + spack_path = f"{args.mount}/config".replace("//", "/") + meta["views"]["spack"] = { + "activate": "/dev/null", + "description": "configure spack upstream", + "root": spack_path, + "env": { + "version": 1, + "type": "augment", + "values": { + "list": {}, + "scalar": { + "UENV_SPACK_CONFIG_PATH": spack_path, + "UENV_SPACK_COMMIT": spack_version, + "UENV_SPACK_URL": spack_url, + }, + }, + }, + } + # update the uenv meta data file with the new env. variable description + with open(meta_path, "w") as fid: + # write updated meta data + json.dump(meta, fid) + fid.write("\n") + print(f"wrote the uenv meta data {meta_path}") + + +if __name__ == "__main__": + # parse CLI arguments + parser = argparse.ArgumentParser() + subparsers = parser.add_subparsers(dest="command") + view_parser = subparsers.add_parser( + "view", formatter_class=argparse.RawDescriptionHelpFormatter, help="generate env.json for a view" + ) + view_parser.add_argument("root", help="root path of the view", type=str) + view_parser.add_argument("build_path", help="build_path", type=str) + view_parser.add_argument( + "--prefix_paths", help="a list of relative prefix path searchs of the form X=y:z,Y=p:q", default="", type=str + ) + # only add compilers if this argument is passed + view_parser.add_argument("--compilers", help="path of the compilers.yaml file", type=str, default=None) + + uenv_parser = subparsers.add_parser( + "uenv", + formatter_class=argparse.RawDescriptionHelpFormatter, + help="generate meta.json meta data file for a uenv.", + ) + uenv_parser.add_argument("mount", help="mount point of the image", type=str) + uenv_parser.add_argument("--modules", help="configure a module view", action="store_true") + uenv_parser.add_argument( + "--spack", help='configure a spack view. Format is "spack_url,git_commit"', type=str, default=None + ) + + args = parser.parse_args() + + if args.command == "uenv": + print("!!! running meta") + meta_impl(args) + elif args.command == "view": + print("!!! running view") + view_impl(args) diff --git a/uenv-impl b/uenv-impl index 1f7bc38..cbd02d2 100755 --- a/uenv-impl +++ b/uenv-impl @@ -408,11 +408,22 @@ class uenv: views = [] for name, info in vlist.items(): - views.append({ - "name": name, - "activate": info["activate"], - "description": info["description"], - "root": info["root"]}) + if "env" in info: + if info["env"]["version"] != 1: + terminal.warning(f"this uenv image is too recent for uenv - please upgrade.") + return [] + views.append({ + "name": name, + "activate": info["activate"], + "description": info["description"], + "root": info["root"], + "env": info["env"]["values"]}) + else: + views.append({ + "name": name, + "activate": info["activate"], + "description": info["description"], + "root": info["root"]}) return views @@ -796,10 +807,21 @@ def generate_view_command(args, env, env_vars): terminal.error(f'the view "{requested_view}" is not available', abort=False) return help_message(shell_error) - path = next((v['activate'] for v in uenv.views if v['name']==vname)) + # load the raw view data + view_meta = next((v for v in uenv.views if v["name"]==vname)) + + # check whether the view provides environment variable information + terminal.info(f"view meta: {view_meta}") + if "env" in view_meta: + terminal.info(f"the image provides environment variable meta data") + env_vars.update(ev.read_dictionary(view_meta["env"])) + terminal.info(f"{env_vars}") + # else use the raw activate.sh + else: + path = view_meta['activate'] + terminal.info(f"full path for view activation script: '{path}'") + env_vars.update(ev.read_activation_script(path)) - terminal.info(f"full path for view activation script: '{path}'") - env_vars.update(ev.read_activation_script(path)) env_vars.set_scalar("UENV_VIEW", f"{uenv.mount}:{uenv.name}:{vname}") env_vars.set_post(False)