diff --git a/shpc/main/container/base.py b/shpc/main/container/base.py index fcb571314..ad09a4494 100644 --- a/shpc/main/container/base.py +++ b/shpc/main/container/base.py @@ -46,6 +46,9 @@ class ContainerTechnology: # The module technology adds extensions here modulefile = "module" + # Wrapper scripts are stored in this subdirectory of the module directory + wrapper_subdir = "bin" + # By default, no extra features features = {} @@ -57,7 +60,7 @@ def __init__(self): self.settings = SettingsBase() - def add(self, sif, module_name, modulefile, template, **kwargs): + def add(self, sif, module_name, modulefile, template, wrapper_template, **kwargs): """ Manually add a registry container. """ @@ -175,5 +178,28 @@ def get_features(self, config_features, settings_features, extra=None): return features + def _generate_wrapper_scripts(self, wrapper_template, aliases, module_dir, features, command=None, container_sif=None, image=None, tty=None): + """ + Generate wrapper scripts for commands (when wrapper_scripts setting enabled) + """ + wrapper_dir = os.path.join(module_dir, self.wrapper_subdir) + shpc.utils.mkdirp([wrapper_dir]) + for alias in aliases: + wrapper_path = os.path.join(wrapper_dir, alias['name']) + out = wrapper_template.render( + alias=alias, + bindpaths=self.settings.bindpaths, + command=command, + container_sif=container_sif, + envfile=self.settings.environment_file, + features=features, + image=image, + module_dir=module_dir, + tty=tty, + wrapper_shell=self.settings.wrapper_shell, + ) + shpc.utils.write_file(wrapper_path, out, exec=True) + return + def __str__(self): return str(self.__class__.__name__) diff --git a/shpc/main/container/docker.py b/shpc/main/container/docker.py index df6685ef6..aaf023760 100644 --- a/shpc/main/container/docker.py +++ b/shpc/main/container/docker.py @@ -189,6 +189,7 @@ def install( version=None, config_features=None, features=None, + wrapper_template=None, ): """Install a general container path to a module @@ -214,6 +215,19 @@ def install( # If there's a tag in the name, don't use it name = name.split(":", 1)[0] + # Option to create wrapper scripts for commands + module_dir = os.path.dirname(module_path) + if self.settings.wrapper_scripts and aliases: + self._generate_wrapper_scripts( + wrapper_template, + aliases, + module_dir, + features, + command=self.command, + image=container_path, + tty=self.settings.enable_tty, + ) + # Make sure to render all values! out = template.render( podman_module=self.settings.podman_module, @@ -223,7 +237,7 @@ def install( else self.settings.docker_shell, image=container_path, description=description, - module_dir=os.path.dirname(module_path), + module_dir=module_dir, aliases=aliases, url=url, features=features, @@ -238,5 +252,7 @@ def install( envfile=self.settings.environment_file, command=self.command, tty=self.settings.enable_tty, + wrapper_scripts=self.settings.wrapper_scripts, + wrapper_subdir=self.wrapper_subdir, ) shpc.utils.write_file(module_path, out) diff --git a/shpc/main/container/singularity.py b/shpc/main/container/singularity.py index 1bd4ac46c..e71a54cc9 100644 --- a/shpc/main/container/singularity.py +++ b/shpc/main/container/singularity.py @@ -71,7 +71,7 @@ def get(self, module_name, env_file=False): logger.exit("Found more than one sif in module folder.") return sif[0] - def add(self, sif, module_name, modulefile, template, **kwargs): + def add(self, sif, module_name, modulefile, template, wrapper_template, **kwargs): """ Manually add a registry container. """ @@ -112,6 +112,7 @@ def add(self, sif, module_name, modulefile, template, **kwargs): template, parsed_name=parsed_name, features=kwargs.get("features"), + wrapper_template=wrapper_template, ) self.add_environment(module_dir, {}, self.settings.environment_file) logger.info("Module %s was created." % (module_name)) @@ -130,6 +131,7 @@ def install( config_features=None, features=None, version=None, + wrapper_template=None, ): """Install a general container path to a module @@ -166,6 +168,17 @@ def install( labels = {} logger.warning("Singularity is not installed, skipping metadata.") + # Option to create wrapper scripts for commands + module_dir = os.path.dirname(module_path) + if self.settings.wrapper_scripts and aliases: + self._generate_wrapper_scripts( + wrapper_template, + aliases, + module_dir, + features, + container_sif=container_path, + ) + # Make sure to render all values! out = template.render( singularity_module=self.settings.singularity_module, @@ -177,7 +190,7 @@ def install( url=url, features=features, version=version, - module_dir=os.path.dirname(module_path), + module_dir=module_dir, labels=labels, deffile=deffile, creation_date=datetime.now(), @@ -186,6 +199,8 @@ def install( registry=parsed_name.registry, repository=parsed_name.repository, envfile=self.settings.environment_file, + wrapper_scripts=self.settings.wrapper_scripts, + wrapper_subdir=self.wrapper_subdir, ) shpc.utils.write_file(module_path, out) diff --git a/shpc/main/modules/__init__.py b/shpc/main/modules/__init__.py index 9b8fd40c4..4c71ab64b 100644 --- a/shpc/main/modules/__init__.py +++ b/shpc/main/modules/__init__.py @@ -79,6 +79,10 @@ def modulefile(self): def templatefile(self): return "%s.%s" % (self.container.templatefile, self.module_extension) + @property + def wrappertemplatefile(self): + return "%s.%s" % (self.container.templatefile, "sh") + def uninstall(self, name, force=False): """ Given a unique resource identifier, uninstall a module @@ -147,7 +151,8 @@ def add(self, sif, module_name, **kwargs): module_name = self.add_namespace(module_name) template = self._load_template(self.templatefile) modulefile = os.path.join(self.settings.module_base, module_name.replace(":", os.sep), self.modulefile) - self.container.add(sif, module_name, modulefile, template, **kwargs) + wrapper_template = self._load_template(self.wrappertemplatefile) + self.container.add(sif, module_name, modulefile, template, wrapper_template, **kwargs) def get(self, module_name, env_file=False): """ @@ -325,6 +330,7 @@ def install(self, name, tag=None, **kwargs): # Get the template based on the module and container type template = self._load_template(self.templatefile) module_path = os.path.join(module_dir, self.modulefile) + wrapper_template = self._load_template(self.wrappertemplatefile) # If the module has a version, overrides version version = tag.name @@ -344,6 +350,7 @@ def install(self, name, tag=None, **kwargs): version=version, config_features=config.features, features=kwargs.get("features"), + wrapper_template=wrapper_template, ) # If the container tech does not need storage, clean up diff --git a/shpc/main/modules/templates/docker.lua b/shpc/main/modules/templates/docker.lua index 967f1958a..2c559530e 100644 --- a/shpc/main/modules/templates/docker.lua +++ b/shpc/main/modules/templates/docker.lua @@ -58,16 +58,16 @@ set_shell_function("{|module_name|}-shell", shellCmd, shellCmd) -- conflict with modules with the same name conflict("{{ tool }}"{% if name != tool %},"{{ name }}"{% endif %}{% if aliases %}{% for alias in aliases %}{% if alias.name != tool %},"{{ alias.name }}"{% endif %}{% endfor %}{% endif %}) --- exec functions to provide "alias" to module commands -{% if aliases %}{% for alias in aliases %} -set_shell_function("{{ alias.name }}", execCmd .. {% if alias.docker_options %} "{{ alias.docker_options }} " .. {% endif %} " --entrypoint {{ alias.entrypoint }} " .. containerPath .. " {{ alias.args }} \"$@\"", execCmd .. {% if alias.docker_options %} "{{ alias.docker_options }} " .. {% endif %} " --entrypoint {{ alias.entrypoint }} " .. containerPath .. " {{ alias.args }}") +-- "aliases" to module commands +{% if wrapper_scripts %}{% if aliases %}prepend_path("PATH", pathJoin(myFileName():match("(.*[/])") or ".", "{{ wrapper_subdir }}")){% endif %} +{% else %}{% if aliases %}{% for alias in aliases %}set_shell_function("{{ alias.name }}", execCmd .. {% if alias.docker_options %} "{{ alias.docker_options }} " .. {% endif %} " --entrypoint {{ alias.entrypoint }} " .. containerPath .. " {{ alias.args }} \"$@\"", execCmd .. {% if alias.docker_options %} "{{ alias.docker_options }} " .. {% endif %} " --entrypoint {{ alias.entrypoint }} " .. containerPath .. " {{ alias.args }}") {% endfor %}{% endif %} {% if aliases %} if (myShellName() == "bash") then {% for alias in aliases %}execute{cmd="export -f {{ alias.name }}", modeA={"load"}} {% endfor %} -end{% endif %} +end{% endif %}{% endif %} -- A customizable exec function set_shell_function("{|module_name|}-exec", execCmd .. " --entrypoint \"\" " .. containerPath .. " \"$@\"", execCmd .. " --entrypoint \"\" " .. containerPath) diff --git a/shpc/main/modules/templates/docker.sh b/shpc/main/modules/templates/docker.sh new file mode 100755 index 000000000..4e4f61bee --- /dev/null +++ b/shpc/main/modules/templates/docker.sh @@ -0,0 +1,3 @@ +#!{{ wrapper_shell }} + +{{ command }} ${PODMAN_OPTS} run ${PODMAN_COMMAND_OPTS} -i{% if tty %}t{% endif %} -u `id -u`:`id -g` --rm {% if envfile %}--env-file {{ module_dir }}/{{ envfile }}{% endif %} {% if bindpaths %}-v {{ bindpaths }} {% endif %}{% if features.home %}-v {{ features.home }} {% endif %} -v ${PWD} -w ${PWD} {% if alias.docker_options %} {{ alias.docker_options }} {% endif %} --entrypoint {{ alias.entrypoint }} {{ image }} {{ alias.args }} {% if '/sh' in wrapper_shell or '/bash' in wrapper_shell %}"$@"{% elif '/csh' in wrapper_shell %}$argv:q{% endif %} diff --git a/shpc/main/modules/templates/docker.tcl b/shpc/main/modules/templates/docker.tcl index 71aac44f7..9d71f8c0e 100644 --- a/shpc/main/modules/templates/docker.tcl +++ b/shpc/main/modules/templates/docker.tcl @@ -73,9 +73,9 @@ set inspectCmd "{{ command }} \${PODMAN_OPTS} inspect ${containerPath}" # set_shell_function takes bashStr and cshStr set-alias {|module_name|}-shell "${shellCmd}" -# exec functions to provide "alias" to module commands -{% if aliases %} -if { [ module-info shell bash ] } { +# "aliases" to module commands +{% if wrapper_scripts %}{% if aliases %}prepend-path PATH "[file dirname ${ModulesCurrentModulefile}]/{{ wrapper_subdir }}"{% endif %} +{% else %}{% if aliases %}if { [ module-info shell bash ] } { if { [ module-info mode load ] } { {% for alias in aliases %} puts stdout "function {{ alias.name }}() { ${execCmd} {% if alias.docker_options %} {{ alias.docker_options | replace("$", "\$") }} {% endif %} --entrypoint {{ alias.entrypoint | replace("$", "\$") }} ${containerPath} {{ alias.args | replace("$", "\$") }} \"\$@\"; }; export -f {{ alias.name }};" {% endfor %} @@ -87,8 +87,7 @@ if { [ module-info shell bash ] } { } else { {% for alias in aliases %} set-alias {{ alias.name }} "${execCmd} {% if alias.docker_options %} {{ alias.docker_options | replace("$", "\$") }} {% endif %} --entrypoint {{ alias.entrypoint | replace("$", "\$") }} ${containerPath} {{ alias.args | replace("$", "\$") }}" {% endfor %} -} -{% endif %} +}{% endif %}{% endif %} # A customizable exec function if { [ module-info shell bash ] } { diff --git a/shpc/main/modules/templates/singularity.lua b/shpc/main/modules/templates/singularity.lua index 4aad55f45..d3674ce33 100644 --- a/shpc/main/modules/templates/singularity.lua +++ b/shpc/main/modules/templates/singularity.lua @@ -61,16 +61,16 @@ set_shell_function("{|module_name|}-shell", shellCmd, shellCmd) -- conflict with modules with the same name conflict("{{ tool }}"{% if name != tool %},"{{ name }}"{% endif %}{% if aliases %}{% for alias in aliases %}{% if alias.name != tool %},"{{ alias.name }}"{% endif %}{% endfor %}{% endif %}) --- exec functions to provide "alias" to module commands -{% if aliases %}{% for alias in aliases %} -set_shell_function("{{ alias.name }}", execCmd .. {% if alias.singularity_options %} "{{ alias.singularity_options }} " .. {% endif %} containerPath .. " {{ alias.command }} \"$@\"", execCmd .. {% if alias.singularity_options %} "{{ alias.singularity_options }} " .. {% endif %} containerPath .. " {{ alias.command }}") +-- "aliases" to module commands +{% if wrapper_scripts %}{% if aliases %}prepend_path("PATH",pathJoin(myFileName():match("(.*[/])") or ".", "{{ wrapper_subdir }}")){% endif %} +{% else %}{% if aliases %}{% for alias in aliases %}set_shell_function("{{ alias.name }}", execCmd .. {% if alias.singularity_options %} "{{ alias.singularity_options }} " .. {% endif %} containerPath .. " {{ alias.command }} \"$@\"", execCmd .. {% if alias.singularity_options %} "{{ alias.singularity_options }} " .. {% endif %} containerPath .. " {{ alias.command }}") {% endfor %}{% endif %} {% if aliases %} if (myShellName() == "bash") then {% for alias in aliases %}execute{cmd="export -f {{ alias.name }}", modeA={"load"}} {% endfor %} -end{% endif %} +end{% endif %}{% endif %} -- A customizable exec function set_shell_function("{|module_name|}-exec", execCmd .. containerPath .. " \"$@\"", execCmd .. containerPath) diff --git a/shpc/main/modules/templates/singularity.sh b/shpc/main/modules/templates/singularity.sh new file mode 100755 index 000000000..9b12aa58d --- /dev/null +++ b/shpc/main/modules/templates/singularity.sh @@ -0,0 +1,3 @@ +#!{{ wrapper_shell }} + +singularity ${SINGULARITY_OPTS} exec ${SINGULARITY_COMMAND_OPTS} {% if features.gpu %}{{ features.gpu }} {% endif %}{% if features.home %}-B {{ features.home }} --home {{ features.home }} {% endif %}{% if features.x11 %}-B {{ features.x11 }} {% endif %}{% if envfile %}-B {{ module_dir }}/{{ envfile }}:/.singularity.d/env/{{ envfile }}{% endif %} {% if bindpaths %}-B {{ bindpaths }}{% endif %} {% if alias.singularity_options %} {{ alias.singularity_options }} {% endif %} {{ container_sif }} {{ alias.command }} {% if '/sh' in wrapper_shell or '/bash' in wrapper_shell %}"$@"{% elif '/csh' in wrapper_shell %}$argv:q{% endif %} \ No newline at end of file diff --git a/shpc/main/modules/templates/singularity.tcl b/shpc/main/modules/templates/singularity.tcl index 50314efe6..a69868212 100644 --- a/shpc/main/modules/templates/singularity.tcl +++ b/shpc/main/modules/templates/singularity.tcl @@ -78,9 +78,9 @@ set inspectCmd "singularity \${SINGULARITY_OPTS} inspect \${SINGULARITY_COMMAND_ # set_shell_function takes bashStr and cshStr set-alias {|module_name|}-shell "${shellCmd}" -# exec functions to provide "alias" to module commands -{% if aliases %} -if { [ module-info shell bash ] } { +# "aliases" to module commands +{% if wrapper_scripts %}{% if aliases %}prepend-path PATH "[file dirname ${ModulesCurrentModulefile}]/{{ wrapper_subdir }}"{% endif %} +{% else %}{% if aliases %}if { [ module-info shell bash ] } { if { [ module-info mode load ] } { {% for alias in aliases %} puts stdout "function {{ alias.name }}() { ${execCmd} {% if alias.singularity_options %} {{ alias.singularity_options | replace("$", "\$") }} {% endif %} ${containerPath} {{ alias.command | replace("$", "\$") }} \"\$@\"; }; export -f {{ alias.name }};" {% endfor %} @@ -92,8 +92,7 @@ if { [ module-info shell bash ] } { } else { {% for alias in aliases %} set-alias {{ alias.name }} "${execCmd} {% if alias.singularity_options %} {{ alias.singularity_options | replace("$", "\$") }} {% endif %} ${containerPath} {{ alias.command | replace("$", "\$") }}" {% endfor %} -} -{% endif %} +}{% endif %}{% endif %} # A customizable exec function if { [ module-info shell bash ] } { diff --git a/shpc/main/schemas.py b/shpc/main/schemas.py index 8810c64ec..5deaa0efc 100644 --- a/shpc/main/schemas.py +++ b/shpc/main/schemas.py @@ -125,11 +125,13 @@ "environment_file": {"type": "string"}, "default_version": {"type": "boolean"}, "enable_tty": {"type": "boolean"}, + "wrapper_scripts": {"type": "boolean"}, "container_tech": {"type": "string", "enum": ["singularity", "podman", "docker"]}, "singularity_shell": {"type": "string", "enum": shells}, "podman_shell": {"type": "string", "enum": shells}, "docker_shell": {"type": "string", "enum": shells}, "test_shell": {"type": "string", "enum": shells}, + "wrapper_shell": {"type": "string", "enum": shells}, "module_sys": {"type": "string", "enum": ["lmod", "tcl", None]}, "container_features": container_features, } diff --git a/shpc/settings.yml b/shpc/settings.yml index ce9687f16..e40d7b371 100644 --- a/shpc/settings.yml +++ b/shpc/settings.yml @@ -51,6 +51,12 @@ docker_shell: /bin/sh # shell for test.sh file test_shell: /bin/bash +# shell for wrapper shellscripts +wrapper_shell: /bin/bash + +# for container aliases, use wrapper shellscripts instead of shell aliases, defaults to false +wrapper_scripts: false + # the module namespace you want to install from. E.g., if you have ghcr.io/autamus/clingo # and you set the namespace to ghcr.io/autamus, you can just do: shpc install clingo. namespace: diff --git a/shpc/utils/fileio.py b/shpc/utils/fileio.py index 8c854efdf..b8c5bb30f 100644 --- a/shpc/utils/fileio.py +++ b/shpc/utils/fileio.py @@ -5,6 +5,7 @@ import hashlib import errno import os +import stat import re import shutil import tempfile @@ -114,12 +115,15 @@ def copyfile(source, destination, force=True): return destination -def write_file(filename, content, mode="w"): +def write_file(filename, content, mode="w", exec=False): """ Write content to a filename """ with open(filename, mode) as filey: filey.writelines(content) + if exec: + st = os.stat(filename) + os.chmod(filename, st.st_mode | stat.S_IEXEC) return filename