entitysdk 0.6.2__tar.gz → 0.7.1__tar.gz

{entitysdk-0.6.2/src/entitysdk.egg-info → entitysdk-0.7.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: entitysdk
-Version: 0.6.2
+Version: 0.7.1
 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.2 → entitysdk-0.7.1}/pyproject.toml

@@ -46,6 +46,7 @@ addopts = [
 [tool.ruff]
 line-length = 100
 target-version = "py311"
+include = ["pyproject.toml", "src/**/*.py", "tests/**/*.py", "examples/**/*.py"]
 
 [tool.ruff.lint]
 select = [

{entitysdk-0.6.2 → entitysdk-0.7.1}/src/entitysdk/client.py

@@ -1,9 +1,10 @@
 """Identifiable SDK client."""
 
+import concurrent.futures
 import io
 import os
 from pathlib import Path
-from typing import Any, cast
+from typing import Any, TypeVar, cast
 
 import httpx
 
@@ -14,7 +15,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 +25,8 @@ from entitysdk.util import (
 )
 from entitysdk.utils.asset import filter_assets
 
+TEntity = TypeVar("TEntity", bound=Entity)
+
 
 class Client:
     """Client for entitysdk."""
@@ -94,9 +97,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,7 +217,7 @@ class Client:
         file_content_type: str,
         file_name: str | None = None,
         file_metadata: dict | None = None,
-        asset_label: str | None = None,
+        asset_label: str,
         project_context: ProjectContext | None = None,
     ) -> Asset:
         """Upload asset to an existing entity's endpoint from a file path."""
@@ -250,7 +253,7 @@ class Client:
         file_name: str,
         file_content_type: str,
         file_metadata: dict | None = None,
-        asset_label: str | None = None,
+        asset_label: str,
         project_context: ProjectContext | None = None,
     ) -> Asset:
         """Upload asset to an existing entity's endpoint from a file-like object."""
@@ -284,7 +287,7 @@ class Client:
         name: str,
         paths: dict[os.PathLike, os.PathLike],
         metadata: dict | None = None,
-        label: str | None = None,
+        label: str,
         project_context: ProjectContext | None = None,
     ) -> Asset:
         """Attach directory to an entity from with a group of paths."""
@@ -347,8 +350,9 @@ class Client:
         output_path: os.PathLike,
         project_context: ProjectContext | None = None,
         ignore_directory_name: bool = False,
+        max_concurrent: int = 1,
     ) -> list[Path]:
-        """List directory existing entity's endpoint from a directory path."""
+        """Download directory of assets."""
         output_path = Path(output_path)
 
         if output_path.exists() and output_path.is_file():
@@ -357,34 +361,35 @@ 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,
         )
 
-        paths = []
-        for path in contents.files:
-            paths.append(
+        if max_concurrent == 1:
+            paths = [
                 self.download_file(
                     entity_id=entity_id,
                     entity_type=entity_type,
@@ -393,7 +398,24 @@ class Client:
                     asset_path=path,
                     project_context=context,
                 )
-            )
+                for path in contents.files
+            ]
+        else:
+            with concurrent.futures.ThreadPoolExecutor(max_workers=max_concurrent) as executor:
+                futures = [
+                    executor.submit(
+                        self.download_file,
+                        entity_id=entity_id,
+                        entity_type=entity_type,
+                        asset_id=asset if asset else asset_id,
+                        output_path=output_path / path,
+                        asset_path=path,
+                        project_context=context,
+                    )
+                    for path in contents.files
+                ]
+                result = concurrent.futures.wait(futures)
+                paths = [res.result() for res in result.done]
 
         return paths
 
@@ -457,6 +479,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,
@@ -498,6 +521,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]],
@@ -520,9 +548,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)
@@ -587,7 +615,7 @@ class Client:
 
         Note: This operation is not atomic. Deletion can succeed and upload can fail.
         """
-        self.delete_asset(
+        deleted_asset = self.delete_asset(
             entity_id=entity_id,
             entity_type=entity_type,
             asset_id=asset_id,
@@ -601,4 +629,5 @@ class Client:
             file_name=file_name,
             file_metadata=file_metadata,
             project_context=project_context,
+            asset_label=deleted_asset.label,
         )

{entitysdk-0.6.2 → entitysdk-0.7.1}/src/entitysdk/core.py

@@ -187,7 +187,7 @@ def upload_asset_directory(
     name: str,
     paths: dict[Path, Path],
     metadata: dict | None = None,
-    label: str | None = None,
+    label: str,
     project_context: ProjectContext,
     token: str,
     http_client: httpx.Client | None = None,

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

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

entitysdk-0.7.1/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.2 → entitysdk-0.7.1}/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.2 → entitysdk-0.7.1}/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.2 → entitysdk-0.7.1}/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.1/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.1/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.2 → entitysdk-0.7.1}/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.2 → entitysdk-0.7.1}/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(
@@ -62,7 +65,7 @@ class Asset(Identifiable):
         dict,
         Field(description="Asset json metadata."),
     ] = {}
-    label: Annotated[str | None, Field(description="Optional asset label.")] = None
+    label: Annotated[str, Field(description="Asset label.")]
 
 
 class LocalAssetMetadata(BaseModel):
@@ -93,7 +96,7 @@ class LocalAssetMetadata(BaseModel):
             description="The metadata of the asset.",
         ),
     ] = None
-    label: Annotated[str | None, Field(description="Optional asset label.")] = None
+    label: Annotated[str, Field(description="Optional asset label.")]
 
 
 class DetailedFile(BaseModel):

entitysdk-0.7.1/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.1/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.1/src/entitysdk/staging/circuit.py

@@ -0,0 +1,44 @@
+"""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,
+            "label": "sonata_circuit",
+        },
+    ).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

entitysdk-0.7.1/src/entitysdk/staging/simulation.py

@@ -0,0 +1,142 @@
+"""Staging functions for Simulation."""
+
+import logging
+from copy import deepcopy
+from pathlib import Path
+
+from entitysdk.client import Client
+from entitysdk.downloaders.simulation import (
+    download_node_sets_file,
+    download_simulation_config_content,
+    download_spike_replay_files,
+)
+from entitysdk.exception import StagingError
+from entitysdk.models import Circuit, Simulation
+from entitysdk.staging.circuit import stage_circuit
+from entitysdk.types import StrOrPath
+from entitysdk.utils.filesystem import create_dir
+from entitysdk.utils.io import write_json
+
+L = logging.getLogger(__name__)
+
+DEFAULT_NODE_SETS_FILENAME = "node_sets.json"
+DEFAULT_SIMULATION_CONFIG_FILENAME = "simulation_config.json"
+DEFAULT_CIRCUIT_DIR = "circuit"
+
+
+def stage_simulation(
+    client: Client,
+    *,
+    model: Simulation,
+    output_dir: StrOrPath,
+    circuit_config_path: Path | None = None,
+    override_results_dir: Path | None = None,
+) -> Path:
+    """Stage a simulation entity into output_dir.
+
+    Args:
+        client: The client to use to stage the simulation.
+        model: The simulation entity to stage.
+        output_dir: The directory to stage the simulation into.
+        circuit_config_path: The path to the circuit config file.
+            If not provided, the circuit will be staged from metadata.
+        override_results_dir: Directory to update the simulation config section to point to.
+
+    Returns:
+        The path to the staged simulation config file.
+    """
+    output_dir = create_dir(output_dir).resolve()
+
+    simulation_config: dict = download_simulation_config_content(client, model=model)
+    node_sets_file: Path = download_node_sets_file(
+        client,
+        model=model,
+        output_path=output_dir / DEFAULT_NODE_SETS_FILENAME,
+    )
+    spike_paths: list[Path] = download_spike_replay_files(
+        client,
+        model=model,
+        output_dir=output_dir,
+    )
+    if circuit_config_path is None:
+        L.info(
+            "Circuit config path was not provided. Circuit is going to be staged from metadata. "
+            "Circuit id to be staged: %s"
+        )
+        circuit_config_path = stage_circuit(
+            client,
+            model=client.get_entity(
+                entity_id=model.entity_id,
+                entity_type=Circuit,
+            ),
+            output_dir=create_dir(output_dir / DEFAULT_CIRCUIT_DIR),
+        )
+
+    transformed_simulation_config: dict = _transform_simulation_config(
+        simulation_config=simulation_config,
+        circuit_config_path=circuit_config_path,
+        node_sets_path=node_sets_file,
+        spike_paths=spike_paths,
+        output_dir=output_dir,
+        override_results_dir=override_results_dir,
+    )
+
+    output_simulation_config_file = output_dir / DEFAULT_SIMULATION_CONFIG_FILENAME
+
+    write_json(
+        data=transformed_simulation_config,
+        path=output_simulation_config_file,
+    )
+
+    L.info("Staged Simulation %s at %s", model.id, output_dir)
+
+    return output_simulation_config_file
+
+
+def _transform_simulation_config(
+    simulation_config: dict,
+    circuit_config_path: Path,
+    node_sets_path: Path,
+    spike_paths: list[Path],
+    output_dir: Path,
+    override_results_dir: Path | None,
+) -> dict:
+    return simulation_config | {
+        "network": str(circuit_config_path),
+        "node_sets_file": str(node_sets_path.relative_to(output_dir)),
+        "inputs": _transform_inputs(simulation_config["inputs"], spike_paths),
+        "output": _transform_output(simulation_config["output"], override_results_dir),
+    }
+
+
+def _transform_inputs(inputs: dict, spike_paths: list[Path]) -> dict:
+    expected_spike_filenames = {p.name for p in spike_paths}
+
+    transformed_inputs = deepcopy(inputs)
+    for values in transformed_inputs.values():
+        if values["input_type"] == "spikes":
+            path = Path(values["spike_file"]).name
+
+            if path not in expected_spike_filenames:
+                raise StagingError(
+                    f"Spike file name in config is not present in spike asset file names.\n"
+                    f"Config file name: {path}\n"
+                    f"Asset file names: {expected_spike_filenames}"
+                )
+
+            values["spike_file"] = str(path)
+            L.debug("Spike file %s -> %s", values["spike_file"], path)
+
+    return transformed_inputs
+
+
+def _transform_output(output: dict, override_results_dir: StrOrPath | None) -> dict:
+    if override_results_dir is None:
+        return output
+
+    path = Path(override_results_dir)
+
+    return {
+        "output_dir": str(path),
+        "spikes_file": str(path / "spikes.h5"),
+    }