Support environment with restricted commands

environments/payu-dev/config.sh

@@ -27,4 +27,15 @@ declare -a rpms_to_remove=()
 declare -a replace_from_apps=()
 declare -a outside_commands_to_include=( "pbs_tmrsh" )
 declare -a outside_files_to_copy=()
-declare -a replace_with_external=()
+declare -a replace_with_external=()
+
+declare -a launcher_commands=(
+    "payu-run"
+    "payu-collate"
+    "payu-sync"
+    "payu"
+    "payu-branch"
+    "payu-checkout"
+    "payu-clone"
+    "payu-profile"
+)

modules/common_cmd_v1

@@ -0,0 +1,32 @@
+#%Module1.0
+
+set prefix __CONDA_BASE__/__APPS_SUBDIR__
+set package __CONDA_INSTALL_BASENAME__
+
+# Prevent running this module with a running payu module
+conflict payu
+
+# Name of this module's environment
+lassign [split [module-info name] {/}] module_name module_version
+if {[string match *-lite $module_version]} {
+    # Remove the '-lite' suffix for the conda environment name
+    set condaenv_version [string range $module_version 0 end-5]
+} else {
+    set condaenv_version $module_version
+}
+set basedir "$prefix/$package/envs/$condaenv_version"
+if {![ file exists $basedir ]} {
+    # For modulenames which are $ENVIRONMENT/$VERSION-lite, rather than conda/$ENIRONMENT-$VERSION-lite
+    set condaenv "${module_name}-${condaenv_version}"
+    set basedir "$prefix/$package/envs/$condaenv"
+}
+
+set myscripts [ file normalize __CONDA_BASE__/__SCRIPT_SUBDIR__/$condaenv-lite.d/bin ]
+set overlay_path [ string map {/conda/ /envs/} $basedir ].sqsh
+
+module load singularity
+
+prepend-path CONTAINER_OVERLAY_PATH $overlay_path
+
+# Add launcher script directory to PATH
+prepend-path PATH $myscripts

scripts/build.sh

@@ -120,9 +120,10 @@ function inner() {
     done
     mkdir -p "${CONDA_MODULE_PATH}"
     copy_and_replace "${SCRIPT_DIR}"/../modules/common_v3 "${CONDA_MODULE_PATH}"/.common_v3       CONDA_BASE APPS_SUBDIR CONDA_INSTALL_BASENAME SCRIPT_SUBDIR
-    copy_and_replace "${SCRIPT_DIR}"/launcher_conf.sh "${CONDA_SCRIPT_PATH}"/launcher_conf.sh CONDA_BASE APPS_SUBDIR CONDA_INSTALL_BASENAME CONDA_SCRIPT_PATH FULLENV
+    copy_and_replace "${SCRIPT_DIR}"/../modules/common_cmd_v1 "${CONDA_MODULE_PATH}"/.common_cmd_v1       CONDA_BASE APPS_SUBDIR CONDA_INSTALL_BASENAME SCRIPT_SUBDIR
+    copy_and_replace "${SCRIPT_DIR}"/launcher_conf.sh "${CONDA_SCRIPT_PATH}"/launcher_conf.sh CONDA_BASE APPS_SUBDIR CONDA_INSTALL_BASENAME
 
-    ### Create symlink tree
+    ### Create symlink tree for full environment
     mkdir -p "${CONDA_SCRIPT_PATH}"/"${FULLENV}".d/{bin,overrides}
     cp "${CONDA_SCRIPT_PATH}"/{launcher.sh,launcher_conf.sh} "${CONDA_SCRIPT_PATH}"/"${FULLENV}".d/bin
     pushd "${CONDA_SCRIPT_PATH}"/"${FULLENV}".d/bin
@@ -143,6 +144,29 @@ function inner() {
     done
     popd
 
+    ### Create symlink tree for lite environment with limited launcher commands
+    if [ "${#launcher_commands[@]}" -gt 0 ]; then
+        mkdir -p "${CONDA_SCRIPT_PATH}"/"${FULLENV}"-lite.d/{bin,overrides}
+        cp "${CONDA_SCRIPT_PATH}"/{launcher.sh,launcher_conf.sh} "${CONDA_SCRIPT_PATH}"/"${FULLENV}"-lite.d/bin
+        pushd "${CONDA_SCRIPT_PATH}"/"${FULLENV}"-lite.d/bin
+        for cmd in "${launcher_commands[@]}"; do
+            ln -s launcher.sh $cmd
+        done
+
+        ### Add in the outside commands
+        for i in "${outside_commands_to_include[@]}"; do
+            ln -s launcher.sh $i
+        done
+        popd
+
+        ### Add in the override and config scripts
+        pushd "${CONDA_SCRIPT_PATH}"/"${FULLENV}-lite".d/overrides
+        for i in ../../overrides/*; do
+            ln -s ${i}
+        done
+        popd
+    fi
+
     if [[ -e "${SCRIPT_DIR}"/../environments/"${CONDA_ENVIRONMENT}"/build_inner.sh ]]; then
         source "${SCRIPT_DIR}"/../environments/"${CONDA_ENVIRONMENT}"/build_inner.sh
     fi
@@ -226,6 +250,9 @@ fi
 
 if [[ "${DO_UPDATE}" == "--install" ]]; then
     ln -s .common_v3 "${CONDA_OUTER_BASE}"/"${MODULE_SUBDIR}"/"${MODULE_NAME}"/"${MODULE_VERSION}"
+    if [ "${#launcher_commands[@]}" -gt 0 ]; then
+        ln -s .common_cmd_v1 "${CONDA_OUTER_BASE}"/"${MODULE_SUBDIR}"/"${MODULE_NAME}"/"${MODULE_VERSION}"-lite
+    fi
 fi
 
 pushd "${CONDA_TEMP_PATH}"
@@ -249,6 +276,10 @@ ln -s /opt/conda/"${FULLENV}" "${CONDA_OUTER_BASE}"/"${APPS_SUBDIR}"/"${CONDA_IN
 ### Can't use ${CONDA_SCRIPT_PATH} or "${CONDA_INSTALLATION_PATH}" due to the need to string match on those paths
 ### which they won't with the '/./' part required for arcane rsync magic
 construct_module_insert "${SINGULARITY_BINARY_PATH}" "${OVERLAY_BASE}" "${my_container}" "${BUILD_STAGE_DIR}"/"${FULLENV}".sqsh.tmp "${SCRIPT_DIR}"/condaenv.sh "${CONDA_INSTALLATION_PATH}" /opt/conda/"${FULLENV}" "${CONDA_BASE}"/"${SCRIPT_SUBDIR}"/"${FULLENV}".d/bin "${CONDA_OUTER_BASE}"/"${MODULE_SUBDIR}"/"${MODULE_NAME}"/."${MODULE_VERSION}"
+construct_launcher_insert "${SINGULARITY_BINARY_PATH}" "${OVERLAY_BASE}" "${my_container}" "${BUILD_STAGE_DIR}"/"${FULLENV}".sqsh.tmp "${SCRIPT_DIR}"/condaenv.sh "${CONDA_INSTALLATION_PATH}" /opt/conda/"${FULLENV}" "${CONDA_BASE}"/"${SCRIPT_SUBDIR}"/"${FULLENV}".d/bin "${CONDA_OUTER_BASE}"/"${SCRIPT_SUBDIR}"/"${FULLENV}".d/bin/launcher_activate.sh
+if [ "${#launcher_commands[@]}" -gt 0 ]; then
+    cp "${CONDA_OUTER_BASE}"/"${SCRIPT_SUBDIR}"/"${FULLENV}".d/bin/launcher_activate.sh "${CONDA_OUTER_BASE}"/"${SCRIPT_SUBDIR}"/"${FULLENV}"-lite.d/bin/launcher_activate.sh
+fi
 
 ### Set permissions on base environment
 set_apps_perms "${CONDA_OUTER_BASE}"

scripts/deploy.sh

@@ -36,6 +36,9 @@ rsync --archive --verbose --partial --progress --one-file-system --itemize-chang
 
 echo "Make sure anything deleted from this environments scripts directory is also deleted from the prod copy"
 rsync --archive --verbose --partial --progress --one-file-system --itemize-changes --hard-links --acls --relative --delete -- "${CONDA_TEMP_PATH}"/"${APPS_SUBDIR}"/./"${CONDA_SCRIPTS_BASENAME}"/"${FULLENV}".d "${CONDA_BASE}"/"${APPS_SUBDIR}"
+if [ "${#launcher_commands[@]}" -gt 0 ]; then
+    rsync --archive --verbose --partial --progress --one-file-system --itemize-changes --hard-links --acls --relative --delete -- "${CONDA_TEMP_PATH}"/"${APPS_SUBDIR}"/./"${CONDA_SCRIPTS_BASENAME}"/"${FULLENV}"-lite.d "${CONDA_BASE}"/"${APPS_SUBDIR}"
+fi
 
 [[ -e "${CONDA_INSTALLATION_PATH}"/envs/"${FULLENV}".sqsh ]] && cp "${CONDA_INSTALLATION_PATH}"/envs/"${FULLENV}".sqsh "${ADMIN_DIR}"/"${FULLENV}".sqsh.bak
 mv "${BUILD_STAGE_DIR}"/"${FULLENV}".sqsh.tmp "${CONDA_INSTALLATION_PATH}"/envs/"${FULLENV}".sqsh

scripts/functions.sh

@@ -172,6 +172,57 @@ function construct_module_insert() {
 
 }
 
+function construct_launcher_insert() {
+
+    singularity_exec="${1}"
+    overlay_path="${2}"
+    container_path="${3}"
+    squashfs_path="${4}"
+    env_script="${5}"
+    rootdir="${6}"
+    condaenv="${7}"
+    script_path="${8}"
+    module_path="${9}"
+
+    declare -a discard_paths=( "/bin" "/usr/bin" "/condabin" )
+    declare -a discard_vars=( "MODULEPATH" "_" "PWD" "SHLVL" )
+
+    while read line; do
+        key="${line%%=*}"
+        value="${line#*=}"
+        ### Skip these environment variables
+        in_array "${discard_vars[@]}" "${key}" && continue
+        ### Prepend to these variables
+        if [[ $key =~ .PATH$ ]]; then
+            echo $key=\"$value:\$$key\"
+            echo export $key=\"\$\{$key\%\:\}\"
+        ### Prepend to Modulefile variables that work like a path
+        elif in_array "_LMFILES_" "LOADEDMODULES" "${key}"; then
+            echo $key=\"$value:\$$key\"
+            echo export $key=\"\$\{$key\%\:\}\"
+        ### Treat path specially - remove system paths and retain order
+        elif [[ "${key}" == "PATH" ]]; then
+            while IFS= read -r -d: entry; do
+                in_array "${discard_paths[@]}" "${entry}" && continue
+                echo PATH=\"$entry:\$PATH\"
+                echo export PATH=\"\$\{PATH\%\:\}\"
+                echo SINGULARITYENV_PREPEND_PATH=\"$entry:\$SINGULARITYENV_PREPEND_PATH\"
+                echo export SINGULARITYENV_PREPEND_PATH=\"\$\{SINGULARITYENV_PREPEND_PATH\%\:\}\"
+            done<<<"${value%:}:"
+        elif [[ "${key}" =~ ^alias\  ]]; then
+            echo alias "${key//alias /}"=\"${value//\'/}\"
+        else
+            if [[ "${value}" ]]; then
+                echo export $key=\"$value\"
+            else
+                echo export $key=\"\"
+            fi
+        fi
+
+    done < <( "${singularity_exec}" -s exec --bind /etc,/half-root,/local,/ram,/run,/system,/usr,/var/lib/sss,/var/run/munge,/var/lib/rpm,"${overlay_path}":/g --overlay="${squashfs_path}"  "${container_path}" /bin/env -i "${env_script}" "${rootdir}" "${condaenv}" ) > "${module_path}"
+
+}
+
 function copy_and_replace() {
     ### Copies the file in $1 to the location in $2 and replaces any occurence
     ### of __${3}__, __${4}__... with the contents of those environment variables

scripts/launcher.sh

@@ -29,6 +29,20 @@ $debug "conf_file = " "${conf_file}"
 
 source "${conf_file}"
 
+# Set up the environment launcher script path - useful for scripts, such as payu,
+# to know how launch the container directly
+export ENV_LAUNCHER_SCRIPT_PATH="${wrapper_bin%/}"/launcher.sh
+$debug "ENV_LAUNCHER_SCRIPT_PATH = " "${ENV_LAUNCHER_SCRIPT_PATH}"
+
+# Check if conda environment already activated
+myenv=$( basename "${wrapper_bin%/*}" ".d" )
+myenv="${myenv%-lite}"
+if [[ "${CONDA_DEFAULT_ENV}" != "/opt/conda/${myenv}" ]]; then
+    activate_script="${wrapper_bin}"/launcher_activate.sh
+    $debug "activate_script = " "${activate_script}"
+    source "${activate_script}"
+fi
+
 ### Add some complicated arguments that are never meant to be used by humans
 declare -a PROG_ARGS=()
 while [[ $# -gt 0 ]]; do
@@ -97,7 +111,6 @@ $debug "cmd_to_run = " "${cmd_to_run[@]}"
 ### Reminder: The --overlay argument that appears LAST takes priority, so put the
 ### default container first, that way if we're intentionally trying to use it from
 ### somewhere else (e.g. jobfs), the one on gdata will be mounted but not used.
-myenv=$( basename "${wrapper_bin%/*}" ".d" )
 if ! [[ "${CONTAINER_OVERLAY_PATH_OVERRIDE}" ]]; then
     if ! [[ :"${CONTAINER_OVERLAY_PATH}": =~ :"${CONDA_BASE_ENV_PATH}"/envs/"${myenv}".sqsh: ]]; then
         [[ -r "${CONDA_BASE_ENV_PATH}"/envs/"${myenv}".sqsh ]] && export CONTAINER_OVERLAY_PATH="${CONDA_BASE_ENV_PATH}"/envs/"${myenv}".sqsh:${CONTAINER_OVERLAY_PATH}
@@ -140,6 +153,7 @@ while IFS= read -r -d: i; do
     in_array "${singularity_default_path[@]}" "${i}" && continue
     [[ "${i}" == "/opt/singularity/bin" ]] && continue
     [[ "${i}" == "${wrapper_bin}" ]] && continue
+    [[ ":${SINGULARITYENV_PREPEND_PATH}:" == *":${i}:"* ]] && continue
     SINGULARITYENV_PREPEND_PATH="${SINGULARITYENV_PREPEND_PATH}:${i}"
 done<<<"${PATH%:}:"
 export SINGULARITYENV_PREPEND_PATH=${SINGULARITYENV_PREPEND_PATH#:*}

scripts/launcher_conf.sh

@@ -5,7 +5,3 @@ export CONTAINER_PATH="__CONDA_BASE__/__APPS_SUBDIR__/__CONDA_INSTALL_BASENAME__
 export CONDA_BASE_ENV_PATH="__CONDA_BASE__/__APPS_SUBDIR__/__CONDA_INSTALL_BASENAME__"
 
 declare -a bind_dirs=( "/etc" "/half-root" "/local" "/ram" "/run" "/system" "/usr" "/var/lib/sss" "/var/run/munge" "/sys/fs/cgroup" "/iointensive" )
-
-# Set up the environment launcher script path - useful for scripts, such as payu,
-# to know how launch the container directly
-export ENV_LAUNCHER_SCRIPT_PATH="__CONDA_SCRIPT_PATH__/__FULLENV__.d/bin/launcher.sh"