Dataset Upload with Force Replace

packages/syft/src/syft/client/datasite_client.py

@@ -24,6 +24,7 @@
 from ..service.dataset.dataset import Contributor
 from ..service.dataset.dataset import CreateAsset
 from ..service.dataset.dataset import CreateDataset
+from ..service.dataset.dataset import Dataset
 from ..service.dataset.dataset import _check_asset_must_contain_mock
 from ..service.migration.object_migration_state import MigrationData
 from ..service.request.request import Request
@@ -36,6 +37,7 @@
 from ..service.user.user import UserView
 from ..types.blob_storage import BlobFile
 from ..types.errors import SyftException
+from ..types.twin_object import TwinObject
 from ..types.uid import UID
 from ..util.misc_objs import HTMLObject
 from ..util.util import get_mb_size
@@ -98,27 +100,24 @@ class DatasiteClient(SyftClient):
     def __repr__(self) -> str:
         return f"<DatasiteClient: {self.name}>"
 
-    def upload_dataset(self, dataset: CreateDataset) -> SyftSuccess:
-        # relative
-        from ..types.twin_object import TwinObject
-
+    def upload_dataset(
+        self, dataset: CreateDataset, force_replace: bool = False
+    ) -> SyftSuccess | SyftError:
         if self.users is None:
             raise SyftException(public_message=f"can't get user service for {self}")
 
         user = self.users.get_current_user()
+
         if user.role not in [ServiceRole.DATA_OWNER, ServiceRole.ADMIN]:
             return SyftError(message="You don't have permission to upload datasets.")
+
         dataset = add_default_uploader(user, dataset)
 
         for i in range(len(dataset.asset_list)):
             asset = dataset.asset_list[i]
             dataset.asset_list[i] = add_default_uploader(user, asset)
 
-        # dataset._check_asset_must_contain_mock()
-        dataset_size: float = 0.0
-
         # TODO: Refactor so that object can also be passed to generate warnings
-
         self.api.connection = cast(ServerConnection, self.api.connection)
 
         metadata = self.api.connection.get_server_metadata(self.api.signing_key)
@@ -134,10 +133,27 @@ def upload_dataset(self, dataset: CreateDataset) -> SyftSuccess:
             )
             prompt_warning_message(message=message, confirm=True)
 
-        with tqdm(
-            total=len(dataset.asset_list), colour="green", desc="Uploading"
-        ) as pbar:
-            for asset in dataset.asset_list:
+        # check if the a dataset with the same name already exists
+        search_res = self.api.services.dataset.search(dataset.name)
+        dataset_exists: bool = len(search_res) > 0
+
+        if not dataset_exists:
+            return self._upload_new_dataset(dataset)
+
+        existed_dataset: Dataset = search_res[0]
+        if not force_replace:
+            return SyftError(
+                message=f"Dataset with name the '{dataset.name}' already exists. "
+                "Please use `upload_dataset(dataset, force_replace=True)` to overwrite."
+            )
+        return self._replace_dataset(existed_dataset, dataset)
+
+    def _upload_assets(self, assets: list[CreateAsset]) -> float | SyftError:
+        total_assets_size: float = 0.0
+
+        with tqdm(total=len(assets), colour="green", desc="Uploading") as pbar:
+            for asset in assets:
+                # create and save a twin object representing the asset to the blob store
                 try:
                     contains_empty: bool = asset.contains_empty()
                     twin = TwinObject(
@@ -163,16 +179,55 @@ def upload_dataset(self, dataset: CreateDataset) -> SyftSuccess:
 
                 asset.action_id = twin.id
                 asset.server_uid = self.id
-                dataset_size += get_mb_size(asset.data)
+
+                total_assets_size += get_mb_size(asset.data)
 
                 # Update the progress bar and set the dynamic description
                 pbar.set_description(f"Uploading: {asset.name}")
                 pbar.update(1)
 
-        dataset.mb_size = dataset_size
+        return total_assets_size
+
+    def _upload_new_dataset(self, dataset: CreateDataset) -> SyftSuccess | SyftError:
+        # upload the assets
+        total_assets_size: float | SyftError = self._upload_assets(dataset.asset_list)
+        if isinstance(total_assets_size, SyftError):
+            return total_assets_size
+
+        # check if the types of the assets are valid
+        dataset.mb_size = total_assets_size
         _check_asset_must_contain_mock(dataset.asset_list)
-        dataset.check()
-        return self.api.services.dataset.add(dataset=dataset)
+        valid = dataset.check()
+        if isinstance(valid, SyftError):
+            return valid
+
+        # add the dataset object to the dataset store
+        try:
+            return self.api.services.dataset.add(dataset=dataset)
+        except Exception as e:
+            return SyftError(message=f"Failed to upload dataset. {e}")
+
+    def _replace_dataset(
+        self, existed_dataset: Dataset, dataset: CreateDataset
+    ) -> SyftSuccess | SyftError:
+        # TODO: is there a way to check if the assets already exist and have not changed,
+        # since if uploading the assets will have different UIDs
+        total_assets_size: float | SyftError = self._upload_assets(dataset.asset_list)
+        if isinstance(total_assets_size, SyftError):
+            return total_assets_size
+
+        # check if the types of the assets are valid
+        dataset.mb_size = total_assets_size
+        valid = dataset.check()
+        _check_asset_must_contain_mock(dataset.asset_list)
+        if isinstance(valid, SyftError):
+            return valid
+        try:
+            return self.api.services.dataset.replace(
+                existed_dataset_uid=existed_dataset.id, dataset=dataset
+            )
+        except Exception as e:
+            return SyftError(message=f"Failed to replace dataset. {e}")
 
     def forgot_password(self, email: str) -> SyftSuccess | SyftError:
         return self.connection.forgot_password(email=email)

packages/syft/src/syft/orchestra.py

@@ -382,7 +382,7 @@ def launch(
             display(
                 SyftInfo(
                     message=f"You have launched a development server at http://{host}:{server_handle.port}."
-                    + "It is intended only for local use."
+                    + " It is intended only for local use."
                 )
             )
             return server_handle

packages/syft/src/syft/protocol/data_protocol.py

@@ -77,12 +77,27 @@ def handle_union_type_klass_name(type_klass_name: str) -> str:
     return type_klass_name
 
 
+def get_klass_or_canonical_name(arg: Any) -> str:
+    """Get the class name or canonical name of the object.
+
+    If the object is a subclass of SyftBaseObject, then use canonical name
+    to identify the object.
+
+    """
+
+    return (
+        arg.__canonical_name__  # If SyftBaseObject subclass, ignore class name
+        if hasattr(arg, "__canonical_name__")
+        else getattr(arg, "__name__", str(arg))
+    )
+
+
 def handle_annotation_repr_(annotation: type) -> str:
     """Handle typing representation."""
     origin = typing.get_origin(annotation)
     args = typing.get_args(annotation)
     if origin and args:
-        args_repr = ", ".join(getattr(arg, "__name__", str(arg)) for arg in args)
+        args_repr = ", ".join(get_klass_or_canonical_name(arg) for arg in args)
         origin_repr = getattr(origin, "__name__", str(origin))
 
         # Handle typing.Union and types.UnionType

packages/syft/src/syft/protocol/protocol_version.json

@@ -1,5 +1,16 @@
 {
   "1": {
     "release_name": "0.9.1.json"
+  },
+  "dev": {
+    "object_versions": {
+      "CreateAsset": {
+        "2": {
+          "version": 2,
+          "hash": "1637c1e35c8cb65c9d667ad91c824b2cc5cf5b281e93770915a68c4926e6a567",
+          "action": "add"
+        }
+      }
+    }
   }
 }

packages/syft/src/syft/service/dataset/dataset.py

@@ -98,7 +98,6 @@ class Asset(SyftObject):
     shape: tuple | None = None
     created_at: DateTime = DateTime.now()
     uploader: Contributor | None = None
-
     # _kwarg_name and _dataset_name are set by the UserCode.assets
     _kwarg_name: str | None = None
     _dataset_name: str | None = None
@@ -314,7 +313,7 @@ def check_mock(data: Any, mock: Any) -> bool:
 
 
 @serializable()
-class CreateAsset(SyftObject):
+class CreateAssetV1(SyftObject):
     # version
     __canonical_name__ = "CreateAsset"
     __version__ = SYFT_OBJECT_VERSION_1
@@ -336,6 +335,30 @@ class CreateAsset(SyftObject):
     __repr_attrs__ = ["name"]
     model_config = ConfigDict(validate_assignment=True, extra="forbid")
 
+
+@serializable()
+class CreateAsset(SyftObject):
+    # version
+    __canonical_name__ = "CreateAsset"
+    __version__ = SYFT_OBJECT_VERSION_2
+
+    id: UID | None = None  # type:ignore[assignment]
+    name: str
+    description: MarkdownDescription | None = None
+    contributors: set[Contributor] = set()
+    data_subjects: list[DataSubjectCreate] = []
+    server_uid: UID | None = None
+    action_id: UID | None = None
+    data: Any | None = None
+    mock: Any | None = None
+    shape: tuple | None = None
+    mock_is_real: bool = False
+    created_at: DateTime | None = None
+    uploader: Contributor | None = None
+
+    __repr_attrs__ = ["name"]
+    model_config = ConfigDict(validate_assignment=True, extra="forbid")
+
     def __init__(self, description: str | None = None, **data: Any) -> None:
         if isinstance(description, str):
             description = MarkdownDescription(text=description)
@@ -514,9 +537,7 @@ def _repr_html_(self) -> Any:
             """
         else:
             description_info_message = ""
-        if self.to_be_deleted:
-            return "This dataset has been marked for deletion. The underlying data may be not available."
-        return f"""
+        repr_html = f"""
             <div class='syft-dataset'>
             <h1>{self.name}</h1>
             <h2><strong><span class='pr-8'>Summary</span></strong></h2>
@@ -529,9 +550,19 @@ def _repr_html_(self) -> Any:
             </span></strong><a href='{self.url}'>{self.url}</a></p>
             <p class='paragraph-sm'><strong><span class='pr-8'>Contributors:</span></strong>
             To see full details call <strong>dataset.contributors</strong>.</p>
-            <h2><strong><span class='pr-8'>Assets</span></strong></h2>
-            {self.assets._repr_html_()}
-            """
+        """
+        if self.to_be_deleted:
+            repr_html += (
+                "<h2><strong><span class='pr-8'>"
+                "This dataset has been marked for deletion. The underlying data may be not available"
+                "</span></strong></h2>"
+            )
+        else:
+            repr_html += f"""
+                        <h2><strong><span class='pr-8'>Assets</span></strong></h2>
+                        {self.assets._repr_html_()}
+                    """
+        return repr_html
 
     def action_ids(self) -> list[UID]:
         return [asset.action_id for asset in self.asset_list if asset.action_id]
@@ -874,3 +905,10 @@ class DatasetUpdate(PartialSyftObject):
 
     name: str
     to_be_deleted: bool
+    asset_list: list[Asset]
+    contributors: set[Contributor]
+    citation: str
+    url: str
+    description: MarkdownDescription
+    uploader: Contributor
+    summary: str

packages/syft/src/syft/service/dataset/dataset_service.py

@@ -115,15 +115,16 @@ def get_all(
         context: AuthedServiceContext,
         page_size: int | None = 0,
         page_index: int | None = 0,
+        include_deleted: bool = False,
     ) -> DatasetPageView | DictTuple[str, Dataset]:
         """Get a Dataset"""
-        datasets = self.stash.get_all(context.credentials).unwrap()
+        datasets = self.stash.get_all(
+            context.credentials, include_deleted=include_deleted
+        ).unwrap()
 
         for dataset in datasets:
             if context.server is not None:
                 dataset.server_uid = context.server.id
-            if dataset.to_be_deleted:
-                datasets.remove(dataset)
 
         return _paginate_dataset_collection(
             datasets=datasets, page_size=page_size, page_index=page_index
@@ -141,9 +142,7 @@ def search(
         results = self.get_all(context)
 
         filtered_results = [
-            dataset
-            for dataset_name, dataset in results.items()
-            if name in dataset_name and not dataset.to_be_deleted
+            dataset for dataset_name, dataset in results.items() if name in dataset_name
         ]
 
         return _paginate_dataset_collection(
@@ -242,6 +241,28 @@ def delete(
         return_msg.append(f"Dataset with id '{uid}' successfully deleted.")
         return SyftSuccess(message="\n".join(return_msg))
 
+    @service_method(
+        path="dataset.replace",
+        name="replace",
+        roles=DATA_OWNER_ROLE_LEVEL,
+        unwrap_on_success=False,
+    )
+    def replace(
+        self,
+        context: AuthedServiceContext,
+        existed_dataset_uid: UID,
+        dataset: CreateDataset,
+    ) -> SyftSuccess:
+        dataset = dataset.to(Dataset, context=context)
+        dataset.id = existed_dataset_uid
+        self.stash.update(
+            credentials=context.credentials, dataset_update=dataset
+        ).unwrap()
+        # TODO: should we delete the existed dataset's asssets after force replace?
+        return SyftSuccess(
+            message=f"Dataset with id '{existed_dataset_uid}' successfully replaced."
+        )
+
 
 TYPE_TO_SERVICE[Dataset] = DatasetService
 SERVICE_TO_TYPES[DatasetService].update({Dataset})