[PLT-773] Allow users to create configs and associate them with projects through SDK. (#1591)

Author: sfendell-labelbox

docs/labelbox/index.rst

@@ -29,12 +29,14 @@ Labelbox Python SDK Documentation
     labeling-frontend-options
     labeling-parameter-override
     model
+    model-config
     model-run
     ontology
     ontology_kind
     organization
     pagination
     project
+    project-model-config
     quality-mode
     resource-tag
     review

docs/labelbox/model-config.rst

@@ -0,0 +1,6 @@
+Model Config
+===============================================================================================
+
+.. automodule:: labelbox.schema.model_config
+   :members:
+   :show-inheritance:

docs/labelbox/project-model-config.rst

@@ -0,0 +1,6 @@
+Project Model Config
+===============================================================================================
+
+.. automodule:: labelbox.schema.project_model_config
+   :members:
+   :show-inheritance:

libs/labelbox/src/labelbox/__init__.py

@@ -5,6 +5,7 @@
 from labelbox.client import Client
 from labelbox.schema.project import Project
 from labelbox.schema.model import Model
+from labelbox.schema.model_config import ModelConfig
 from labelbox.schema.bulk_import_request import BulkImportRequest
 from labelbox.schema.annotation_import import MALPredictionImport, MEAPredictionImport, LabelImport, MEAToMALPredictionImport
 from labelbox.schema.dataset import Dataset
@@ -29,6 +30,7 @@
 from labelbox.schema.benchmark import Benchmark
 from labelbox.schema.iam_integration import IAMIntegration
 from labelbox.schema.resource_tag import ResourceTag
+from labelbox.schema.project_model_config import ProjectModelConfig
 from labelbox.schema.project_resource_tag import ProjectResourceTag
 from labelbox.schema.media_type import MediaType
 from labelbox.schema.slice import Slice, CatalogSlice, ModelSlice

libs/labelbox/src/labelbox/client.py

@@ -38,6 +38,7 @@
 from labelbox.schema.labeling_frontend import LabelingFrontend
 from labelbox.schema.media_type import MediaType, get_media_type_validation_error
 from labelbox.schema.model import Model
+from labelbox.schema.model_config import ModelConfig
 from labelbox.schema.model_run import ModelRun
 from labelbox.schema.ontology import Ontology, DeleteFeatureFromOntologyResult
 from labelbox.schema.ontology import Tool, Classification, FeatureSchema
@@ -595,6 +596,56 @@ def _create(self, db_object_type, data, extra_params={}):
         res = res["create%s" % db_object_type.type_name()]
         return db_object_type(self, res)
 
+    def create_model_config(self, name: str, model_id: str, inference_params: dict) -> ModelConfig:
+        """ Creates a new model config with the given params.
+            Model configs are scoped to organizations, and can be reused between projects.
+
+        Args:
+            name (str): Name of the model config
+            model_id (str): ID of model to configure
+            inference_params (dict): JSON of model configuration parameters.
+
+        Returns:
+            str, id of the created model config
+        """
+
+        query = """mutation CreateModelConfigPyApi($modelId: ID!, $inferenceParams: Json!, $name: String!)  {
+                    createModelConfig(input: {modelId: $modelId, inferenceParams: $inferenceParams, name: $name}) {
+                        modelId
+                        inferenceParams
+                        id
+                        name
+                    }
+                }"""
+        params = {
+            "modelId": model_id,
+            "inferenceParams": inference_params,
+            "name": name
+        }
+        result = self.execute(query, params)
+        return ModelConfig(self, result['createModelConfig'])
+
+    def delete_model_config(self, id: str) -> bool:
+        """ Deletes an existing model config with the given id
+
+        Args:
+            id (str): ID of existing model config
+
+        Returns:
+            bool, indicates if the operation was a success.
+        """
+
+        query = """mutation DeleteModelConfigPyApi($id: ID!)  {
+                    deleteModelConfig(input: {id: $id}) {
+                        success
+                    }
+                }"""
+        params = {
+            "id": id
+        }
+        result = self.execute(query, params)
+        return result['deleteModelConfig']['success']
+
     def create_dataset(self,
                        iam_integration=IAMIntegration._DEFAULT,
                        **kwargs) -> Dataset:

libs/labelbox/src/labelbox/orm/model.py

@@ -362,6 +362,7 @@ class Entity(metaclass=EntityMeta):
         Task: Type[labelbox.Task]
         AssetAttachment: Type[labelbox.AssetAttachment]
         ModelRun: Type[labelbox.ModelRun]
+        ModelConfig: Type[labelbox.ModelConfig]
         Review: Type[labelbox.Review]
         User: Type[labelbox.User]
         LabelingFrontend: Type[labelbox.LabelingFrontend]
@@ -375,6 +376,7 @@ class Entity(metaclass=EntityMeta):
         Invite: Type[labelbox.Invite]
         InviteLimit: Type[labelbox.InviteLimit]
         ProjectRole: Type[labelbox.ProjectRole]
+        ProjectModelConfig: Type[labelbox.ProjectModelConfig]
         Project: Type[labelbox.Project]
         Batch: Type[labelbox.Batch]
         CatalogSlice: Type[labelbox.CatalogSlice]

libs/labelbox/src/labelbox/schema/model_config.py

@@ -0,0 +1,16 @@
+from labelbox.orm.db_object import DbObject
+from labelbox.orm.model import Field
+
+
+class ModelConfig(DbObject):
+    """ A ModelConfig represents a set of inference params configured for a model
+
+    Attributes: 
+        inference_params (JSON): Dict of inference params
+        model_id (str): ID of the model to configure
+        name (str): Name of config
+    """
+
+    inference_params = Field.Json("inference_params", "inferenceParams")
+    model_id = Field.String("model_id", "modelId")
+    name = Field.String("name", "name")

libs/labelbox/src/labelbox/schema/project.py

@@ -33,6 +33,7 @@
 from labelbox.schema.identifiable import DataRowIdentifier, GlobalKey, UniqueId
 from labelbox.schema.identifiables import DataRowIdentifiers, UniqueIds
 from labelbox.schema.media_type import MediaType
+from labelbox.schema.project_model_config import ProjectModelConfig
 from labelbox.schema.queue_mode import QueueMode
 from labelbox.schema.resource_tag import ResourceTag
 from labelbox.schema.task import Task
@@ -136,6 +137,28 @@ class Project(DbObject, Updateable, Deletable):
     def is_chat_evaluation(self) -> bool:
         return self.media_type == MediaType.Conversational and self.editor_task_type == EditorTaskType.ModelChatEvaluation
 
+    def project_model_configs(self):
+        query_str = """query ProjectModelConfigsPyApi($id: ID!) {
+            project(where: {id : $id}) {
+                projectModelConfigs {
+                    id
+                    modelConfigId
+                    modelConfig {
+                        id
+                        modelId
+                        inferenceParams
+                    }
+                    projectId
+                }
+            }
+        }"""
+        data = {"id": self.uid}
+        res = self.client.execute(query_str, data)
+        return [
+            ProjectModelConfig(self.client, projectModelConfig)
+            for projectModelConfig in res["project"]["projectModelConfigs"]
+        ]
+
     def update(self, **kwargs):
         """ Updates this project with the specified attributes
 
@@ -1226,6 +1249,50 @@ def get_queue_mode(self) -> "QueueMode":
         else:
             raise ValueError("Status not known")
 
+    def add_model_config(self, model_config_id: str) -> str:
+        """ Adds a model config to this project.
+
+        Args:
+            model_config_id (str): ID of a model config to add to this project.
+
+        Returns:
+            str, ID of the project model config association. This is needed for updating and deleting associations.
+        """
+
+        query = """mutation CreateProjectModelConfigPyApi($projectId: ID!, $modelConfigId: ID!)  {
+                    createProjectModelConfig(input: {projectId: $projectId, modelConfigId: $modelConfigId}) {
+                        projectModelConfigId
+                    }
+                }"""
+
+        params = {
+            "projectId": self.uid,
+            "modelConfigId": model_config_id,
+        }
+        result = self.client.execute(query, params)
+        return result["createProjectModelConfig"]["projectModelConfigId"]
+
+    def delete_project_model_config(self, project_model_config_id: str) -> bool:
+        """ Deletes the association between a model config and this project.
+
+        Args:
+            project_model_config_id (str): ID of a project model config association to delete for this project.
+
+        Returns:
+            bool, indicates if the operation was a success.
+        """
+        query = """mutation DeleteProjectModelConfigPyApi($id: ID!)  {
+                    deleteProjectModelConfig(input: {id: $id}) {
+                        success
+                    }
+                }"""
+
+        params = {
+            "id": project_model_config_id,
+        }
+        result = self.client.execute(query, params)
+        return result["deleteProjectModelConfig"]["success"]
+
     def set_labeling_parameter_overrides(
             self, data: List[LabelingParameterOverrideInput]) -> bool:
         """ Adds labeling parameter overrides to this project.

libs/labelbox/src/labelbox/schema/project_model_config.py

@@ -0,0 +1,34 @@
+from labelbox.orm.db_object import DbObject, Deletable
+from labelbox.orm.model import Field, Relationship
+
+
+class ProjectModelConfig(DbObject):
+    """ A ProjectModelConfig represents an association between a project and a single model config.
+
+    Attributes: 
+        project_id (str): ID of project to associate
+        model_config_id (str): ID of the model configuration
+        model_config (ModelConfig): Configuration for model
+    """
+
+    project_id = Field.String("project_id", "projectId")
+    model_config_id = Field.String("model_config_id", "modelConfigId")
+    model_config = Relationship.ToOne("ModelConfig", False, "model_config")
+
+    def delete(self) -> bool:
+        """ Deletes this association between a model config and this project.
+
+        Returns:
+            bool, indicates if the operation was a success.
+        """
+        query = """mutation DeleteProjectModelConfigPyApi($id: ID!)  {
+                    deleteProjectModelConfig(input: {id: $id}) {
+                        success
+                    }
+                }"""
+
+        params = {
+            "id": self.uid,
+        }
+        result = self.client.execute(query, params)
+        return result["deleteProjectModelConfig"]["success"]

libs/labelbox/tests/conftest.py

@@ -428,6 +428,12 @@ def consensus_project(client, rand_gen):
     project.delete()
 
 
+@pytest.fixture
+def model_config(client, rand_gen, valid_model_id):
+    model_config = client.create_model_config(name=rand_gen(str), model_id=valid_model_id, inference_params = {"param": "value"})
+    yield model_config
+    client.delete_model_config(model_config.uid)
+
 @pytest.fixture
 def consensus_project_with_batch(consensus_project, initial_dataset, rand_gen,
                                  image_url):
@@ -1043,3 +1049,7 @@ def embedding(client: Client):
     embedding = client.create_embedding(f"sdk-int-{uuid_str}", 8)
     yield embedding
     embedding.delete()
+
+@pytest.fixture
+def valid_model_id():
+    return "2c903542-d1da-48fd-9db1-8c62571bd3d2"