diff --git a/ersilia/cli/commands/fetch.py b/ersilia/cli/commands/fetch.py index af731ec80..dbf68546a 100644 --- a/ersilia/cli/commands/fetch.py +++ b/ersilia/cli/commands/fetch.py @@ -68,6 +68,12 @@ def _fetch(mf, model_id): default=False, help="Force fetch from DockerHub", ) + @click.option( + "--version", + default=None, + type=click.STRING, + help="Version of the model to fetch, when fetching a model from DockerHub", + ) @click.option( "--from_s3", is_flag=True, default=False, help="Force fetch from AWS S3 bucket" ) @@ -101,6 +107,7 @@ def fetch( from_dir, from_github, from_dockerhub, + version, from_s3, from_hosted, hosted_url, @@ -125,6 +132,7 @@ def fetch( force_from_github=from_github, force_from_s3=from_s3, force_from_dockerhub=from_dockerhub, + img_version=version, force_from_hosted=from_hosted, force_with_bentoml=with_bentoml, force_with_fastapi=with_fastapi, diff --git a/ersilia/core/modelbase.py b/ersilia/core/modelbase.py index 2c1c49c9e..07ac00254 100644 --- a/ersilia/core/modelbase.py +++ b/ersilia/core/modelbase.py @@ -2,7 +2,7 @@ import os from .. import ErsiliaBase, throw_ersilia_exception -from ..default import IS_FETCHED_FROM_DOCKERHUB_FILE +from ..default import DOCKER_INFO_FILE from ..hub.content.slug import Slug from ..hub.fetch import DONE_TAG, STATUS_FILE from ..utils.exceptions_utils.exceptions import InvalidModelIdentifierError @@ -104,7 +104,7 @@ def _is_available_locally_from_status(self): def _is_available_locally_from_dockerhub(self): from_dockerhub_file = os.path.join( - self._dest_dir, self.model_id, IS_FETCHED_FROM_DOCKERHUB_FILE + self._dest_dir, self.model_id, DOCKER_INFO_FILE ) if not os.path.exists(from_dockerhub_file): return False @@ -138,7 +138,7 @@ def was_fetched_from_dockerhub(self): True if the model was fetched from DockerHub, False otherwise. """ from_dockerhub_file = os.path.join( - self._dest_dir, self.model_id, IS_FETCHED_FROM_DOCKERHUB_FILE + self._dest_dir, self.model_id, DOCKER_INFO_FILE ) if not os.path.exists(from_dockerhub_file): return False diff --git a/ersilia/db/environments/managers.py b/ersilia/db/environments/managers.py index 592a9e07c..5a5b75c8b 100644 --- a/ersilia/db/environments/managers.py +++ b/ersilia/db/environments/managers.py @@ -6,7 +6,7 @@ from ...core.base import ErsiliaBase from ...default import DOCKERHUB_LATEST_TAG, DOCKERHUB_ORG from ...setup.requirements.docker import DockerRequirement -from ...utils.docker import SimpleDocker, resolve_platform +from ...utils.docker import SimpleDocker, model_image_version_reader, resolve_platform from ...utils.identifiers.short import ShortIdentifier from ...utils.logging import make_temp_dir from ...utils.paths import Paths @@ -235,7 +235,7 @@ def containers_of_model(self, model_id, only_run, only_latest=True): cnt_dict[k] = v return cnt_dict - def build_with_bentoml(self, model_id, use_cache=True): + def build_with_bentoml(self, model_id, use_cache=True): # Ignore for versioning """ Builds a Docker image for the model using BentoML. @@ -295,7 +295,9 @@ def _build_ersilia_base(self): ) run_command(cmd) - def build_with_ersilia(self, model_id, docker_user, docker_pwd): + def build_with_ersilia( + self, model_id, docker_user, docker_pwd + ): # Ignore for versioning """ Builds a Docker image for the model using Ersilia's base image. @@ -383,10 +385,12 @@ def remove(self, model_id): model_id : str Identifier of the model. """ - self.docker.delete(org=DOCKERHUB_ORG, img=model_id, tag=DOCKERHUB_LATEST_TAG) + bundle_path = self._get_bundle_location(model_id) + docker_tag = model_image_version_reader(bundle_path) + self.docker.delete(org=DOCKERHUB_ORG, img=model_id, tag=docker_tag) self.db.delete( model_id=model_id, - env="{0}/{1}:{2}".format(DOCKERHUB_ORG, model_id, DOCKERHUB_LATEST_TAG), + env="{0}/{1}:{2}".format(DOCKERHUB_ORG, model_id, docker_tag), ) def run(self, model_id, workers=1, enable_microbatch=True, memory=None): diff --git a/ersilia/default.py b/ersilia/default.py index bcf0fe2fd..23131762c 100644 --- a/ersilia/default.py +++ b/ersilia/default.py @@ -57,7 +57,7 @@ MODEL_SOURCE_FILE = "model_source.txt" APIS_LIST_FILE = "apis_list.txt" INFORMATION_FILE = "information.json" -IS_FETCHED_FROM_DOCKERHUB_FILE = "from_dockerhub.json" +DOCKER_INFO_FILE = "from_dockerhub.json" IS_FETCHED_FROM_HOSTED_FILE = "from_hosted.json" DEFAULT_UDOCKER_USERNAME = "udockerusername" DEFAULT_UDOCKER_PASSWORD = "udockerpassword" diff --git a/ersilia/hub/bundle/status.py b/ersilia/hub/bundle/status.py index 143562c83..3d61bdfcc 100644 --- a/ersilia/hub/bundle/status.py +++ b/ersilia/hub/bundle/status.py @@ -4,7 +4,7 @@ from ... import ErsiliaBase from ...db.environments.localdb import EnvironmentDb -from ...default import IS_FETCHED_FROM_DOCKERHUB_FILE +from ...default import DOCKER_INFO_FILE from ...utils.conda import SimpleConda from ...utils.docker import SimpleDocker @@ -89,7 +89,7 @@ def is_pulled_docker(self, model_id: str) -> bool: True if the Docker image has been pulled, False otherwise. """ model_dir = os.path.join(self._model_path(model_id=model_id)) - json_file = os.path.join(model_dir, IS_FETCHED_FROM_DOCKERHUB_FILE) + json_file = os.path.join(model_dir, DOCKER_INFO_FILE) if not os.path.exists(json_file): return False with open(json_file, "r") as f: diff --git a/ersilia/hub/fetch/fetch.py b/ersilia/hub/fetch/fetch.py index 3e750b546..220355fe8 100644 --- a/ersilia/hub/fetch/fetch.py +++ b/ersilia/hub/fetch/fetch.py @@ -50,6 +50,8 @@ class ModelFetcher(ErsiliaBase): Whether to force fetching from S3. force_from_dockerhub : bool, optional Whether to force fetching from DockerHub. + img_version : str, optional + Version of the model image. force_from_hosted : bool, optional Whether to force fetching from hosted services. force_with_bentoml : bool, optional @@ -81,6 +83,7 @@ def __init__( force_from_github: bool = False, force_from_s3: bool = False, force_from_dockerhub: bool = False, + img_version: str = None, force_from_hosted: bool = False, force_with_bentoml: bool = False, force_with_fastapi: bool = False, @@ -100,7 +103,7 @@ def __init__( dockerize = True self.do_docker = dockerize self.model_dockerhub_fetcher = ModelDockerHubFetcher( - overwrite=self.overwrite, config_json=self.config_json + overwrite=self.overwrite, config_json=self.config_json, img_tag=img_version ) self.is_docker_installed = self.model_dockerhub_fetcher.is_docker_installed() self.is_docker_active = self.model_dockerhub_fetcher.is_docker_active() diff --git a/ersilia/hub/fetch/lazy_fetchers/dockerhub.py b/ersilia/hub/fetch/lazy_fetchers/dockerhub.py index 73839facf..a99000f4e 100644 --- a/ersilia/hub/fetch/lazy_fetchers/dockerhub.py +++ b/ersilia/hub/fetch/lazy_fetchers/dockerhub.py @@ -59,10 +59,12 @@ class ModelDockerHubFetcher(ErsiliaBase): Fetch the model from DockerHub. """ - def __init__(self, overwrite=None, config_json=None): + def __init__(self, overwrite=None, config_json=None, img_tag=None): super().__init__(config_json=config_json, credentials_json=None) self.simple_docker = SimpleDocker() self.overwrite = overwrite + self.img_tag = img_tag or DOCKERHUB_LATEST_TAG + self.pack_method = None def is_docker_installed(self) -> bool: """ @@ -101,7 +103,10 @@ def is_available(self, model_id: str) -> bool: True if the model is available, False otherwise. """ mp = ModelPuller( - model_id=model_id, overwrite=self.overwrite, config_json=self.config_json + model_id=model_id, + overwrite=self.overwrite, + config_json=self.config_json, + docker_tag=self.img_tag, ) if mp.is_available_locally(): return True @@ -144,7 +149,7 @@ async def _copy_from_bentoml_image(self, model_id: str, file: str): local_path=to_file, org=DOCKERHUB_ORG, img=model_id, - tag=DOCKERHUB_LATEST_TAG, + tag=self.img_tag, ) except Exception as e: self.logger.error(f"Exception when copying: {e}") @@ -167,7 +172,7 @@ async def _copy_from_ersiliapack_image(self, model_id: str, file: str): local_path=to_file, org=DOCKERHUB_ORG, img=model_id, - tag=DOCKERHUB_LATEST_TAG, + tag=self.img_tag, ) async def _copy_from_image_to_local(self, model_id: str, file: str): @@ -181,8 +186,12 @@ async def _copy_from_image_to_local(self, model_id: str, file: str): file : str Name of the file to copy. """ - pack_method = resolve_pack_method_docker(model_id) - if pack_method == PACK_METHOD_BENTOML: + if not self.pack_method: + self.logger.debug("Resolving pack method") + self.pack_method = resolve_pack_method_docker(model_id) + self.logger.debug(f"Resolved pack method: {self.pack_method}") + + if self.pack_method == PACK_METHOD_BENTOML: await self._copy_from_bentoml_image(model_id, file) else: await self._copy_from_ersiliapack_image(model_id, file) @@ -245,7 +254,9 @@ async def modify_information(self, model_id: str): ID of the model. """ information_file = os.path.join(self._model_path(model_id), INFORMATION_FILE) - mp = ModelPuller(model_id=model_id, config_json=self.config_json) + mp = ModelPuller( + model_id=model_id, config_json=self.config_json, docker_tag=self.img_tag + ) try: with open(information_file, "r") as infile: data = json.load(infile) @@ -268,13 +279,15 @@ async def fetch(self, model_id: str): model_id : str ID of the model. """ - mp = ModelPuller(model_id=model_id, config_json=self.config_json) + mp = ModelPuller( + model_id=model_id, config_json=self.config_json, docker_tag=self.img_tag + ) self.logger.debug("Pulling model image from DockerHub") await mp.async_pull() mr = ModelRegisterer(model_id=model_id, config_json=self.config_json) self.logger.debug("Asynchronous and concurrent execution started!") await asyncio.gather( - mr.register(is_from_dockerhub=True), + mr.register(is_from_dockerhub=True, img_tag=self.img_tag), self.write_apis(model_id), self.copy_information(model_id), self.modify_information(model_id), diff --git a/ersilia/hub/fetch/register/register.py b/ersilia/hub/fetch/register/register.py index d6b0b4cf5..d9a578f92 100644 --- a/ersilia/hub/fetch/register/register.py +++ b/ersilia/hub/fetch/register/register.py @@ -7,7 +7,7 @@ from .... import EOS, ErsiliaBase, throw_ersilia_exception from ....default import ( - IS_FETCHED_FROM_DOCKERHUB_FILE, + DOCKER_INFO_FILE, IS_FETCHED_FROM_HOSTED_FILE, SERVICE_CLASS_FILE, ) @@ -40,13 +40,17 @@ def __init__(self, model_id: str, config_json: dict): ErsiliaBase.__init__(self, config_json=config_json, credentials_json=None) self.model_id = model_id - def register_from_dockerhub(self): + def register_from_dockerhub(self, **kwargs): """ Register the model from DockerHub. This method registers the model in the file system indicating it was fetched from DockerHub. """ - data = {"docker_hub": True} + if 'img_tag' in kwargs: + img_tag = kwargs['img_tag'] + data = {"docker_hub": True, "tag": img_tag} + else: + data = {"docker_hub": True} self.logger.debug( "Registering model {0} in the file system".format(self.model_id) ) @@ -55,7 +59,7 @@ def register_from_dockerhub(self): if os.path.exists(path): shutil.rmtree(path) os.mkdir(path) - file_name = os.path.join(path, IS_FETCHED_FROM_DOCKERHUB_FILE) + file_name = os.path.join(path, DOCKER_INFO_FILE) self.logger.debug(file_name) with open(file_name, "w") as f: json.dump(data, f) @@ -66,7 +70,7 @@ def register_from_dockerhub(self): shutil.rmtree(path) path = os.path.join(path, folder_name) os.makedirs(path) - file_name = os.path.join(path, IS_FETCHED_FROM_DOCKERHUB_FILE) + file_name = os.path.join(path, DOCKER_INFO_FILE) with open(file_name, "w") as f: json.dump(data, f) file_name = os.path.join(path, SERVICE_CLASS_FILE) @@ -82,11 +86,11 @@ def register_not_from_dockerhub(self): """ data = {"docker_hub": False} path = self._model_path(self.model_id) - file_name = os.path.join(path, IS_FETCHED_FROM_DOCKERHUB_FILE) + file_name = os.path.join(path, DOCKER_INFO_FILE) with open(file_name, "w") as f: json.dump(data, f) path = self._get_bundle_location(model_id=self.model_id) - file_name = os.path.join(path, IS_FETCHED_FROM_DOCKERHUB_FILE) + file_name = os.path.join(path, DOCKER_INFO_FILE) with open(file_name, "w") as f: json.dump(data, f) @@ -174,7 +178,7 @@ def register_not_from_hosted(self): json.dump(data, f) async def register( - self, is_from_dockerhub: bool = False, is_from_hosted: bool = False + self, is_from_dockerhub: bool = False, is_from_hosted: bool = False, **kwargs ): """ Register the model based on its source. @@ -205,7 +209,7 @@ async def register( if is_from_dockerhub and is_from_hosted: raise ValueError("Model cannot be from both DockerHub and hosted") elif is_from_dockerhub and not is_from_hosted: - self.register_from_dockerhub() + self.register_from_dockerhub(**kwargs) self.register_not_from_hosted() elif not is_from_dockerhub and is_from_hosted: self.register_from_hosted() diff --git a/ersilia/hub/pull/pull.py b/ersilia/hub/pull/pull.py index 38c326ea2..8a906b0a2 100644 --- a/ersilia/hub/pull/pull.py +++ b/ersilia/hub/pull/pull.py @@ -41,12 +41,19 @@ class ModelPuller(ErsiliaBase): await puller.async_pull() """ - def __init__(self, model_id: str, overwrite: bool = None, config_json: dict = None): + def __init__( + self, + model_id: str, + overwrite: bool = None, + config_json: dict = None, + docker_tag: str = None, + ): ErsiliaBase.__init__(self, config_json=config_json, credentials_json=None) self.simple_docker = SimpleDocker() self.model_id = model_id + self.docker_tag = docker_tag or DOCKERHUB_LATEST_TAG self.image_name = "{0}/{1}:{2}".format( - DOCKERHUB_ORG, self.model_id, DOCKERHUB_LATEST_TAG + DOCKERHUB_ORG, self.model_id, self.docker_tag ) self.overwrite = overwrite @@ -60,7 +67,7 @@ def is_available_locally(self) -> bool: True if the image is available locally, False otherwise. """ is_available = self.simple_docker.exists( - DOCKERHUB_ORG, self.model_id, DOCKERHUB_LATEST_TAG + DOCKERHUB_ORG, self.model_id, self.docker_tag ) if is_available: self.logger.debug("Image {0} is available locally".format(self.image_name)) @@ -81,7 +88,7 @@ def is_available_in_dockerhub(self) -> bool: True if the image is available in DockerHub, False otherwise. """ url = "https://hub.docker.com/v2/repositories/{0}/{1}/tags/{2}".format( - DOCKERHUB_ORG, self.model_id, DOCKERHUB_LATEST_TAG + DOCKERHUB_ORG, self.model_id, self.docker_tag ) response = requests.get(url) if response.status_code == 200: @@ -102,13 +109,13 @@ def _delete(self): "Deleting locally available image {0}".format(self.image_name) ) self.simple_docker.delete( - org=DOCKERHUB_ORG, img=self.model_id, tag=DOCKERHUB_LATEST_TAG + org=DOCKERHUB_ORG, img=self.model_id, tag=self.docker_tag ) def _get_size_of_local_docker_image_in_mb(self) -> float: try: image_name = "{0}/{1}:{2}".format( - DOCKERHUB_ORG, self.model_id, DOCKERHUB_LATEST_TAG + DOCKERHUB_ORG, self.model_id, self.docker_tag ) result = subprocess.check_output( ["docker", "image", "inspect", image_name, "--format", "{{.Size}}"] @@ -152,7 +159,9 @@ async def async_pull(self): "Trying to pull image {0}/{1}".format(DOCKERHUB_ORG, self.model_id) ) - pull_command = f"docker pull {DOCKERHUB_ORG}/{self.model_id}:{DOCKERHUB_LATEST_TAG}" + pull_command = ( + f"docker pull {DOCKERHUB_ORG}/{self.model_id}:{self.docker_tag}" + ) process = await asyncio.create_subprocess_shell( pull_command, @@ -185,7 +194,7 @@ async def log_stream(stream, log_method): self.logger.warning( "Conventional pull did not work, Ersilia is now forcing linux/amd64 architecture" ) - force_pull_command = f"docker pull {DOCKERHUB_ORG}/{self.model_id}:{DOCKERHUB_LATEST_TAG} --platform linux/amd64" + force_pull_command = f"docker pull {DOCKERHUB_ORG}/{self.model_id}:{self.docker_tag} --platform linux/amd64" process = await asyncio.create_subprocess_shell( force_pull_command, @@ -257,7 +266,7 @@ def pull(self): self.logger.debug("Keeping logs of pull in {0}".format(tmp_file)) run_command( "docker pull {0}/{1}:{2} > {3} 2>&1".format( - DOCKERHUB_ORG, self.model_id, DOCKERHUB_LATEST_TAG, tmp_file + DOCKERHUB_ORG, self.model_id, self.docker_tag, tmp_file ) ) with open(tmp_file, "r") as f: @@ -275,7 +284,7 @@ def pull(self): ) run_command( "docker pull {0}/{1}:{2} --platform linux/amd64".format( - DOCKERHUB_ORG, self.model_id, DOCKERHUB_LATEST_TAG + DOCKERHUB_ORG, self.model_id, self.docker_tag ) ) size = self._get_size_of_local_docker_image_in_mb() diff --git a/ersilia/serve/autoservice.py b/ersilia/serve/autoservice.py index 2470faeb0..c5a6a5667 100644 --- a/ersilia/serve/autoservice.py +++ b/ersilia/serve/autoservice.py @@ -8,7 +8,7 @@ from ..default import ( APIS_LIST_FILE, DEFAULT_BATCH_SIZE, - IS_FETCHED_FROM_DOCKERHUB_FILE, + DOCKER_INFO_FILE, IS_FETCHED_FROM_HOSTED_FILE, SERVICE_CLASS_FILE, ) @@ -217,7 +217,7 @@ def __init__( def _was_fetched_from_dockerhub(self): from_dockerhub_file = os.path.join( - self._dest_dir, self.model_id, IS_FETCHED_FROM_DOCKERHUB_FILE + self._dest_dir, self.model_id, DOCKER_INFO_FILE ) if not os.path.exists(from_dockerhub_file): return False diff --git a/ersilia/serve/services.py b/ersilia/serve/services.py index 7efa562d7..a53abb68f 100644 --- a/ersilia/serve/services.py +++ b/ersilia/serve/services.py @@ -14,7 +14,6 @@ APIS_LIST_FILE, CONTAINER_LOGS_TMP_DIR, DEFAULT_VENV, - DOCKERHUB_LATEST_TAG, DOCKERHUB_ORG, INFORMATION_FILE, IS_FETCHED_FROM_HOSTED_FILE, @@ -25,7 +24,7 @@ from ..setup.requirements.conda import CondaRequirement from ..setup.requirements.docker import DockerRequirement from ..utils.conda import SimpleConda, StandaloneConda -from ..utils.docker import SimpleDocker +from ..utils.docker import SimpleDocker, model_image_version_reader from ..utils.exceptions_utils.serve_exceptions import ( BadGatewayError, DockerNotActiveError, @@ -1120,9 +1119,11 @@ def __init__(self, model_id, config_json=None, preferred_port=None, url=None): self.port = find_free_port() else: self.port = preferred_port + bundle_path = self._model_path(model_id) + self.docker_tag = model_image_version_reader(bundle_path) self.logger.debug("Using port {0}".format(self.port)) self.image_name = "{0}/{1}:{2}".format( - DOCKERHUB_ORG, self.model_id, DOCKERHUB_LATEST_TAG + DOCKERHUB_ORG, self.model_id, self.docker_tag ) self.logger.debug("Starting Docker Daemon service") self.container_tmp_logs = os.path.join( @@ -1205,7 +1206,7 @@ def is_available(self) -> bool: True if the service is available, False otherwise. """ is_available = self.simple_docker.exists( - DOCKERHUB_ORG, self.model_id, DOCKERHUB_LATEST_TAG + DOCKERHUB_ORG, self.model_id, self.docker_tag ) if is_available: self.logger.debug("Image {0} is available locally".format(self.image_name)) diff --git a/ersilia/utils/docker.py b/ersilia/utils/docker.py index 13e3a09b7..c4a8a8490 100644 --- a/ersilia/utils/docker.py +++ b/ersilia/utils/docker.py @@ -1,3 +1,4 @@ +import json import os import subprocess import threading @@ -10,8 +11,10 @@ from ..default import ( DEFAULT_DOCKER_PLATFORM, DEFAULT_UDOCKER_USERNAME, + DOCKER_INFO_FILE, DOCKERHUB_LATEST_TAG, DOCKERHUB_ORG, + EOS, PACK_METHOD_BENTOML, PACK_METHOD_FASTAPI, ) @@ -23,9 +26,9 @@ def resolve_pack_method_docker(model_id): client = docker.from_env() - model_image = client.images.get( - f"{DOCKERHUB_ORG}/{model_id}:{DOCKERHUB_LATEST_TAG}" - ) + bundle_path = f"{EOS}/dest/{model_id}" + docker_tag = model_image_version_reader(bundle_path) + model_image = client.images.get(f"{DOCKERHUB_ORG}/{model_id}:{docker_tag}") image_history = model_image.history() for hist in image_history: # Very hacky, but works bec we don't have nginx in ersilia-pack images @@ -62,6 +65,18 @@ def is_udocker_installed(): return False +def model_image_version_reader(dir): + """ + Read the requested model image version from a file. + """ + if os.path.exists(os.path.join(dir, DOCKER_INFO_FILE)): + with open(os.path.join(dir, DOCKER_INFO_FILE), "r") as f: + data = json.load(f) + if "tag" in data: + return data["tag"] + return DOCKERHUB_LATEST_TAG + + class SimpleDocker(object): """ A class to manage Docker containers and images.