[Launch] add bundle_id for model creation (#270)

Author: syandroo

nucleus/__init__.py

@@ -479,7 +479,11 @@ def add_model(
         return self.create_model(name, reference_id, metadata)
 
     def create_model(
-        self, name: str, reference_id: str, metadata: Optional[Dict] = None
+        self,
+        name: str,
+        reference_id: str,
+        metadata: Optional[Dict] = None,
+        bundle_name: Optional[str] = None,
     ) -> Model:
         """Adds a :class:`Model` to Nucleus.
 
@@ -496,14 +500,23 @@ def create_model(
             :class:`Model`: The newly created model as an object.
         """
         response = self.make_request(
-            construct_model_creation_payload(name, reference_id, metadata),
+            construct_model_creation_payload(
+                name, reference_id, metadata, bundle_name
+            ),
             "models/add",
         )
         model_id = response.get("model_id", None)
         if not model_id:
             raise ModelCreationError(response.get("error"))
 
-        return Model(model_id, name, reference_id, metadata, self)
+        return Model(
+            model_id=model_id,
+            name=name,
+            reference_id=reference_id,
+            metadata=metadata,
+            bundle_name=bundle_name,
+            client=self,
+        )
 
     @deprecated(
         "Model runs have been deprecated and will be removed. Use a Model instead"

nucleus/constants.py

@@ -88,6 +88,7 @@
 MAX_PAYLOAD_SIZE = 0x1FFFFFE8  # Set to max string size since we currently convert payloads to strings for processing on the server-side
 MESSAGE_KEY = "message"
 METADATA_KEY = "metadata"
+MODEL_BUNDLE_NAME_KEY = "bundle_name"
 MODEL_ID_KEY = "model_id"
 MODEL_RUN_ID_KEY = "model_run_id"
 NAME_KEY = "name"

nucleus/model.py

@@ -92,22 +92,26 @@ class Model:
     endpoint, using :meth:`NucleusClient.add_model`.
     """
 
-    def __init__(self, model_id, name, reference_id, metadata, client):
+    def __init__(
+        self, model_id, name, reference_id, metadata, client, bundle_name=None
+    ):
         self.id = model_id
         self.name = name
         self.reference_id = reference_id
         self.metadata = metadata
+        self.bundle_name = bundle_name
         self._client = client
 
     def __repr__(self):
-        return f"Model(model_id='{self.id}', name='{self.name}', reference_id='{self.reference_id}', metadata={self.metadata}, client={self._client})"
+        return f"Model(model_id='{self.id}', name='{self.name}', reference_id='{self.reference_id}', metadata={self.metadata}, bundle_name={self.bundle_name}, client={self._client})"
 
     def __eq__(self, other):
         return (
             (self.id == other.id)
             and (self.name == other.name)
             and (self.metadata == other.metadata)
             and (self._client == other._client)
+            and (self.bundle_name == other.bundle_name)
         )
 
     def __hash__(self):
@@ -187,3 +191,25 @@ def evaluate(self, scenario_test_names: List[str]) -> AsyncJob:
             requests_command=requests.post,
         )
         return AsyncJob.from_json(response, self._client)
+
+    def run(self, dataset_id: str, slice_id: Optional[str]) -> str:
+        """Runs inference on the bundle associated with the model on the dataset. ::
+
+            import nucleus
+            client = nucleus.NucleusClient("YOUR_SCALE_API_KEY")
+            model = client.list_models()[0]
+
+            model.run("ds_123456")
+
+        Args:
+            dataset_id: id of dataset to run inference on
+            job_id: nucleus job used to track async job progress
+            slice_id: (optional) id of slice of the dataset to run inference on
+        """
+        response = self._client.make_request(
+            {"dataset_id": dataset_id, "slice_id": slice_id},
+            f"model/run/{self.id}/",
+            requests_command=requests.post,
+        )
+
+        return response

nucleus/payload_constructor.py

@@ -15,6 +15,7 @@
     ITEMS_KEY,
     LABELS_KEY,
     METADATA_KEY,
+    MODEL_BUNDLE_NAME_KEY,
     MODEL_ID_KEY,
     NAME_KEY,
     REFERENCE_ID_KEY,
@@ -122,14 +123,22 @@ def construct_box_predictions_payload(
 
 
 def construct_model_creation_payload(
-    name: str, reference_id: str, metadata: Optional[Dict]
+    name: str,
+    reference_id: str,
+    metadata: Optional[Dict],
+    bundle_name: Optional[str],
 ) -> dict:
-    return {
+    payload = {
         NAME_KEY: name,
         REFERENCE_ID_KEY: reference_id,
         METADATA_KEY: metadata if metadata else {},
     }
 
+    if bundle_name:
+        payload[MODEL_BUNDLE_NAME_KEY] = bundle_name
+
+    return payload
+
 
 def construct_model_run_creation_payload(
     name: str,