entitysdk 0.6.1__tar.gz → 0.7.0__tar.gz

{entitysdk-0.6.1/src/entitysdk.egg-info → entitysdk-0.7.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: entitysdk
-Version: 0.6.1
+Version: 0.7.0
 Summary: Python library for interacting with the entitycore service
 Author-email: Open Brain Institute <info@openbraininstitute.org>
 Maintainer-email: Open Brain Institute <info@openbraininstitute.org>

{entitysdk-0.6.1 → entitysdk-0.7.0}/examples/02_morphology.ipynb

@@ -276,7 +276,8 @@
     "        entity_id=registered.id,\n",
     "        entity_type=ReconstructionMorphology,\n",
     "        file_path=file1,\n",
-    "        file_content_type=\"application/h5\",\n",
+    "        file_content_type=\"application/x-hdf5\",\n",
+    "        asset_label=\"hdf5\",\n",
     "    )\n",
     "    rprint(asset1)\n",
     "\n",
@@ -289,6 +290,7 @@
     "        file_content=buffer,\n",
     "        file_name=\"buffer.h5\",\n",
     "        file_content_type=\"application/swc\",\n",
+    "        asset_label=\"swc\",\n",
     "    )\n",
     "    rprint(asset2)"
    ]
@@ -408,8 +410,22 @@
   }
  ],
  "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
   "language_info": {
-   "name": "python"
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.10"
   }
  },
  "nbformat": 4,

{entitysdk-0.6.1 → entitysdk-0.7.0}/src/entitysdk/client.py

@@ -3,7 +3,7 @@
 import io
 import os
 from pathlib import Path
-from typing import Any, cast
+from typing import Any, TypeVar, cast
 
 import httpx
 
@@ -14,7 +14,7 @@ from entitysdk.models.asset import Asset, DetailedFileList, LocalAssetMetadata
 from entitysdk.models.core import Identifiable
 from entitysdk.models.entity import Entity
 from entitysdk.result import IteratorResult
-from entitysdk.schemas.asset import DownloadedAsset
+from entitysdk.schemas.asset import DownloadedAssetFile
 from entitysdk.token_manager import TokenFromValue, TokenManager
 from entitysdk.types import ID, DeploymentEnvironment, Token
 from entitysdk.util import (
@@ -24,6 +24,8 @@ from entitysdk.util import (
 )
 from entitysdk.utils.asset import filter_assets
 
+TEntity = TypeVar("TEntity", bound=Entity)
+
 
 class Client:
     """Client for entitysdk."""
@@ -94,9 +96,9 @@ class Client:
         self,
         entity_id: ID,
         *,
-        entity_type: type[Identifiable],
+        entity_type: type[TEntity],
         project_context: ProjectContext | None = None,
-    ) -> Identifiable:
+    ) -> TEntity:
         """Get entity from resource id.
 
         Args:
@@ -214,6 +216,7 @@ class Client:
         file_content_type: str,
         file_name: str | None = None,
         file_metadata: dict | None = None,
+        asset_label: str | None = None,
         project_context: ProjectContext | None = None,
     ) -> Asset:
         """Upload asset to an existing entity's endpoint from a file path."""
@@ -229,6 +232,7 @@ class Client:
             file_name=file_name or path.name,
             content_type=file_content_type,
             metadata=file_metadata,
+            label=asset_label,
         )
         return core.upload_asset_file(
             url=url,
@@ -248,6 +252,7 @@ class Client:
         file_name: str,
         file_content_type: str,
         file_metadata: dict | None = None,
+        asset_label: str | None = None,
         project_context: ProjectContext | None = None,
     ) -> Asset:
         """Upload asset to an existing entity's endpoint from a file-like object."""
@@ -261,6 +266,7 @@ class Client:
             file_name=file_name,
             content_type=file_content_type,
             metadata=file_metadata or {},
+            label=asset_label,
         )
         context = self._required_user_context(override_context=project_context)
         return core.upload_asset_content(
@@ -353,28 +359,30 @@ class Client:
 
         context = self._optional_user_context(override_context=project_context)
 
-        asset = None
+        asset = cast(Asset, asset_id) if isinstance(asset_id, Asset) else None
+
         if not ignore_directory_name:
-            asset_endpoint = route.get_assets_endpoint(
-                api_url=self.api_url,
-                entity_type=entity_type,
-                entity_id=entity_id,
-                asset_id=asset_id,
-            )
-            asset = core.get_entity(
-                asset_endpoint,
-                entity_type=Asset,
-                project_context=context,
-                http_client=self._http_client,
-                token=self._token_manager.get_token(),
-            )
+            if asset is None:
+                asset_endpoint = route.get_assets_endpoint(
+                    api_url=self.api_url,
+                    entity_type=entity_type,
+                    entity_id=cast(ID, entity_id),
+                    asset_id=asset_id,
+                )
+                asset = core.get_entity(
+                    asset_endpoint,
+                    entity_type=Asset,
+                    project_context=context,
+                    http_client=self._http_client,
+                    token=self._token_manager.get_token(),
+                )
 
             output_path /= asset.path
 
         contents = self.list_directory(
             entity_id=entity_id,
             entity_type=entity_type,
-            asset_id=asset_id,
+            asset_id=asset_id if isinstance(asset_id, ID) else asset.id,
             project_context=project_context,
         )
 
@@ -453,6 +461,7 @@ class Client:
             Output file path.
         """
         context = self._optional_user_context(override_context=project_context)
+
         asset_endpoint = route.get_assets_endpoint(
             api_url=self.api_url,
             entity_type=entity_type,
@@ -494,6 +503,11 @@ class Client:
             token=self._token_manager.get_token(),
         )
 
+    @staticmethod
+    def select_assets(entity: Entity, selection: dict) -> IteratorResult:
+        """Select assets from entity based on selection."""
+        return IteratorResult(filter_assets(entity.assets, selection))
+
     def download_assets(
         self,
         entity_or_id: Entity | tuple[ID, type[Entity]],
@@ -516,9 +530,9 @@ class Client:
                     project_context=context,
                 )
 
-            return DownloadedAsset(
+            return DownloadedAssetFile(
                 asset=asset,
-                output_path=path,
+                path=path,
             )
 
         context = self._optional_user_context(override_context=project_context)

{entitysdk-0.6.1 → entitysdk-0.7.0}/src/entitysdk/core.py

@@ -1,6 +1,7 @@
 """Core SDK operations."""
 
 import io
+import logging
 import os
 from collections.abc import Iterator
 from pathlib import Path
@@ -10,11 +11,13 @@ import httpx
 
 from entitysdk import serdes
 from entitysdk.common import ProjectContext
+from entitysdk.exception import EntitySDKError
 from entitysdk.models.asset import Asset, DetailedFileList, LocalAssetMetadata
 from entitysdk.models.core import Identifiable
 from entitysdk.result import IteratorResult
 from entitysdk.util import make_db_api_request, stream_paginated_request
 
+L = logging.getLogger(__name__)
 TIdentifiable = TypeVar("TIdentifiable", bound=Identifiable)
 
 
@@ -170,6 +173,7 @@ def upload_asset_content(
         url=url,
         method="POST",
         files=files,
+        data={"label": asset_metadata.label} if asset_metadata.label else None,
         project_context=project_context,
         token=token,
         http_client=http_client,
@@ -192,7 +196,7 @@ def upload_asset_directory(
     for concrete_path in paths.values():
         if not concrete_path.exists():
             msg = f"Path {concrete_path} does not exist"
-            raise Exception(msg)
+            raise EntitySDKError(msg)
 
     response = make_db_api_request(
         url=url,
@@ -214,11 +218,20 @@ def upload_asset_directory(
         failed = {}
         for path, url in to_upload.items():
             with open(paths[Path(path)], "rb") as fd:
-                response = http_client.request(
-                    method="PUT", url=url, content=fd, follow_redirects=True
-                )
-            if response.status_code != 200:
-                failed[path] = url
+                try:
+                    response = http_client.request(
+                        method="PUT",
+                        url=url,
+                        content=fd,
+                        follow_redirects=True,
+                        timeout=20,
+                    )
+                except httpx.HTTPError:
+                    L.exception("Upload failed, will retry again")
+                    failed[path] = url
+                else:
+                    if response.status_code != 200:
+                        failed[path] = url
         return failed
 
     to_upload = js["files"]
@@ -228,7 +241,7 @@ def upload_asset_directory(
             break
 
     if to_upload:
-        raise Exception(f"Uploading these files failed: {to_upload}")
+        raise EntitySDKError(f"Uploading these files failed: {to_upload}")
 
     return serdes.deserialize_model(js["asset"], Asset)
 

entitysdk-0.7.0/src/entitysdk/dependencies/__init__.py

@@ -0,0 +1 @@
+"""Dependencies."""

entitysdk-0.7.0/src/entitysdk/dependencies/entity.py

@@ -0,0 +1,18 @@
+"""Entity dependencies."""
+
+from entitysdk.exception import DependencyError
+from entitysdk.models.entity import Entity
+
+
+def ensure_has_id(model: Entity) -> Entity:
+    """Ensure entity has id."""
+    if model.id is None:
+        raise DependencyError(f"Model has no id: {repr(model)}")
+    return model
+
+
+def ensure_has_assets(model: Entity) -> Entity:
+    """Ensure entity has assets."""
+    if not model.assets:
+        raise DependencyError(f"Model has no assets: {repr(model)}")
+    return model

{entitysdk-0.6.1 → entitysdk-0.7.0}/src/entitysdk/downloaders/emodel.py

@@ -26,4 +26,4 @@ def download_hoc(
         output_path=output_dir,
     ).one()
 
-    return asset.output_path
+    return asset.path

{entitysdk-0.6.1 → entitysdk-0.7.0}/src/entitysdk/downloaders/ion_channel_model.py

@@ -26,4 +26,4 @@ def download_ion_channel_mechanism(
         output_path=output_dir,
     ).one()
 
-    return asset.output_path
+    return asset.path

{entitysdk-0.6.1 → entitysdk-0.7.0}/src/entitysdk/downloaders/morphology.py

@@ -33,4 +33,4 @@ def download_morphology(
         output_path=output_dir,
     ).one()
 
-    return asset.output_path
+    return asset.path

entitysdk-0.7.0/src/entitysdk/downloaders/simulation.py

@@ -0,0 +1,78 @@
+"""Downloading functions for Simulation."""
+
+import json
+import logging
+from pathlib import Path
+from typing import cast
+
+from entitysdk.client import Client
+from entitysdk.dependencies.entity import ensure_has_assets, ensure_has_id
+from entitysdk.models import Simulation
+from entitysdk.types import ID
+
+L = logging.getLogger(__name__)
+
+
+def download_simulation_config_content(client: Client, *, model: Simulation) -> dict:
+    """Download the the simulation config json into a dictionary."""
+    ensure_has_id(model)
+    ensure_has_assets(model)
+
+    asset = client.select_assets(
+        model,
+        selection={"label": "sonata_simulation_config"},
+    ).one()
+
+    json_content: bytes = client.download_content(
+        entity_id=cast(ID, model.id),
+        entity_type=Simulation,
+        asset_id=asset.id,
+    )
+
+    return json.loads(json_content)
+
+
+def download_node_sets_file(client: Client, *, model: Simulation, output_path: Path) -> Path:
+    """Download the node sets file from simulation's assets."""
+    ensure_has_id(model)
+    ensure_has_assets(model)
+
+    asset = client.select_assets(
+        model,
+        selection={"label": "custom_node_sets"},
+    ).one()
+
+    path = client.download_file(
+        entity_id=cast(ID, model.id),
+        entity_type=Simulation,
+        asset_id=asset,
+        output_path=output_path,
+    )
+
+    L.info("Node sets file downloaded at %s", path)
+
+    return path
+
+
+def download_spike_replay_files(
+    client: Client, *, model: Simulation, output_dir: Path
+) -> list[Path]:
+    """Download the spike replay files from simualtion's assets."""
+    ensure_has_id(model)
+    ensure_has_assets(model)
+
+    assets = client.select_assets(model, selection={"label": "replay_spikes"}).all()
+
+    spike_files: list[Path] = [
+        client.download_file(
+            entity_id=cast(ID, model.id),
+            entity_type=Simulation,
+            asset_id=asset,
+            output_path=output_dir / asset.path,
+        )
+        for asset in assets
+    ]
+
+    L.info("Downloaded %d spike replay files: %s", len(spike_files), spike_files)
+
+    return spike_files

entitysdk-0.7.0/src/entitysdk/downloaders/simulation_result.py

@@ -0,0 +1,61 @@
+"""Downloading functions for SimulationResult."""
+
+import logging
+from pathlib import Path
+from typing import cast
+
+from entitysdk.client import Client
+from entitysdk.dependencies.entity import ensure_has_assets, ensure_has_id
+from entitysdk.models import SimulationResult
+from entitysdk.types import ID
+
+L = logging.getLogger(__name__)
+
+
+def download_spike_report_file(
+    client: Client, *, model: SimulationResult, output_path: Path
+) -> Path:
+    """Download spike report file from SimulationResult entity."""
+    ensure_has_id(model)
+    ensure_has_assets(model)
+
+    asset = client.select_assets(
+        model,
+        selection={"label": "spike_report"},
+    ).one()
+
+    path = client.download_file(
+        entity_id=cast(ID, model.id),
+        entity_type=SimulationResult,
+        asset_id=asset,
+        output_path=output_path / asset.path if output_path.is_dir() else output_path,
+    )
+    L.info("Spike report file downloaded at %s", path)
+    return path
+
+
+def download_voltage_report_files(
+    client: Client, *, model: SimulationResult, output_dir: Path
+) -> list[Path]:
+    """Download voltage report files from SimulationResult entity."""
+    ensure_has_id(model)
+    ensure_has_assets(model)
+
+    assets = client.select_assets(
+        model,
+        selection={"label": "voltage_report"},
+    ).all()
+
+    files: list[Path] = [
+        client.download_file(
+            entity_id=cast(ID, model.id),
+            entity_type=SimulationResult,
+            asset_id=asset,
+            output_path=output_dir / asset.path,
+        )
+        for asset in assets
+    ]
+
+    L.info("Downloaded voltage report files: %s", files)
+
+    return files

{entitysdk-0.6.1 → entitysdk-0.7.0}/src/entitysdk/exception.py

@@ -11,3 +11,11 @@ class RouteNotFoundError(EntitySDKError):
 
 class IteratorResultError(EntitySDKError):
     """Raised when the result of an iterator is not as expected."""
+
+
+class DependencyError(EntitySDKError):
+    """Raised when a dependency check fails."""
+
+
+class StagingError(EntitySDKError):
+    """Raised when a staging operation has failed."""

{entitysdk-0.6.1 → entitysdk-0.7.0}/src/entitysdk/models/asset.py

@@ -8,11 +8,14 @@ from pydantic import ConfigDict, Field
 
 from entitysdk.models.base import BaseModel
 from entitysdk.models.core import Identifiable
+from entitysdk.types import ID
 
 
 class Asset(Identifiable):
     """Asset."""
 
+    id: ID
+
     path: Annotated[
         str,
         Field(
@@ -93,6 +96,7 @@ class LocalAssetMetadata(BaseModel):
             description="The metadata of the asset.",
         ),
     ] = None
+    label: Annotated[str | None, Field(description="Optional asset label.")] = None
 
 
 class DetailedFile(BaseModel):

entitysdk-0.7.0/src/entitysdk/schemas/asset.py

@@ -0,0 +1,20 @@
+"""Asset related schemas."""
+
+from pathlib import Path
+
+from entitysdk.models.asset import Asset
+from entitysdk.schemas.base import Schema
+
+
+class DownloadedAssetFile(Schema):
+    """Downloaded asset file."""
+
+    asset: Asset
+    path: Path
+
+
+class DownloadedAssetContent(Schema):
+    """Downloaded asset content."""
+
+    asset: Asset
+    content: bytes

entitysdk-0.7.0/src/entitysdk/staging/__init__.py

@@ -0,0 +1,7 @@
+"""Staging functions."""
+
+from entitysdk.staging.circuit import stage_circuit
+from entitysdk.staging.simulation import stage_simulation
+from entitysdk.staging.simulation_result import stage_simulation_result
+
+__all__ = ["stage_circuit", "stage_simulation", "stage_simulation_result"]

entitysdk-0.7.0/src/entitysdk/staging/circuit.py

@@ -0,0 +1,43 @@
+"""Staging functions for Circuit."""
+
+import logging
+from pathlib import Path
+from typing import cast
+
+from entitysdk.client import Client
+from entitysdk.dependencies.entity import ensure_has_assets, ensure_has_id
+from entitysdk.models import Circuit
+from entitysdk.types import ID
+
+L = logging.getLogger(__name__)
+
+
+def stage_circuit(client: Client, *, model: Circuit, output_dir: Path) -> Path:
+    """Stage a Circuit directory into output_dir."""
+    ensure_has_id(model)
+    ensure_has_assets(model)
+
+    asset = client.select_assets(
+        model,
+        selection={
+            "content_type": "application/vnd.directory",
+            "is_directory": True,
+        },
+    ).one()
+
+    paths = client.download_directory(
+        entity_id=cast(ID, model.id),
+        entity_type=Circuit,
+        asset_id=asset,
+        output_path=output_dir,
+        ignore_directory_name=True,
+    )
+
+    L.debug("Downloaded circuit %s paths: %s", model.id, paths)
+
+    circuit_config_path = output_dir / "circuit_config.json"
+    assert circuit_config_path in paths
+
+    L.info("Circuit %s staged at %s", model.id, circuit_config_path)
+
+    return circuit_config_path