Merge pull request #9133 from OpenMined/tauquir/fix-model-size-repr

Fix model size calculation and improve repr
OpenMined · Aug 6, 2024 · 714964f · 714964f
2 parents 25ec85e + 7f9ae71
commit 714964f
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 3 deletions.
diff --git a/notebooks/experimental/enclaves/V3/V3-Enclave-Model-HostingSingle-Notebook.ipynb b/notebooks/experimental/enclaves/V3/V3-Enclave-Model-HostingSingle-Notebook.ipynb
@@ -496,6 +496,16 @@
     "do_canada_client.upload_model(model)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6b0340bc-ab78-4e49-8cd6-6d3e3ef370e8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "do_canada_client.models"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,

diff --git a/packages/syft/src/syft/client/datasite_client.py b/packages/syft/src/syft/client/datasite_client.py
@@ -34,8 +34,10 @@
 from ..service.user.roles import Roles
 from ..service.user.user import UserView
 from ..types.blob_storage import BlobFile
+from ..types.file import SyftFolder
 from ..types.uid import UID
 from ..util.misc_objs import HTMLObject
+from ..util.util import get_mb_serialized_size
 from ..util.util import get_mb_size
 from ..util.util import prompt_warning_message
 from .api import APIModule
@@ -157,7 +159,11 @@ def upload_model(self, model: CreateModel) -> SyftSuccess | SyftError:
 
                 asset.action_id = twin.id
                 asset.server_uid = self.id
-                model_size += get_mb_size(asset.data)
+                model_size += (
+                    asset.data.size_mb
+                    if isinstance(asset.data, SyftFolder)
+                    else get_mb_serialized_size(asset.data)
+                )
                 model_ref_action_ids.append(twin.id)
 
                 # Clear the Data and Mock , as they are uploaded as twin object

diff --git a/packages/syft/src/syft/service/model/model.py b/packages/syft/src/syft/service/model/model.py
@@ -462,7 +462,7 @@ def _coll_repr_(self) -> dict[str, Any]:
             "Name": self.name,
             "Assets": len(self.asset_list),
             "Url": self.url,
-            "Size": f"{self.mb_size} (MB)",
+            "Size": f"{self.mb_size:.2f} (MB)",
             "created at": str(self.created_at),
         }
 
@@ -494,7 +494,7 @@ def _repr_html_(self) -> Any:
         button_html = CopyIDButton(copy_text=str(self.id), max_width=60).to_html()
 
         attrs = {
-            "Mb Size": str(self.mb_size),
+            "Size": f"{self.mb_size:.2f} (MB)",
             "URL": str(self.url),
             "Created at": str(self.created_at),
             "Updated at": self.updated_at,