def delete_model(args: argparse.Namespace, backend: StorageBackend,
log: logging.Logger):
"""
Delete a model.
:param args: :class:`argparse.Namespace` with "input", "backend", "args", "meta", \
"update_default", "username", "password", "remote_repo", \
"template_model", "template_readme" and "log_level".
:param backend: Backend which is responsible for working with model files.
:param log: Logger supplied by supply_backend
:return: None
"""
try:
meta = backend.index.remove_model(args.input)
template_readme = backend.index.load_template(args.template_readme)
backend.index.update_readme(template_readme)
except ValueError:
return 1
backend.delete_model(meta)
log.info("Updating the models index...")
try:
backend.index.upload("delete", meta)
except ValueError: # TODO: replace with PorcelainError
return 1
log.info("Successfully deleted.")
def publish_model(args: argparse.Namespace, backend: StorageBackend,
log: logging.Logger):
"""
Pushes the model to Google Cloud Storage and updates the index file.
:param args: :class:`argparse.Namespace` with "model", "backend", "args", "force" and \
"update_default".
:return: None if successful, 1 otherwise.
"""
path = os.path.abspath(args.model)
try:
model = GenericModel(source=path, dummy=True)
except ValueError as e:
log.critical('"model" must be a path: %s', e)
return 1
except Exception as e:
log.critical("Failed to load the model: %s: %s" %
(type(e).__name__, e))
return 1
meta = model.meta
with backend.lock():
model_url = backend.upload_model(path, meta, args.force)
log.info("Uploaded as %s", model_url)
log.info("Updating the models index...")
index = backend.fetch_index()
index["models"].setdefault(meta["model"], {})[meta["uuid"]] = \
extract_index_meta(meta, model_url)
if args.update_default:
index["models"][meta["model"]][Model.DEFAULT_NAME] = meta["uuid"]
backend.upload_index(index)
def initialize_registry(args: argparse.Namespace, backend: StorageBackend, log: logging.Logger):
"""
Initialize the registry and the index.
:param args: :class:`argparse.Namespace` with "backend", "args", "force" and "log_level".
:param backend: Backend which is responsible for working with model files.
:param log: Logger supplied by supply_backend
:return: None
"""
try:
backend.reset(args.force)
except ExistingBackendError:
return 1
log.info("Resetting the index ...")
backend.index.reset()
try:
backend.index.upload("initialize", {})
except ValueError:
return 1
log.info("Successfully initialized")
def initialize_registry(args: argparse.Namespace, backend: StorageBackend,
log: logging.Logger):
"""
Initialize the registry - list and publish will fail otherwise.
:param args: :class:`argparse.Namespace` with "backend", "args" and "force".
:return: None
"""
try:
backend.fetch_index()
if not args.force:
log.warning("Registry is already initialized")
return
except FileNotFoundError:
pass
# The lock is not needed here, but upload_index() will raise otherwise
with backend.lock():
backend.upload_index({"models": {}})
log.info("Successfully initialized")
def publish_model(args: argparse.Namespace, backend: StorageBackend,
log: logging.Logger):
"""
Push the model to Google Cloud Storage and updates the index file.
:param args: :class:`argparse.Namespace` with "model", "backend", "args", "force", "meta" \
"update_default", "username", "password", "remote_repo", "template_model", \
"template_readme" and "log_level".
:param backend: Backend which is responsible for working with model files.
:param log: Logger supplied by supply_backend
:return: None if successful, 1 otherwise.
"""
path = os.path.abspath(args.model)
try:
model = GenericModel(source=path, dummy=True)
except ValueError as e:
log.critical('"model" must be a path: %s', e)
return 1
except Exception as e:
log.critical("Failed to load the model: %s: %s" %
(type(e).__name__, e))
return 1
base_meta = model.meta
try:
model_url = backend.upload_model(path, base_meta, args.force)
except ModelAlreadyExistsError:
return 1
log.info("Uploaded as %s", model_url)
with open(os.path.join(args.meta), encoding="utf-8") as _in:
extra_meta = json.load(_in)
model_type, model_uuid = base_meta["model"], base_meta["uuid"]
meta = extract_model_meta(base_meta, extra_meta, model_url)
log.info("Updating the models index...")
try:
template_model = backend.index.load_template(args.template_model)
template_readme = backend.index.load_template(args.template_readme)
except ValueError:
return 1
backend.index.add_model(model_type, model_uuid, meta, template_model,
args.update_default)
backend.index.update_readme(template_readme)
try:
backend.index.upload("add", {"model": model_type, "uuid": model_uuid})
except ValueError: # TODO: replace with PorcelainError, see related TODO in index.py:181
return 1
log.info("Successfully published.")
def list_models(args: argparse.Namespace, backend: StorageBackend,
log: logging.Logger):
"""
Outputs the list of known models in the registry.
:param args: :class:`argparse.Namespace` with "backend" and "args".
:return: None
"""
index = backend.fetch_index()
for key, val in index["models"].items():
print(key)
default = None
for mid, meta in val.items():
if mid == "default":
default = meta
break
for mid, meta in sorted(
[m for m in val.items() if m[1] != default],
key=lambda m: parse_datetime(m[1]["created_at"]),
reverse=True):
print(" %s %s" % ("*" if default == mid else " ", mid),
meta["created_at"])
def load(self,
source: Union[str, BinaryIO, "Model"] = None,
cache_dir: str = None,
backend: StorageBackend = None,
lazy=False) -> "Model":
"""
Build a new Model instance.
:param source: UUID, file system path, file object or an URL; None means auto.
:param cache_dir: The directory where to store the downloaded model.
:param backend: Remote storage backend to use if ``source`` is a UUID or a URL.
:param lazy: Do not really load numpy arrays into memory. Instead, mmap() them. \
User is expected to call Model.close() when the tree is no longer needed.
"""
if isinstance(source, Model):
if not isinstance(source, type(self)):
raise TypeError("Incompatible model instance: %s <> %s" %
(type(source), type(self)))
self.__dict__ = source.__dict__
return self
if backend is not None and not isinstance(backend, StorageBackend):
raise TypeError("backend must be an instance of "
"modelforge.storage_backend.StorageBackend")
self._source = str(source)
generic = self.NAME == self.GENERIC_NAME
try:
if source is None or (isinstance(source, str)
and not os.path.isfile(source)):
if cache_dir is None:
if not generic:
cache_dir = os.path.join(self.cache_dir(), self.NAME)
else:
cache_dir = tempfile.mkdtemp(prefix="modelforge-")
try:
uuid.UUID(source)
is_uuid = True
except (TypeError, ValueError):
is_uuid = False
model_id = self.DEFAULT_NAME if not is_uuid else source
file_name = model_id + self.DEFAULT_FILE_EXT
file_name = os.path.join(os.path.expanduser(cache_dir),
file_name)
if os.path.exists(file_name) and (not source or
not os.path.exists(source)):
source = file_name
elif source is None or is_uuid:
if backend is None:
raise ValueError(
"The backend must be set to load a UUID or the default "
"model.")
index = backend.index.contents
config = index["models"]
if not generic:
if not is_uuid:
model_id = index["meta"][self.NAME][model_id]
source = config[self.NAME][model_id]
else:
if not is_uuid:
raise ValueError(
"File path, URL or UUID is needed.")
for models in config.values():
if source in models:
source = models[source]
break
else:
raise FileNotFoundError("Model %s not found." %
source)
source = source["url"]
if re.match(r"\w+://", source):
if backend is None:
raise ValueError(
"The backend must be set to load a URL.")
backend.fetch_model(source, file_name)
self._source = source
source = file_name
if isinstance(source, str):
size = os.stat(source).st_size
else:
self._source = "<file object>"
pos = source.tell()
size = source.seek(0, os.SEEK_END) - pos
source.seek(pos, os.SEEK_SET)
self._log.info("Reading %s (%s)...", source,
humanize.naturalsize(size))
model = asdf.open(source, copy_arrays=not lazy, lazy_load=lazy)
try:
tree = model.tree
self._meta = tree["meta"]
self._initial_version = list(self.version)
if not generic:
meta_name = self._meta["model"]
matched = self.NAME == meta_name
if not matched:
needed = {self.NAME}
for child in type(self).__subclasses__():
needed.add(child.NAME)
matched |= child.NAME == meta_name
if not matched:
raise ValueError(
"The supplied model is of the wrong type: needed "
"%s, got %s." % (needed, meta_name))
self._load_tree(tree)
finally:
if not lazy:
model.close()
else:
self._asdf = model
finally:
if generic and cache_dir is not None:
shutil.rmtree(cache_dir)
self._size = size
return self
def load(self,
source: Union[str, BinaryIO, "Model"] = None,
cache_dir: str = None,
backend: StorageBackend = None) -> "Model":
"""
Initializes a new Model instance.
:param source: UUID, file system path, file object or an URL; None means auto.
:param cache_dir: The directory where to store the downloaded model.
:param backend: Remote storage backend to use if ``source`` is a UUID or a URL.
"""
if isinstance(source, Model):
if not isinstance(source, type(self)):
raise TypeError("Incompatible model instance: %s <> %s" %
(type(source), type(self)))
self.__dict__ = source.__dict__
return self
if backend is not None and not isinstance(backend, StorageBackend):
raise TypeError("backend must be an instance of "
"modelforge.storage_backend.StorageBackend")
self._source = str(source)
try:
if source is None or (isinstance(source, str)
and not os.path.isfile(source)):
if cache_dir is None:
if self.NAME is not None:
cache_dir = os.path.join(self.cache_dir(), self.NAME)
else:
cache_dir = tempfile.mkdtemp(prefix="modelforge-")
try:
uuid.UUID(source)
is_uuid = True
except (TypeError, ValueError):
is_uuid = False
model_id = self.DEFAULT_NAME if not is_uuid else source
file_name = model_id + self.DEFAULT_FILE_EXT
file_name = os.path.join(os.path.expanduser(cache_dir),
file_name)
if os.path.exists(file_name) and (not source or
not os.path.exists(source)):
source = file_name
elif source is None or is_uuid:
if backend is None:
raise ValueError(
"The backend must be set to load a UUID or the default "
"model.")
index = backend.index.contents
config = index["models"]
if self.NAME is not None:
if not is_uuid:
model_id = index["meta"][self.NAME][model_id]
source = config[self.NAME][model_id]
else:
if not is_uuid:
raise ValueError(
"File path, URL or UUID is needed.")
for models in config.values():
if source in models:
source = models[source]
break
else:
raise FileNotFoundError("Model %s not found." %
source)
source = source["url"]
if re.match(r"\w+://", source):
if backend is None:
raise ValueError(
"The backend must be set to load a URL.")
backend.fetch_model(source, file_name)
source = file_name
self._log.info("Reading %s...", source)
with asdf.open(source) as model:
tree = model.tree
self._meta = tree["meta"]
if self.NAME is not None:
meta_name = self._meta["model"]
matched = self.NAME == meta_name
if not matched:
needed = {self.NAME}
for child in type(self).__subclasses__():
needed.add(child.NAME)
matched |= child.NAME == meta_name
if not matched:
raise ValueError(
"The supplied model is of the wrong type: needed "
"%s, got %s." % (needed, meta_name))
self._load_tree(tree)
finally:
if self.NAME is None and cache_dir is not None:
shutil.rmtree(cache_dir)
return self