Make PUT request for reference dataset upload

Author: gustavocidornelas

openlayer/__init__.py

@@ -1648,9 +1648,24 @@ def upload_reference_dataset(
             {"task_type": task_type.value, **dataset_config}
         )
 
-        # TODO: Make POST request to upload dataset
-        print("Uploading reference dataset...")
-        print(dataset_data)
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            # Copy relevant files to tmp dir
+            utils.write_yaml(dataset_data, f"{tmp_dir}/dataset_config.yaml")
+            shutil.copy(file_path, tmp_dir)
+
+            tar_file_path = os.path.join(tmp_dir, "tarfile")
+            with tarfile.open(tar_file_path, mode="w:gz") as tar:
+                tar.add(tmp_dir, arcname=os.path.basename("reference_dataset"))
+
+            self.api.upload(
+                endpoint=f"inference-pipelines/{inference_pipeline_id}",
+                file_path=tar_file_path,
+                object_name="tarfile",
+                body={},
+                storage_uri_key="referenceDatasetUri",
+                method="PUT",
+            )
+        print("Referece dataset uploaded!")
 
     def upload_reference_dataframe(
         self,

openlayer/api.py

@@ -181,7 +181,15 @@ def put_request(self, endpoint: str, body=None, files=None, data=None):
             data=data,
         )
 
-    def upload(self, endpoint: str, file_path: str, object_name: str = None, body=None):
+    def upload(
+        self,
+        endpoint: str,
+        file_path: str,
+        object_name: str = None,
+        body=None,
+        method: str = "POST",
+        storage_uri_key: str = "storageUri",
+    ):
         """Generic method to upload data to the default storage medium and create the
         appropriate resource in the backend.
         """
@@ -198,10 +206,18 @@ def upload(self, endpoint: str, file_path: str, object_name: str = None, body=No
             file_path=file_path,
             object_name=object_name,
             body=body,
+            method=method,
+            storage_uri_key=storage_uri_key,
         )
 
     def upload_blob_s3(
-        self, endpoint: str, file_path: str, object_name: str = None, body=None
+        self,
+        endpoint: str,
+        file_path: str,
+        object_name: str = None,
+        body=None,
+        method: str = "POST",
+        storage_uri_key: str = "storageUri",
     ):
         """Generic method to upload data to S3 storage and create the appropriate resource
         in the backend.
@@ -235,13 +251,22 @@ def upload_blob_s3(
                 )
 
         if res.ok:
-            body["storageUri"] = presigned_json["storageUri"]
-            return self.post_request(f"{endpoint}", body=body)
+            body[storage_uri_key] = presigned_json["storageUri"]
+            if method == "POST":
+                return self.post_request(f"{endpoint}", body=body)
+            elif method == "PUT":
+                return self.put_request(f"{endpoint}", body=body)
         else:
             self._raise_on_respose(res)
 
     def upload_blob_gcs(
-        self, endpoint: str, file_path: str, object_name: str = None, body=None
+        self,
+        endpoint: str,
+        file_path: str,
+        object_name: str = None,
+        body=None,
+        method: str = "POST",
+        storage_uri_key: str = "storageUri",
     ):
         """Generic method to upload data to Google Cloud Storage and create the
         appropriate resource in the backend.
@@ -265,13 +290,22 @@ def upload_blob_gcs(
                     timeout=constants.REQUESTS_TIMEOUT,
                 )
         if res.ok:
-            body["storageUri"] = presigned_json["storageUri"]
-            return self.post_request(f"{endpoint}", body=body)
+            body[storage_uri_key] = presigned_json["storageUri"]
+            if method == "POST":
+                return self.post_request(f"{endpoint}", body=body)
+            elif method == "PUT":
+                return self.put_request(f"{endpoint}", body=body)
         else:
             self._raise_on_respose(res)
 
     def upload_blob_azure(
-        self, endpoint: str, file_path: str, object_name: str = None, body=None
+        self,
+        endpoint: str,
+        file_path: str,
+        object_name: str = None,
+        body=None,
+        method: str = "POST",
+        storage_uri_key: str = "storageUri",
     ):
         """Generic method to upload data to Azure Blob Storage and create the
         appropriate resource in the backend.
@@ -298,12 +332,23 @@ def upload_blob_azure(
                     timeout=constants.REQUESTS_TIMEOUT,
                 )
         if res.ok:
-            body["storageUri"] = presigned_json["storageUri"]
-            return self.post_request(f"{endpoint}", body=body)
+            body[storage_uri_key] = presigned_json["storageUri"]
+            if method == "POST":
+                return self.post_request(f"{endpoint}", body=body)
+            elif method == "PUT":
+                return self.put_request(f"{endpoint}", body=body)
         else:
             self._raise_on_respose(res)
 
-    def transfer_blob(self, endpoint: str, file_path: str, object_name: str, body=None):
+    def transfer_blob(
+        self,
+        endpoint: str,
+        file_path: str,
+        object_name: str,
+        body=None,
+        method: str = "POST",
+        storage_uri_key: str = "storageUri",
+    ):
         """Generic method to transfer data to the openlayer folder and create the appropriate
         resource in the backend when using a local deployment.
         """
@@ -317,5 +362,8 @@ def transfer_blob(self, endpoint: str, file_path: str, object_name: str, body=No
         except OSError as exc:
             raise OpenlayerException(f"Directory {dir_path} cannot be created") from exc
         shutil.copyfile(file_path, blob_path)
-        body["storageUri"] = presigned_json["storageUri"]
-        return self.post_request(f"{endpoint}", body=body)
+        body[storage_uri_key] = presigned_json["storageUri"]
+        if method == "POST":
+            return self.post_request(f"{endpoint}", body=body)
+        elif method == "PUT":
+            return self.put_request(f"{endpoint}", body=body)

openlayer/inference_pipelines.py

@@ -10,6 +10,7 @@ class InferencePipeline:
     def __init__(self, json, upload, client, task_type):
         self._json = json
         self.id = json["id"]
+        self.project_id = json["projectId"]
         self.upload = upload
         self.client = client
         self.taskType = task_type