Add tags to Models (#316)

Author: jean-lucas

CHANGELOG.md

@@ -5,11 +5,30 @@ All notable changes to the [Nucleus Python Client](https://github.com/scaleapi/n
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+
+## [0.14.0](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.14.0) - 2022-06-16
+
+### Added
+
+- Allow creation/deletion of model tags on new and existing models, eg:
+```python
+# on model creation
+model = client.create_model(name="foo_model", reference_id="foo-model-ref", tags=["some tag"])
+
+# on existing models
+existing_model = client.models[0]
+existing_model.add_tags(['tag a', 'tag b'])
+
+# remove tag
+existing_model.remove_tags(['tag a'])
+```
+
 ## [0.13.5](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.13.4) - 2022-06-15
 
 ### Fixed
 - Guard against invalid skeleton indexes in KeypointsAnnotation
 
+
 ## [0.13.4](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.13.4) - 2022-06-09
 
 ### Fixed
@@ -37,7 +56,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Segmentation functions to Validate API
- 
+
 ## [0.12.4](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.12.4) - 2022-06-02
 
 ### Fixed

nucleus/__init__.py

@@ -91,6 +91,7 @@
     KEEP_HISTORY_KEY,
     MESSAGE_KEY,
     MODEL_RUN_ID_KEY,
+    MODEL_TAGS_KEY,
     NAME_KEY,
     NUCLEUS_ENDPOINT,
     PREDICTIONS_IGNORED_KEY,
@@ -218,6 +219,7 @@ def models(self) -> List[Model]:
                 reference_id=model["ref_id"],
                 metadata=model["metadata"] or None,
                 client=self,
+                tags=model.get(MODEL_TAGS_KEY, []),
             )
             for model in model_objects["models"]
         ]
@@ -484,6 +486,7 @@ def create_model(
         reference_id: str,
         metadata: Optional[Dict] = None,
         bundle_name: Optional[str] = None,
+        tags: Optional[List[str]] = None,
     ) -> Model:
         """Adds a :class:`Model` to Nucleus.
 
@@ -495,13 +498,15 @@ def create_model(
             metadata: An arbitrary dictionary of additional data about this model
               that can be stored and retrieved. For example, you can store information
               about the hyperparameters used in training this model.
+            bundle_name: Optional name of bundle attached to this model
+            tags: Optional list of tags to attach to this model
 
         Returns:
             :class:`Model`: The newly created model as an object.
         """
         response = self.make_request(
             construct_model_creation_payload(
-                name, reference_id, metadata, bundle_name
+                name, reference_id, metadata, bundle_name, tags
             ),
             "models/add",
         )
@@ -516,6 +521,7 @@ def create_model(
             metadata=metadata,
             bundle_name=bundle_name,
             client=self,
+            tags=tags,
         )
 
     def create_launch_model(

nucleus/constants.py

@@ -90,6 +90,7 @@
 MESSAGE_KEY = "message"
 METADATA_KEY = "metadata"
 MODEL_BUNDLE_NAME_KEY = "bundle_name"
+MODEL_TAGS_KEY = "tags"
 MODEL_ID_KEY = "model_id"
 MODEL_RUN_ID_KEY = "model_run_id"
 NAME_KEY = "name"

nucleus/model.py

@@ -2,7 +2,7 @@
 
 import requests
 
-from .constants import METADATA_KEY, NAME_KEY, REFERENCE_ID_KEY
+from .constants import METADATA_KEY, MODEL_TAGS_KEY, NAME_KEY, REFERENCE_ID_KEY
 from .dataset import Dataset
 from .job import AsyncJob
 from .model_run import ModelRun
@@ -93,13 +93,21 @@ class Model:
     """
 
     def __init__(
-        self, model_id, name, reference_id, metadata, client, bundle_name=None
+        self,
+        model_id,
+        name,
+        reference_id,
+        metadata,
+        client,
+        bundle_name=None,
+        tags: List[str] = None,
     ):
         self.id = model_id
         self.name = name
         self.reference_id = reference_id
         self.metadata = metadata
         self.bundle_name = bundle_name
+        self.tags = tags if tags else []
         self._client = client
 
     def __repr__(self):
@@ -213,3 +221,49 @@ def run(self, dataset_id: str, slice_id: Optional[str]) -> str:
         )
 
         return response
+
+    def add_tags(self, tags: List[str]):
+        """Tag the model with custom tag names. ::
+
+            import nucleus
+            client = nucleus.NucleusClient("YOUR_SCALE_API_KEY")
+            model = client.list_models()[0]
+
+            model.add_tags(["tag_A", "tag_B"])
+
+        Args:
+            tags: list of tag names
+        """
+        response = self._client.make_request(
+            {MODEL_TAGS_KEY: tags},
+            f"model/{self.id}/tag",
+            requests_command=requests.post,
+        )
+
+        if response.get("msg", False):
+            self.tags.extend(tags)
+
+        return response
+
+    def remove_tags(self, tags: List[str]):
+        """Remove tag(s) from the model. ::
+
+            import nucleus
+            client = nucleus.NucleusClient("YOUR_SCALE_API_KEY")
+            model = client.list_models()[0]
+
+            model.remove_tags(["tag_x"])
+
+        Args:
+            tags: list of tag names to remove
+        """
+        response = self._client.make_request(
+            {MODEL_TAGS_KEY: tags},
+            f"model/{self.id}/tag",
+            requests_command=requests.delete,
+        )
+
+        if response.get("msg", False):
+            self.tags = list(filter(lambda t: t not in tags, self.tags))
+
+        return response

nucleus/payload_constructor.py

@@ -17,6 +17,7 @@
     METADATA_KEY,
     MODEL_BUNDLE_NAME_KEY,
     MODEL_ID_KEY,
+    MODEL_TAGS_KEY,
     NAME_KEY,
     REFERENCE_ID_KEY,
     SCENES_KEY,
@@ -127,6 +128,7 @@ def construct_model_creation_payload(
     reference_id: str,
     metadata: Optional[Dict],
     bundle_name: Optional[str],
+    tags: Optional[List[str]],
 ) -> dict:
     payload = {
         NAME_KEY: name,
@@ -136,6 +138,8 @@ def construct_model_creation_payload(
 
     if bundle_name:
         payload[MODEL_BUNDLE_NAME_KEY] = bundle_name
+    if tags:
+        payload[MODEL_TAGS_KEY] = tags
 
     return payload
 

pyproject.toml

@@ -21,7 +21,7 @@ exclude = '''
 
 [tool.poetry]
 name = "scale-nucleus"
-version = "0.13.5"
+version = "0.14.0"
 description = "The official Python client library for Nucleus, the Data Platform for AI"
 license =  "MIT"
 authors = ["Scale AI Nucleus Team <nucleusapi@scaleapi.com>"]

tests/test_models.py

@@ -115,3 +115,46 @@ def test_new_model_endpoints(CLIENT, dataset: Dataset):
         model, predictions[0].reference_id, predictions[0].annotation_id
     )
     assert_box_prediction_matches_dict(prediction_loc, TEST_BOX_PREDICTIONS[0])
+
+
+def test_tag_model(CLIENT, dataset: Dataset):
+    def testing_model(ref_id):
+        models_from_backend = list(
+            filter(lambda m: m.reference_id == ref_id, CLIENT.models)
+        )
+        assert len(models_from_backend) == 1
+        return models_from_backend[0]
+
+    model_reference = "model_" + str(time.time())
+    model = CLIENT.create_model(
+        TEST_MODEL_NAME, model_reference, tags=["first_tag"]
+    )
+
+    model.add_tags(["single tag"])
+    model.add_tags(["tag_a", "tag_b"])
+
+    backend_model = testing_model(model_reference)
+    assert sorted(backend_model.tags) == sorted(
+        ["first_tag", "single tag", "tag_a", "tag_b"]
+    )
+
+    model.remove_tags(["tag_a"])
+    model.remove_tags(["first_tag", "tag_b"])
+
+    backend_model = testing_model(model_reference)
+    assert backend_model.tags == ["single tag"]
+
+
+def test_remove_invalid_tag_from_model(CLIENT, dataset: Dataset):
+
+    model_reference = "model_" + str(time.time())
+    model = CLIENT.create_model(TEST_MODEL_NAME, model_reference)
+
+    model.add_tags(["single tag"])
+
+    response = model.remove_tags(["tag_a"])
+    assert "error" in response
+    assert (
+        response["error"]
+        == "Deleted 0 tags from model. Either the tag(s) did not exist, or you are not the owner of this model project."
+    )