Repository: neuralmagic/sparsify (language: python)

Issue:
`sparsify.package` naive search
This PR is second in a series of PRs scoped out for `sparsify.package`, `alpha` release:

This PR includes the following:

- [X] A naive "working" `package` function that works with the cli + normal invocation
- [X] Support for searching zoo with `task, dataset and a deployment scenario`, (Only `vnni` supported for now, )
- [X] Broke the functionality into three main functions and added tests for all three
`sparsify.package` backend-call
This PR refactors code from #127 into a backend service call and appropriate functions were added in the backend in [diff #1](https://github.com/neuralmagic/model-api-service/pull/1)

Also depends on sparsezoo diff [#235](https://github.com/neuralmagic/sparsezoo/pull/235)


Discussion / hints:

Fix patch:
diff --git a/Makefile b/Makefile
index 66e312a8..6e999527 100644
--- a/Makefile
+++ b/Makefile
@@ -12,7 +12,11 @@ PYTEST_ARGS ?= ""
 INTEGRATION_TEST_ARGS ?= ""
 ifneq ($(findstring auto,$(TARGETS)),auto)
     PYTEST_ARGS := $(PYTEST_ARGS) --ignore tests/sparsify/auto
-	INTEGRATION_TEST_ARGS := $(INTEGRATION_TEST_ARGS) --ignore tests/sparsify/auto
+	INTEGRATION_TEST_ARGS := $(INTEGRATION_TEST_ARGS) --ignore tests/integration/auto
+endif
+ifneq ($(findstring package,$(TARGETS)),package)
+    PYTEST_ARGS := $(PYTEST_ARGS) --ignore tests/sparsify/package
+	INTEGRATION_TEST_ARGS := $(INTEGRATION_TEST_ARGS) --ignore tests/integration/package
 endif
 
 # run checks on all files for the repo
diff --git a/src/sparsify/package/cli.py b/src/sparsify/package/cli.py
index 073ad389..288ed8c0 100644
--- a/src/sparsify/package/cli.py
+++ b/src/sparsify/package/cli.py
@@ -82,6 +82,19 @@
 _LOGGER = logging.getLogger(__name__)
 
 
+def _get_template(results: str):
+    return f"""
+    Relevant Stub: {results}
+    Use sparsezoo to download the deployment directory as follows:
+    ```python
+    from sparsezoo import Model
+    model = Model("{results}")
+    model.deployment.download()
+    print(model.deployment.path)
+    ```
+    """
+
+
 @click.command(context_settings=dict(show_default=True))
 @click.version_option(version=__version__)
 @click.argument(
@@ -126,16 +139,18 @@ def main(**kwargs):
 
     Example for using sparsify.package:
 
-         1) `sparsify.package --task image_classification -m accuracy`
+         1) `sparsify.package --task image_classification \
+            --optimizing_metric accuracy`
 
          2) `sparsify.package --task ic --optimizing_metric accuracy \
          --optimizing_metric compression --target VNNI`
     """
     if not (kwargs.get("task") or kwargs.get("dataset")):
-        raise ValueError("At-least one of the `task` or `dataset`")
+        raise ValueError("At-least one of the `task` or `dataset` must be specified")
     _LOGGER.debug(f"{kwargs}")
     results = package(**kwargs)
-    print(f"Relevant Stubs: {results}")
+
+    print(_get_template(results))
 
 
 if __name__ == "__main__":
diff --git a/src/sparsify/package/config.py b/src/sparsify/package/config.py
new file mode 100644
index 00000000..c136f046
--- /dev/null
+++ b/src/sparsify/package/config.py
@@ -0,0 +1,34 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+
+__all__ = ["get_backend_url", "get_base_url"]
+_END_POINT = "/v1/sparsify/package/recommend-stub"
+_DEFAULT_BASE_URL = "https://api.neuralmagic.com"
+
+
+def get_base_url() -> str:
+    """
+    :return The base url for sparsify.package server
+    """
+    return os.getenv("SPARSIFY_BACKEND_URL", default=_DEFAULT_BASE_URL)
+
+
+def get_backend_url() -> str:
+    """
+    :return The backend url  for sparsify.package server
+    """
+    return get_base_url() + _END_POINT
diff --git a/src/sparsify/package/main.py b/src/sparsify/package/main.py
index c70c0706..6bf6b87f 100644
--- a/src/sparsify/package/main.py
+++ b/src/sparsify/package/main.py
@@ -12,19 +12,56 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
+from typing import Iterable, Optional, Union
+
+import requests
+
+from sparsify.package.config import get_backend_url
+
+
 __all__ = [
     "package",
 ]
 
-import logging
-
-
-_LOGGER = logging.getLogger(__file__)
+_LOGGER = logging.getLogger(__name__)
 
 
-def package(*args, **kwargs):
+def package(
+    task: Optional[str] = None,
+    dataset: Optional[str] = None,
+    scenario: Optional[str] = None,
+    optimizing_metric: Optional[Union[Iterable[str], str]] = None,
+    **kwargs,
+) -> str:
     """
-    A function that returns relevant sparsezoo stubs given the task or dataset,
-    and an optimizing criterion
+    A function that returns appropriate SparseZoo stub or deployment directory given
+    the task or dataset, optimizing criterions and a deployment scenario
+
+    :param directory: str A local existing directory path to download the deployment
+        artifacts to
+    :param task: str A supported task
+    :param dataset: str The public dataset this model was trained for
+    :param scenario: Optional[str] `VNNI` or `vnni for a VNNI compatible machine
+    :param optimizing_metric: Optional[List[str], str] representing different metrics
+        to prioritize for when searching for models
+    :return: The appropriate stub based on specified arguments
     """
-    raise NotImplementedError
+    optimizing_metric = (
+        [optimizing_metric] if isinstance(optimizing_metric, str) else optimizing_metric
+    )
+
+    payload = {
+        "task": task,
+        "dataset": dataset,
+        "scenario": scenario,
+        "optimizing_metric": optimizing_metric,
+    }
+
+    response = requests.get(
+        url=get_backend_url(),
+        headers={"Content-Type": "application/json"},
+        params=payload,
+    )
+
+    return response.json()
diff --git a/src/sparsify/utils/constants.py b/src/sparsify/utils/constants.py
index f4e18a7b..314a32be 100644
--- a/src/sparsify/utils/constants.py
+++ b/src/sparsify/utils/constants.py
@@ -33,8 +33,9 @@
     "f1",
     "recall",
     "mAP",
-    "compression",
     "latency",
+    "throughput",
+    "compression",
     "file_size",
     "memory_usage",
 ]


Test patch:
diff --git a/tests/integration/package/test_cli_run.py b/tests/integration/package/test_cli_run.py
new file mode 100644
index 00000000..455c66e4
--- /dev/null
+++ b/tests/integration/package/test_cli_run.py
@@ -0,0 +1,51 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import subprocess
+
+import pytest
+import requests
+
+from sparsify.package.config import get_base_url
+
+
+smallest_qa_model = (
+    "zoo:nlp/question_answering/bert-base/pytorch/huggingface/"
+    "squad/pruned95_obs_quant-none"
+)
+
+
+@pytest.mark.dependency()
+def test_server_is_up():
+    response = requests.get(get_base_url())
+    assert response.status_code == 200
+
+
+@pytest.mark.dependency(depends=["test_server_is_up"])
+@pytest.mark.slow()
+@pytest.mark.parametrize(
+    "command, expected_stub",
+    [
+        (
+            "sparsify.package --task qa --dataset squad "
+            "--optimizing_metric compression",
+            smallest_qa_model,
+        )
+    ],
+)
+def test_end_to_end_run(command, expected_stub):
+    subprocess.check_call(command.split())
+
+    output = subprocess.check_output(command.split())
+    assert expected_stub in output.decode()