diff --git a/nobrainerzoo/cli.py b/nobrainerzoo/cli.py index 2cbc0be..96f60ae 100755 --- a/nobrainerzoo/cli.py +++ b/nobrainerzoo/cli.py @@ -4,6 +4,7 @@ import subprocess as sp import click import yaml +import re from .utils import ( CACHE_PATH, @@ -14,6 +15,7 @@ get_repo, get_spec, pull_singularity_image, + pull_apptainer_image, _get_model_file, ) @@ -109,7 +111,33 @@ def init(cache): os.makedirs(DATA_PATH, exist_ok=True) # adding trained_model repository model_db_url = "https://github.com/neuronets/trained-models" - if _container_installed("singularity"): + if _container_installed("apptainer"): + # pull the nobrainer image from docker-hub + download_image = IMAGES_PATH / "nobrainer-zoo_zoo.sif" + if not download_image.exists(): + dwnld_cmd = [ + "apptainer", + "pull", + "--dir", + str(IMAGES_PATH), + "docker://neuronets/nobrainer-zoo:zoo", + ] + p0 = sp.run(dwnld_cmd, stdout=sp.PIPE, stderr=sp.STDOUT, text=True) + print(p0.stdout) + # For a robust behavior of model_db, we should clone via datalad. + clone_cmd = [ + "apptainer", + "run", + "-e", + "-B", + CACHE_PATH, + download_image, + "datalad", + "clone", + model_db_url, + MODELS_PATH, + ] + elif _container_installed("singularity"): # pull the nobrainer image from docker-hub download_image = IMAGES_PATH / "nobrainer-zoo_zoo.sif" if not download_image.exists(): @@ -230,7 +258,7 @@ def ls(model, model_type): "--container_type", default="singularity", type=str, - help="Type of the container technology (docker or singularity)", + help="Type of the container technology (docker, singularity, or apptainer)", **_option_kwds, ) @click.option( @@ -329,7 +357,26 @@ def predict(infile, outfile, model, model_type, container_type, cpu, options, ** except NameError: model_cmd = spec["command"] # breakpoint() - if container_type == "singularity": + if container_type == "apptainer": + bind_paths = ",".join(bind_paths) + cmd_options = [ + #"-e", + "--nv", + "-B", + bind_paths, + "-B", + f"{out_path}:/output", + "-W", + "/output", + ] + cmd = ( + ["apptainer", "exec"] + + cmd_options + + [image] + + model_cmd.split() + + model_options + ) + elif container_type == "singularity": bind_paths = ",".join(bind_paths) cmd_options = [ # "-e", @@ -342,7 +389,7 @@ def predict(infile, outfile, model, model_type, container_type, cpu, options, ** "/output", ] cmd = ( - ["singularity", "run"] + ["singularity", "exec"] + cmd_options + [image] + model_cmd.split() @@ -403,7 +450,7 @@ def predict(infile, outfile, model, model_type, container_type, cpu, options, ** "--container_type", default="singularity", type=str, - help="Type of the container technology (docker or singularity)", + help="Type of the container technology (docker, singularity or apptainer)", **_option_kwds, ) @click.option( @@ -436,7 +483,7 @@ def generate(outfile, model, model_type, container_type, cpu, options, **kwrg): spec = get_spec(model, model_type) - # download the model-required docker/singularity image and set the path + # download the model-required docker/singularity/apptainer image and set the path image = _container_check( container_type=container_type, image_spec=spec.get("image") ) @@ -498,7 +545,26 @@ def generate(outfile, model, model_type, container_type, cpu, options, **kwrg): except NameError: model_cmd = spec["command"] - if container_type == "singularity": + if container_type == "apptainer": + bind_paths = ",".join(bind_paths) + cmd_options = [ + #"-e", + "--nv", + "-B", + bind_paths, + "-B", + f"{out_path}:/output", + "-W", + "/output", + ] + cmd = ( + ["apptainer", "exec"] + + cmd_options + + [image] + + model_cmd.split() + + model_options + ) + elif container_type == "singularity": bind_paths = ",".join(bind_paths) cmd_options = [ "-e", @@ -511,7 +577,7 @@ def generate(outfile, model, model_type, container_type, cpu, options, **kwrg): "/output", ] cmd = ( - ["singularity", "run"] + ["singularity", "exec"] + cmd_options + [image] + model_cmd.split() @@ -574,7 +640,7 @@ def generate(outfile, model, model_type, container_type, cpu, options, **kwrg): "--container_type", default="singularity", type=str, - help="Type of the container technology (docker or singularity)", + help="Type of the container technology (docker, singularity or apptainer)", **_option_kwds, ) @click.option( @@ -611,7 +677,7 @@ def register(moving, fixed, moved, model, model_type, container_type, cpu, optio spec = get_spec(model, model_type) - # set the docker/singularity image + # set the docker/singularity/apptainer image image = _container_check( container_type=container_type, image_spec=spec.get("image") ) @@ -675,7 +741,27 @@ def register(moving, fixed, moved, model, model_type, container_type, cpu, optio except NameError: model_cmd = spec["command"] - if container_type == "singularity": + if container_type == "apptainer": + bind_paths = ",".join(bind_paths) + cmd_options = [ + #"-e", + "--nv", + "-B", + bind_paths, + "-B", + f"{out_path}:/output", + "-W", + "/output", + ] + cmd = ( + ["apptainer", "exec"] + + cmd_options + + [image] + + model_cmd.split() + + model_options + ) + + elif container_type == "singularity": bind_paths = ",".join(bind_paths) cmd_options = [ "-e", @@ -688,7 +774,7 @@ def register(moving, fixed, moved, model, model_type, container_type, cpu, optio "/output", ] cmd = ( - ["singularity", "run"] + ["singularity", "exec"] + cmd_options + [image] + model_cmd.split() @@ -749,7 +835,7 @@ def register(moving, fixed, moved, model, model_type, container_type, cpu, optio "--container_type", default="singularity", type=str, - help="Type of the container technology (docker or singularity)", + help="Type of the container technology (docker, singularity or apptainer)", **_option_kwds, ) @click.option( @@ -815,7 +901,7 @@ def fit( """ - # set the docker/singularity image + # set the docker/singularity/apptainer image org, model_nm, ver = model.split("/") if spec_file: @@ -891,7 +977,22 @@ def fit( image = _container_check( container_type=container_type, image_spec=spec.get("image") ) - if container_type == "singularity": + + if container_type == "apptainer": + bind_paths = ",".join(bind_paths) + cmd_options = [ + #"-e", + "--nv", + "-B", + bind_paths, + "-B", + f"{out_path}:/output", + "-W", + "/output", + ] + cmd_cont = ["apptainer", "exec"] + options + [image] + cmd + + elif container_type == "singularity": bind_paths = ",".join(bind_paths) options = [ "-e", @@ -903,7 +1004,7 @@ def fit( "-W", "/output", ] - cmd_cont = ["singularity", "run"] + options + [image] + cmd + cmd_cont = ["singularity", "exec"] + options + [image] + cmd elif container_type == "docker": bind_paths_docker = [] for el in bind_paths: @@ -931,11 +1032,22 @@ def _container_check(container_type, image_spec): if image_spec is None: raise Exception("image not provided in the specification") - # check the installation of singularity or docker + # check the installation of singularity or docker or apptainer if not _container_installed(container_type): raise Exception(f"{container_type} is not installed.") - if container_type == "singularity": + if container_type == "apptainer": + if container_type in image_spec: + pull_apptainer_image(image_spec[container_type], IMAGES_PATH) + image = IMAGES_PATH / image_spec[container_type] + else: + raise Exception( + f"container name for {container_type} is not " + f"provided in the specification, " + f"try using container_type=singularity" + ) + + elif container_type == "singularity": if container_type in image_spec: pull_singularity_image(image_spec[container_type], IMAGES_PATH) image = IMAGES_PATH / image_spec[container_type] @@ -957,7 +1069,7 @@ def _container_check(container_type, image_spec): else: raise Exception( - f"container_type should be docker or singularity, " + f"container_type should be docker, singularity or apptainer, " f"but {container_type} provided" ) @@ -991,12 +1103,36 @@ def _name(**variables): def _container_installed(container_type): """checks singularity or docker is installed.""" - + # Check for singularity installation and version if container_type == "singularity": if shutil.which("singularity") is None: return False else: - return True + try: + # Checks singularity version and apptainer + singularity_version_output = sp.check_output(["singularity", "--version"], text=True) + singularity_version = re.search(r"\d+\.\d+\.\d+", singularity_version_output).group() + + # When singularity and apptainer are loaded in the same env, doing singularity --version will return apptainer's version. Therefore, here I check for both scenarios, where apptainer & singularity are installed or only singularity. + if singularity_version >= "3.7" or (shutil.which("apptainer") is not None and singularity_version >= "1.0"): + return True + else: + return False + except sp.CalledProcessError: + return False + elif container_type == "apptainer": + if shutil.which("apptainer") is None: + return False + else: + try: + app_tainer_version_output = sp.check_output(["apptainer", "--version"], text=True) + app_tainer_version = re.search(r"\d+\.\d+\.\d+", app_tainer_version_output).group() + if app_tainer_version > "1.0": + return True + else: + return False + except sp.CalledProcessError: + return False elif container_type == "docker": if shutil.which("docker") is None or sp.call(["docker", "info"]): @@ -1005,7 +1141,7 @@ def _container_installed(container_type): return True else: raise Exception( - f"container_type should be docker or singularity, " + f"container_type should be docker, singularity or apptainer, " f"but {container_type} is provided." ) diff --git a/nobrainerzoo/utils.py b/nobrainerzoo/utils.py index ff19a23..93ad382 100755 --- a/nobrainerzoo/utils.py +++ b/nobrainerzoo/utils.py @@ -177,6 +177,21 @@ def pull_singularity_image(singularity_image, path): p0 = sp.run(dwnld_cmd, stdout=sp.PIPE, stderr=sp.STDOUT, text=True) print(p0.stdout) +def pull_apptainer_image(apptainer_image, path): + download_image = Path(path) /apptainer_image + if not download_image.exists(): + image_tag, _ = apptainer_image.split("_")[1].split(".") + print("Downloading the container file. it might take a while...") + dwnld_cmd = [ + "apptainer", + "pull", + "--dir", + str(path), + # container images are stored in dockerhub neuronets/nobrainer-zoo + f"docker://neuronets/nobrainer-zoo:{image_tag}", + ] + p0 = sp.run(dwnld_cmd, stdout=sp.PIPE, stderr=sp.STDOUT, text=True) + print(p0.stdout) def _check_model_type(model_name, model_type=None): models = get_model_db(MODELS_PATH, print_models=False) @@ -199,7 +214,33 @@ def _get_model_file(model_path, container_type, ): """downloads the model file.""" parent_dir = Path(__file__).resolve().parent loader = str(parent_dir / "download.py") - if container_type == "singularity": + if container_type == "apptainer": + download_image = IMAGES_PATH / "nobrainer-zoo_zoo.sif" + if not download_image.exists(): + raise Exception( + "'nobrainer-zoo' apptainer image is missing! ", + "Please run 'nobrainer-zoo init'.", + ) + + # mount CACHE_PATH to /cache_dir, I will be using that path in some functions + cmd0 = [ + "apptainer", + "run", + "-e", + "-B", + parent_dir, + "-B", + str(CACHE_PATH), + "-B", + f"{CACHE_PATH}:/cache_dir", + download_image, + "python3", + loader, + MODELS_PATH, + model_path, + ] + # str( parent_dir / "download.py"), "/cache_dir/trained-models", model] + elif container_type == "singularity": download_image = IMAGES_PATH / "nobrainer-zoo_zoo.sif" if not download_image.exists(): raise Exception(