feat: [VRD-909] Add pii to client (#3835)

client/verta/tests/registry/test_model.py

@@ -132,6 +132,24 @@ def test_description(self, client, created_entities):
         registered_model.set_description(desc)
         assert desc == registered_model.get_description()
 
+    def test_pii(self, client, created_entities):
+        # Test setting on creation
+        registered_model = client.get_or_create_registered_model(pii=True)
+        created_entities.append(registered_model)
+        assert registered_model.get_pii() is True
+        registered_model = client.create_registered_model(pii=True)
+        created_entities.append(registered_model)
+        assert registered_model.get_pii() is True
+
+        # Test updating
+        registered_model = client.get_or_create_registered_model()
+        created_entities.append(registered_model)
+        registered_model.set_pii(True)
+        assert registered_model.get_pii() is True
+        registered_model.set_pii(False)
+        assert registered_model.get_pii() is False
+
+
 
 class TestTaskTypes:
     @pytest.mark.parametrize(

client/verta/verta/client.py

@@ -780,6 +780,7 @@ def get_or_create_registered_model(
         id=None,
         task_type=None,
         data_type=None,
+        pii=False,
     ):
         """
         Attaches a registered_model to this Client.
@@ -815,6 +816,8 @@ def get_or_create_registered_model(
             Task type of the registered_model.
         data_type : :mod:`~verta.registry.data_type`, optional
             Data type of the registered_model.
+        pii : bool, default False
+            Whether the registered_model ingests personally identifiable information.
 
         Returns
         -------
@@ -863,6 +866,7 @@ def get_or_create_registered_model(
                     visibility=visibility,
                     task_type=task_type,
                     data_type=data_type,
+                    pii=pii,
                 ),
                 lambda: check_unnecessary_params_warning(
                     resource_name, "name {}".format(name), param_names, params
@@ -1250,6 +1254,7 @@ def create_registered_model(
         visibility=None,
         task_type=None,
         data_type=None,
+        pii=False,
     ):
         """
         Creates a new Registered Model.
@@ -1280,6 +1285,8 @@ def create_registered_model(
             Task type of the registered_model.
         data_type : :mod:`~verta.registry.data_type`, optional
             Data type of the registered_model.
+        pii : bool, default False
+            Whether the registered_model ingests personally identifiable information.
 
         Returns
         -------
@@ -1312,6 +1319,7 @@ def create_registered_model(
             visibility=visibility,
             task_type=task_type,
             data_type=data_type,
+            pii=pii,
         )
 
     def create_endpoint(

client/verta/verta/registry/entities/_model.py

@@ -36,6 +36,8 @@ class RegisteredModel(_entity._ModelDBEntity):
         Verta web app URL.
     versions : iterable of :class:`~verta.registry.entities.RegisteredModelVersion`
         Versions of this RegisteredModel.
+    pii: bool
+         Whether the registered_model ingests personally identifiable information.
 
     """
 
@@ -61,6 +63,7 @@ def __repr__(self):
                 "description: {}".format(msg.description),
                 "labels: {}".format(msg.labels),
                 "id: {}".format(msg.id),
+                "pii: {}".format(msg.pii),
             )
         )
 
@@ -1034,6 +1037,7 @@ def _create_proto_internal(
         visibility=None,
         task_type=None,
         data_type=None,
+        pii=None,
     ):
         if task_type is None:
             task_type = task_type_module._Unknown()
@@ -1048,6 +1052,7 @@ def _create_proto_internal(
             time_updated=date_created,
             task_type=task_type._as_proto(),
             data_type=data_type._as_proto(),
+            pii=pii,
         )
         if (
             public_within_org
@@ -1083,6 +1088,18 @@ def _create_proto_internal(
 
     RegisteredModelMessage = _RegistryService.RegisteredModel
 
+    def set_pii(self, pii):
+        """
+        Note that you *MUST* update any live endpoints running this model in order to propagate this change.
+        """
+        self._fetch_with_no_cache()
+        self._msg.pii = pii
+        self._update(self._msg, method="PUT")
+
+    def get_pii(self):
+        self._refresh_cache()
+        return self._msg.pii
+
     def set_description(self, desc):
         if not desc:
             raise ValueError("desc is not specified")