Skip to content

Commit

Permalink
feat(Version Check): Added Singularity version check (#40)
Browse files Browse the repository at this point in the history
* Added a singularity version check

* Added apptainer version check

* fix(Apptainer): Updated pull and click option info

* feat(apptainer): Updated cli, utils for apptainer

---------

Co-authored-by: GitHub Actions <[email protected]>
  • Loading branch information
gaiborjosue and actions-user authored Oct 23, 2023
1 parent 175052f commit c6230e9
Show file tree
Hide file tree
Showing 2 changed files with 200 additions and 23 deletions.
180 changes: 158 additions & 22 deletions nobrainerzoo/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import subprocess as sp
import click
import yaml
import re

from .utils import (
CACHE_PATH,
Expand All @@ -14,6 +15,7 @@
get_repo,
get_spec,
pull_singularity_image,
pull_apptainer_image,
_get_model_file,
)

Expand Down Expand Up @@ -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():
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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",
Expand All @@ -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()
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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")
)
Expand Down Expand Up @@ -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",
Expand All @@ -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()
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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")
)
Expand Down Expand Up @@ -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",
Expand All @@ -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()
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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",
Expand All @@ -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:
Expand Down Expand Up @@ -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]
Expand All @@ -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"
)

Expand Down Expand Up @@ -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"]):
Expand All @@ -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."
)

Expand Down
43 changes: 42 additions & 1 deletion nobrainerzoo/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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(
Expand Down

0 comments on commit c6230e9

Please sign in to comment.