|
| 1 | +# Author: Samuel Omlin, CSCS (omlins) |
| 2 | +# |
| 3 | +# Description: Creation of an uenv view equivalent to the activation script in JUHPC. |
| 4 | + |
| 5 | + |
| 6 | +using Pkg; Pkg.add("JSON") |
| 7 | +using JSON |
| 8 | + |
| 9 | + |
| 10 | +# Functions |
| 11 | + |
| 12 | +function check_env(vars...) |
| 13 | + for var in vars |
| 14 | + if !haskey(ENV, var) |
| 15 | + error("uenv-view: $var is not set.") |
| 16 | + elseif isempty(ENV[var]) |
| 17 | + error("uenv-view: $var is empty.") |
| 18 | + end |
| 19 | + end |
| 20 | +end |
| 21 | + |
| 22 | +function uenv_env(key::AbstractString) |
| 23 | + s = ENV[key] |
| 24 | + pattern = r"\$\{(.*?)\}" |
| 25 | + for m in eachmatch(pattern, s) |
| 26 | + s = replace(s, m.match => "\${@" * m.captures[1] * "@}") |
| 27 | + end |
| 28 | + return s |
| 29 | +end |
| 30 | + |
| 31 | + |
| 32 | +# Check if the required environment variables are set and not empty |
| 33 | + |
| 34 | +check_env("JULIA_WRAPPER_BINDIR", "JULIAUP_WRAPPER_BINDIR", "JULIAUP_BINDIR", "JULIAUP_DEPOT", "JULIA_DEPOT", "JULIA_PREFDIR", "JULIAUP_INSTALLDIR", "ENV_JSON") |
| 35 | + |
| 36 | + |
| 37 | +# Define the environment variables part of the juliaup view |
| 38 | + |
| 39 | +juliaup_view_env = Dict( |
| 40 | + "values" => Dict( |
| 41 | + "list" => Dict( |
| 42 | + "PATH" => [ |
| 43 | + Dict( |
| 44 | + "value" => [uenv_env("JULIAUP_WRAPPER_BINDIR"), uenv_env("JULIAUP_BINDIR")], # The wrapper must be before the juliaup bindir |
| 45 | + "op" => "prepend", |
| 46 | + ), |
| 47 | + ], |
| 48 | + ), |
| 49 | + "scalar" => Dict( |
| 50 | + "JULIAUP_DEPOT_PATH" => uenv_env("JULIAUP_DEPOT"), |
| 51 | + "JULIA_DEPOT_PATH" => uenv_env("JULIA_DEPOT"), |
| 52 | + "JULIA_LOAD_PATH" => ":$(uenv_env("JULIA_PREFDIR"))", # ":" means appending! |
| 53 | + [key => uenv_env(env_key) for (key, env_key) in |
| 54 | + [("CUDA_HOME", "JUHPC_CUDA_HOME"), |
| 55 | + ("ROCM_PATH", "JUHPC_ROCM_HOME"), |
| 56 | + ("JULIA_ADIOS2_PATH", "JUHPC_ADIOS2_HOME")] |
| 57 | + if haskey(ENV, env_key) && !isempty(uenv_env(env_key))]..., # Conditional inclusion for variables if they are set and not empty |
| 58 | + [("JULIA_CUDA_MEMORY_POOL" => "none") for env_key in ["JUHPC_CUDA_HOME"] if haskey(ENV, env_key) && !isempty(uenv_env(env_key))]... # Conditionally include JULIA_CUDA_MEMORY_POOL |
| 59 | + ) |
| 60 | + ), |
| 61 | + "version" => 1 |
| 62 | +) |
| 63 | + |
| 64 | + |
| 65 | +# Define the environment variables part of the jupyter view |
| 66 | + |
| 67 | +jupyter_view_env = deepcopy(juliaup_view_env) |
| 68 | +jupyter_view_env["values"]["list"]["PATH"] = [ |
| 69 | + Dict( |
| 70 | + "value" => [uenv_env("IJULIA_INSTALLER_BINDIR"), uenv_env("JULIA_WRAPPER_BINDIR"), uenv_env("JULIAUP_WRAPPER_BINDIR"), uenv_env("JULIAUP_BINDIR")], # The wrappers must be before the juliaup bindir |
| 71 | + "op" => "prepend", |
| 72 | + ), |
| 73 | +] |
| 74 | + |
| 75 | + |
| 76 | +# Define the juliaup view |
| 77 | + |
| 78 | +juliaup_view = Dict( |
| 79 | + "juliaup" => Dict( |
| 80 | + "root" => "/user-environment/env/juliaup", |
| 81 | + "env" => juliaup_view_env, |
| 82 | + "activate" => "/dev/null", |
| 83 | + "description" => "description: HPC setup for Juliaup, Julia and some HPC key packages (Juliaup and Julia are installed on first execution of `juliaup`` in $(ENV["JULIAUP_INSTALLDIR"]))", |
| 84 | + "type" => "augment" |
| 85 | + ) |
| 86 | +) |
| 87 | + |
| 88 | + |
| 89 | +# Define the jupyter view |
| 90 | +jupyter_view = Dict( |
| 91 | + "jupyter" => Dict( |
| 92 | + "root" => "/user-environment/env/jupyter", |
| 93 | + "env" => jupyter_view_env, |
| 94 | + "activate" => "/dev/null", |
| 95 | + "description" => "description: HPC setup for Julia in Jupyter: IJulia can be installed from a Jupyter terminal by typing `install_ijulia`. The setup includes furthermore Juliaup, Julia and some HPC key packages (if not present, Juliaup and Julia are installed on first execution of `install_ijulia`/`julia`/`juliaup` in $(ENV["JULIAUP_INSTALLDIR"]))", |
| 96 | + "type" => "augment" |
| 97 | + ) |
| 98 | +) |
| 99 | + |
| 100 | + |
| 101 | +# Merge the juliaup view with the existing views in the JSON file |
| 102 | + |
| 103 | +env = JSON.parsefile(ENV["ENV_JSON"]) |
| 104 | +views = env["views"] |
| 105 | +views = merge(views, juliaup_view) |
| 106 | +views = merge(views, jupyter_view) |
| 107 | +env["views"] = views |
| 108 | +open(ENV["ENV_JSON"],"w") do f |
| 109 | + JSON.print(f, env, 4) |
| 110 | +end |
| 111 | + |
| 112 | + |
| 113 | +# Print the new views to the console |
| 114 | +@info "New views:" |
| 115 | +@info "Juliaup view:" |
| 116 | +JSON.print(stdout, juliaup_view, 4) |
| 117 | +@info "Jupyter view:" |
| 118 | +JSON.print(stdout, jupyter_view, 4) |
| 119 | + |
| 120 | + |
| 121 | +# Remove the added package |
| 122 | +Pkg.rm("JSON") |
0 commit comments