py(core): draft CLI

omlmd/cli.py

@@ -0,0 +1,36 @@
+# Using this to scope CLI targets
+import click
+from omlmd.helpers import Helper
+
+@click.group()
+def cli():
+    pass
+
+@click.command()
+@click.argument('target', required=True)
+@click.option('-o', '--output', default='.', show_default=True)
+@click.option('--media-types', '-m', multiple=True, default=[])
+def pull(target, output, media_types):
+    """Pulls an OCI Artifact containing ML model and metadata, filtering if necessary."""
+    Helper().pull(target, output, media_types)
+
+@click.group()
+def get():
+    pass
+
+@click.command()
+@click.argument('target', required=True)
+def config(target):
+    """Outputs configuration of the given OCI Artifact for ML model and metadata."""
+    click.echo(Helper().get_config(target))
+
+@click.command()
+@click.argument('targets', required=True, nargs=-1)
+def crawl(targets):
+    """Crawls configuration for the given list of OCI Artifact for ML model and metadata."""
+    click.echo(Helper().crawl(targets))
+
+cli.add_command(pull)
+cli.add_command(get)
+get.add_command(config)
+cli.add_command(crawl)

omlmd/helpers.py

@@ -56,7 +56,7 @@ def push(
         write_content_to_file("model_metadata.omlmd.json", lambda: model_metadata.to_json())
         write_content_to_file("model_metadata.omlmd.yaml", lambda: model_metadata.to_yaml())
         files = [
-            f"{path}:application/x-artifact",
+            f"{path}:application/x-mlmodel",
             "model_metadata.omlmd.json:application/x-config",
             "model_metadata.omlmd.yaml:application/x-config",
         ]

omlmd/provider.py

@@ -26,7 +26,7 @@ def download_layers(self, package, download_dir, media_types):
         paths = []
 
         for layer in manifest.get('layers', []):
-            if layer['mediaType'] in media_types:
+            if len(media_types) == 0 or layer['mediaType'] in media_types:
                 artifact = layer['annotations']['org.opencontainers.image.title']
                 outfile = oras.utils.sanitize_path(download_dir, os.path.join(download_dir, artifact))
                 path = self.download_blob(package, layer['digest'], outfile)

pyproject.toml

@@ -10,13 +10,17 @@ python = "^3.10"
 oras = "^0.1.30"
 pyyaml = "^6.0.1"
 nbconvert = "^7.16.4"
+click = "^8.1.7"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.2.2"
 jq = "^1.7.0"
 scikit-learn = "^1.5.0"
 ipykernel = "^6.29.4"
 
+[tool.poetry.scripts]
+omlmd = "omlmd.cli:cli"
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"