Skip to content

Commit

Permalink
first commit, adding library and cli
Browse files Browse the repository at this point in the history
  • Loading branch information
Edgard Marx committed Apr 3, 2024
1 parent 979a41f commit 3e711f0
Show file tree
Hide file tree
Showing 2 changed files with 153 additions and 0 deletions.
70 changes: 70 additions & 0 deletions src/imr-cli.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,70 @@
import click
import yaml
from os.path import expanduser
from imr import IMRRemote, IMRLocal

home = expanduser("~")
imr_dir = home + "/.imr"
imr_local : IMRLocal = IMRLocal(imr_dir)
imr_remote : IMRRemote = None

def loadParams():
with open(imr_dir + "/config.yaml", 'r') as stream:
imrConfig = yaml.safe_load(stream)

@click.group()
def cli():
# loadParams()
pass

@cli.group()
def local():
pass

@local.command()
def list():
for package in imr_local.list():
print(package)

@local.command()
@click.argument('package')
@click.option('-v', '--version', type=str, default='latest', help='version of the model.', show_default=True)
def remove(package, version):
list = local.rm(package, version)

@cli.group()
@click.argument('host')
@click.argument('user')
@click.argument('password')
def remote(host, user, password):
global imr_remote
imr_remote = IMRRemote(host, user, password)

@remote.command()
def list():
list = imr_remote.list()
for p in list:
print(p)

@remote.command()
@click.argument('package')
@click.option('-v', '--version', type=str, default='latest', help='version of the model.', show_default=True)
def rm(package, version):
list = imr_remote.rm(package, version)

@remote.command()
@click.argument('model_directory')
@click.argument('package')
@click.option('-v', '--version', type=str, default='latest', help='version of the model.', show_default=True)
def push(model_directory, package, version):
list = imr_remote.push(model_directory, package, version)

@remote.command()
@click.argument('package')
@click.option('-d', '--dir', type=str, default=imr_dir, help='directory to pull the model in.', show_default=True)
@click.option('-v', '--version', type=str, default='latest', help='version of the model.', show_default=True)
def pull(package, dir, version):
list = imr_remote.pull(dir, package, version)

if __name__ == '__main__':
cli()
83 changes: 83 additions & 0 deletions src/imr.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,83 @@
# Inteligent Model Registry (imr)
from artifactory import ArtifactoryPath
from requests.auth import HTTPBasicAuth
import shutil, os
from os.path import expanduser
from os import listdir

class IMRRemote:

def __init__(self, repo:str, user:str, password:str):
self.repo = repo
self.user = user
self.password = password
self.home = expanduser("~")

def list(self):
packages = []
path = ArtifactoryPath(
self.repo,
auth=(self.user, self.password),
auth_type=HTTPBasicAuth)
for package in path.glob("*/*"):
packages.append(str(package).replace(self.repo + "/",""))
return packages

def push(self, directory: str, package: str, version = "latest"):
path = ArtifactoryPath(
self.repo + "/" + package + "/" + version,
auth=(self.user, self.password),
auth_type=HTTPBasicAuth)
path.mkdir()
shutil.make_archive("model", 'zip', directory)
path.deploy_file("model.zip")

def pull(self, directory: str, package: str, version = "latest"):
path = ArtifactoryPath(
self.repo + "/" + package + "/" + version + "/model.zip",
auth=(self.user, self.password),
auth_type=HTTPBasicAuth)
file_path = directory + "/" + package + "/" + version
if not os.path.exists(file_path):
os.makedirs(file_path)
with path.open() as fd, open(file_path + "/" + "model.zip", "wb") as out:
out.write(fd.read())
shutil.unpack_archive(file_path + "/" + "model.zip", file_path + "/model")

def rm(self, package: str, version = "latest"):
artefact = self.repo + "/" + package
if version is not None:
artefact = self.repo + "/" + package + "/" + version
path = ArtifactoryPath(
artefact,
auth=(self.user, self.password),
auth_type=HTTPBasicAuth)
if path.exists():
path.unlink()

class IMRLocal:

def __init__(self, repo:str | None = None):
self.home = expanduser("~")
if(repo is None) :
self.repo = self.home + "/.imr"
else:
self.repo = repo
if not os.path.exists(self.repo):
os.makedirs(self.repo)

def list(self):
list = []
for entry in os.walk(self.repo):
if(len(entry[1]) == 0 and entry[0] not in self.repo):
list.append(entry[0].replace(self.repo + "/",""))
return list

def push(self, directory: str, package: str, version = "latest"):
shutil.copytree(directory, self.repo + "/" + package + "/" + version)

def rm(self, package: str, version = "latest"):
if version is None:
shutil.rmtree(self.repo + "/" + package)
else:
shutil.rmtree(self.repo + "/" + package + "/" + version)

0 comments on commit 3e711f0

Please sign in to comment.