Update backend [part 4]

modelforge/__main__.py

@@ -4,7 +4,6 @@
 import os
 
 from modelforge.logs import setup_logging
-from modelforge.index import DEFAULT_CACHE
 from modelforge.registry import publish_model, list_models, initialize_registry, dump_model, \
     delete_model
 
@@ -29,8 +28,9 @@ def add_index_args(p):
                        help="Username for the Git repository with the index.")
         p.add_argument("--password", default="",
                        help="Password for the Git repository with the index")
-        p.add_argument("--index-repo", default="", help="Url of the remote Git repository.")
-        p.add_argument("--cache", default=DEFAULT_CACHE,
+        p.add_argument("--index-repo", default=None,
+                       help="Url of the remote Git repository.")
+        p.add_argument("--cache", default=None,
                        help="Path to the folder where the Git repository will be cached.")
 
     def add_templates_args(p):
@@ -43,14 +43,22 @@ def add_templates_args(p):
             default=os.path.join(os.path.dirname(__file__), "templates/template_readme.md.jinja2"),
             help="Path to the jinja2 template used in the index for the readme.")
 
+    # ------------------------------------------------------------------------
+    init_parser = subparsers.add_parser("init", help="Initialize the registry.")
+    init_parser.set_defaults(handler=initialize_registry)
+    init_parser.add_argument("-f", "--force", action="store_true",
+                             help="Destructively initialize the registry.")
+    add_index_args(init_parser)
+    add_backend_args(init_parser)
+    # ------------------------------------------------------------------------
     dump_parser = subparsers.add_parser(
         "dump", help="Print a brief information about the model to stdout.")
     dump_parser.set_defaults(handler=dump_model)
     dump_parser.add_argument(
         "input", help="Path to the model file, URL or UUID.")
-    add_backend_args(dump_parser)
     add_index_args(dump_parser)
-
+    add_backend_args(dump_parser)
+    # ------------------------------------------------------------------------
     publish_parser = subparsers.add_parser(
         "publish", help="Upload the model and update the registry.")
     publish_parser.set_defaults(handler=publish_model)
@@ -63,30 +71,23 @@ def add_templates_args(p):
                                 help="Set this model as the default one.")
     publish_parser.add_argument("-f", "--force", action="store_true",
                                 help="Overwrite existing models.")
-    add_backend_args(publish_parser)
     add_index_args(publish_parser)
+    add_backend_args(publish_parser)
     add_templates_args(publish_parser)
-
+    # ------------------------------------------------------------------------
     list_parser = subparsers.add_parser(
         "list", help="Lists all the models in the registry.")
     list_parser.set_defaults(handler=list_models)
     add_index_args(list_parser)
-
-    init_parser = subparsers.add_parser("init", help="Initialize the registry.")
-    init_parser.set_defaults(handler=initialize_registry)
-    init_parser.add_argument("-f", "--force", action="store_true",
-                             help="Destructively initialize the registry.")
-    add_backend_args(init_parser)
-    add_index_args(init_parser)
-
+    # ------------------------------------------------------------------------
     delete_parser = subparsers.add_parser("delete", help="Delete a model.")
     delete_parser.set_defaults(handler=delete_model)
     delete_parser.add_argument(
         "input", help="UUID of the model to be deleted.")
-    add_backend_args(delete_parser)
     add_index_args(delete_parser)
+    add_backend_args(delete_parser)
     add_templates_args(delete_parser)
-
+    # ------------------------------------------------------------------------
     args = parser.parse_args()
     args.log_level = logging._nameToLevel[args.log_level]
     setup_logging(args.log_level)

modelforge/backends.py

@@ -52,9 +52,13 @@ def supply_backend_inner(func):
         @wraps(func)
         def wrapped_supply_backend(args):
             log = logging.getLogger(func.__name__)
-            git_index = GitIndex(index_repo=args.index_repo, username=args.username,
-                                 password=args.password, cache=args.cache,
-                                 log_level=args.log_level)
+            try:
+                git_index = GitIndex(index_repo=args.index_repo, username=args.username,
+                                     password=args.password, cache=args.cache,
+                                     log_level=args.log_level)
+            except ValueError:
+                return 1
+
             if real_optional and not getattr(args, "backend", False):
                 backend = None
             else:

modelforge/configuration.py

@@ -7,6 +7,7 @@
 BACKEND = os.getenv("MODELFORGE_BACKEND", None)
 BACKEND_ARGS = os.getenv("MODELFORGE_BACKEND_ARGS", "")
 DEFAULT_REPO = os.getenv("MODELFORGE_DEFAULT_REPO", "")
+DEFAULT_CACHE = os.getenv("MODELFORGE_DEFAULT_CACHE", os.path.expanduser("~/.cache"))
 
 OVERRIDE_FILE = "modelforgecfg.py"
 

modelforge/gcs_backend.py

@@ -1,24 +1,25 @@
 import io
-import json
 import logging
 import math
 import os
 import requests
 
+from typing import BinaryIO, Union
 from clint.textui import progress
 from google.cloud.exceptions import NotFound
 
 from modelforge.index import GitIndex
 from modelforge.progress_bar import progress_bar
-from modelforge.storage_backend import StorageBackend, TransactionRequiredError
+from modelforge.storage_backend import StorageBackend, ExistingBackendError, \
+    ModelAlreadyExistsError, BackendRequiredError
 
 
 class Tracker:
     """
     Wrapper around a bytes buffer which follows the file position and updates
     the console progressbar mimicking a file object.
     """
-    def __init__(self, data, logger):
+    def __init__(self, data: memoryview, logger: logging.Logger):
         self._file = io.BytesIO(data)
         self._size = len(data)
         self._enabled = logger.isEnabledFor(logging.INFO)
@@ -27,7 +28,7 @@ def __init__(self, data, logger):
         else:
             logger.debug("Progress indication is not enabled")
 
-    def read(self, size=None):
+    def read(self, size: int=None):
         pos_before = self._file.tell()
         result = self._file.read(size)
         if self._enabled:
@@ -75,31 +76,41 @@ def bucket_name(self):
     def credentials(self):
         return self._credentials
 
-    def fetch_model(self, source: str, file: str) -> None:
-        self._fetch(source, file)
-
-    def connect(self):
-        log = self._log
-        log.info("Connecting to the bucket...")
-
+    def create_client(self):
         # Client should be imported here because grpc starts threads during import
         # and if you call fork after that, a child process will be hang during exit
         from google.cloud.storage import Client
-
         if self.credentials:
             client = Client.from_service_account_json(self.credentials)
         else:
             client = Client()
+        return client
+
+    def connect(self):
+        log = self._log
+        log.info("Connecting to the bucket...")
+        client = self.create_client()
         return client.lookup_bucket(self.bucket_name)
 
-    def upload_model(self, path, meta, force):
+    def reset(self, force):
+        client = self.create_client()
+        bucket = client.lookup_bucket(self.bucket_name)
+        if bucket is not None:
+            if not force:
+                self._log.error("Bucket already exists, aborting.")
+                raise ExistingBackendError
+            self._log.info("Bucket already exists, deleting all content.")
+            bucket.delete(force, client)
+        client.create_bucket(self.bucket_name)
+
+    def upload_model(self, path: str, meta: dict, force: bool):
         bucket = self.connect()
         if bucket is None:
-            raise TransactionRequiredError
+            raise BackendRequiredError
         blob = bucket.blob("models/%s/%s.asdf" % (meta["model"], meta["uuid"]))
         if blob.exists() and not force:
             self._log.error("Model %s already exists, aborted.", meta["uuid"])
-            return 1
+            raise ModelAlreadyExistsError
         self._log.info("Uploading %s from %s...", meta["model"], os.path.abspath(path))
 
         def tracker(data):
@@ -124,23 +135,19 @@ def request_with_progress(method, url, data=None, headers=None, **kwargs):
         blob.make_public()
         return blob.public_url
 
-    def upload_index(self, index):
-        bucket = self.connect()
-        if bucket is None:
-            raise TransactionRequiredError
-        blob = bucket.blob(self.INDEX_FILE)
-        blob.cache_control = "no-cache"
-        blob.upload_from_string(json.dumps(index, indent=4, sort_keys=True))
-        blob.make_public()
-
-    def _fetch(self, url, where, chunk_size=DEFAULT_CHUNK_SIZE):
-        self._log.info("Fetching %s...", url)
-        r = requests.get(url, stream=True)
-        if isinstance(where, str):
-            os.makedirs(os.path.dirname(where), exist_ok=True)
-            f = open(where, "wb")
+    def fetch_model(self, source: str, file: Union[str, BinaryIO],
+                    chunk_size: int=DEFAULT_CHUNK_SIZE) -> None:
+        self._log.info("Fetching %s...", source)
+        r = requests.get(source, stream=True)
+        if r.status_code != 200:
+            self._log.error(
+                "An error occurred while fetching the model, with code %s" % r.status_code)
+            raise ValueError
+        if isinstance(file, str):
+            os.makedirs(os.path.dirname(file), exist_ok=True)
+            f = open(file, "wb")
         else:
-            f = where
+            f = file
         try:
             total_length = int(r.headers.get("content-length"))
             num_chunks = math.ceil(total_length / chunk_size)
@@ -153,13 +160,13 @@ def _fetch(self, url, where, chunk_size=DEFAULT_CHUNK_SIZE):
                     if chunk:
                         f.write(chunk)
         finally:
-            if isinstance(where, str):
+            if isinstance(file, str):
                 f.close()
 
-    def delete_model(self, meta):
+    def delete_model(self, meta: dict):
         bucket = self.connect()
         if bucket is None:
-            raise TransactionRequiredError
+            raise BackendRequiredError
         blob_name = "models/%s/%s.asdf" % (meta["model"], meta["uuid"])
         self._log.info(blob_name)
         try:

modelforge/index.py

@@ -9,9 +9,9 @@
 from dulwich.errors import HangupException, GitProtocolError, NotGitRepository
 from urllib.parse import urlparse
 
+import modelforge.configuration as config
 
 INDEX_FILE = "index.json"  #: Models repository index file name.
-DEFAULT_CACHE = os.path.expanduser("~/.cache")
 REMOTE_URL = "%s://%s%s/%s"  #: Remote repo url
 
 
@@ -23,8 +23,8 @@ class GitIndex:
         "add": "Add {model}/{uuid}",
     }
 
-    def __init__(self, index_repo: str= "", username: str= "", password: str= "",
-                 cache: str=DEFAULT_CACHE, log_level: int=logging.INFO):
+    def __init__(self, index_repo: str="", username: str= "", password: str= "", cache: str="",
+                 log_level: int=logging.INFO):
         """
         Initializes a new instance of :class:`GitIndex`.
 
@@ -37,14 +37,21 @@ def __init__(self, index_repo: str= "", username: str= "", password: str= "",
         """
         self._log = logging.getLogger(type(self).__name__)
         self._log.setLevel(log_level)
+        if index_repo is None and config.DEFAULT_REPO:
+            index_repo = config.DEFAULT_REPO
+        if cache is None:
+            cache = config.DEFAULT_CACHE
         parsed_url = urlparse(index_repo)
         if not parsed_url.scheme or \
                 parsed_url.scheme not in ("git", "git+ssh", "ssh", "http", "https"):
-            raise ValueError("Parsed url does not contain a valid protocol.")
+            self._log.critical("Parsed url does not contain a valid protocol.")
+            raise ValueError
         if not parsed_url.netloc:
-            raise ValueError("Parsed url does not contain a valid domain.")
+            self._log.critical("Parsed url does not contain a valid domain.")
+            raise ValueError
         if not parsed_url.path:
-            raise ValueError("Parsed url does not contain a valid repository path.")
+            self._log.critical("Parsed url does not contain a valid repository path.")
+            raise ValueError
         self.repo = parsed_url.path
         if self.repo.startswith("/"):
             self.repo = self.repo[1:]
@@ -55,21 +62,25 @@ def __init__(self, index_repo: str= "", username: str= "", password: str= "",
             auth = username + ":" + password + "@"
             self.remote_url = REMOTE_URL % (parsed_url.scheme, auth, parsed_url.netloc, self.repo)
         elif username or password:
-            raise ValueError("Both username and password must be supplied to access git with "
-                             "credentials.")
+            self._log.critical("Both username and password must be supplied to access git with "
+                               "credentials.")
+            raise ValueError
         else:
             self.remote_url = index_repo
-        self.content = None
+        self.contents = None
         try:
             self.fetch_index()
         except NotGitRepository as e:
-            raise ValueError("Repository does not exist: %s" % e) from e
+            self._log.critical("Repository does not exist: %s" % e)
+            raise ValueError
         except HangupException as e:
-            raise ValueError("Check SSH is configured, or connection is stable: %s" % e) from e
+            self._log.critical("Check SSH is configured, or connection is stable: %s" % e)
+            raise ValueError
         except GitProtocolError as e:
-            raise ValueError("%s: %s\nCheck your Git credentials." % (type(e), e))
-        self.models = self.content["models"]
-        self.meta = self.content["meta"]
+            self._log.critical("%s: %s\nCheck your Git credentials." % (type(e), e))
+            raise ValueError
+        self.models = self.contents["models"]
+        self.meta = self.contents["meta"]
 
     def fetch_index(self):
         os.makedirs(os.path.dirname(self.cached_repo), exist_ok=True)
@@ -82,7 +93,7 @@ def fetch_index(self):
                 self._log.info("Cached index is not up to date, pulling %s", self. repo)
                 git.pull(self.cached_repo, self.remote_url)
         with open(os.path.join(self.cached_repo, INDEX_FILE), encoding="utf-8") as _in:
-            self.content = json.load(_in)
+            self.contents = json.load(_in)
 
     def remove_model(self, model_uuid: str) -> dict:
         model_type = None
@@ -92,7 +103,8 @@ def remove_model(self, model_uuid: str) -> dict:
                 model_type = key
                 break
         if model_type is None:
-            raise ValueError("Model not found, aborted.")
+            self._log.error("Model not found, aborted.")
+            raise ValueError
         model_directory = os.path.join(self.cached_repo, model_type)
         model_node = self.models[model_type]
         meta_node = self.meta[model_type]
@@ -152,15 +164,15 @@ def initialize_index(self):
                 os.remove(path)
             elif os.path.isdir(path):
                 shutil.rmtree(path)
-        self.content = {"models": {}, "meta": {}}
+        self.contents = {"models": {}, "meta": {}}
 
     def upload_index(self, cmd: str, meta: dict):
         index = os.path.join(self.cached_repo, INDEX_FILE)
         if os.path.exists(index):
             os.remove(index)
         self._log.info("Writing the new index.json ...")
         with open(index, "w") as _out:
-            json.dump(self.content, _out)
+            json.dump(self.contents, _out)
         # implementation of git add --all is pretty bad, changing directory is the easiest way
         os.chdir(self.cached_repo)
         git.add()
@@ -169,15 +181,18 @@ def upload_index(self, cmd: str, meta: dict):
         # TODO: change when https://github.com/dulwich/dulwich/issues/631 gets addressed
         git.push(self.cached_repo, self.remote_url, b"master")
         if self._are_local_and_remote_heads_different():
-            raise ValueError("Push has failed")
+            self._log.error("Push has failed")
+            raise ValueError
 
     def load_template(self, template: str) -> Template:
         env = dict(trim_blocks=True, lstrip_blocks=True, keep_trailing_newline=False)
         jinja2_ext = ".jinja2"
         if not template.endswith(jinja2_ext):
-            raise ValueError("Template file name must end with %s" % jinja2_ext)
+            self._log.error("Template file name must end with %s" % jinja2_ext)
+            raise ValueError
         if not template[:-len(jinja2_ext)].endswith(".md"):
-            raise ValueError("Template file should be a Markdown file.")
+            self._log.error("Template file should be a Markdown file.")
+            raise ValueError
         with open(template, encoding="utf-8") as fin:
             template_obj = Template(fin.read(), **env)
         template_obj.filename = template

modelforge/model.py

@@ -81,7 +81,7 @@ def load(self, source: Union[str, BinaryIO, "Model"]=None, cache_dir: str=None,
                     if backend is None:
                         raise ValueError("The backend must be set to load a UUID or the default "
                                          "model.")
-                    index = backend.index.content
+                    index = backend.index.contents
                     config = index["models"]
                     if self.NAME is not None:
                         if not is_uuid: