entitysdk 0.6.2__tar.gz → 0.7.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: entitysdk
-Version: 0.6.2
+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.2 → 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:
@@ -357,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,
         )
 
@@ -457,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,
@@ -498,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]],
@@ -520,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.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.2 → 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.2 → 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.2 → 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.2 → 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.2 → 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(

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

entitysdk-0.7.0/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"),
+    }

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

@@ -0,0 +1,72 @@
+"""Staging functions for SimulationResult."""
+
+import logging
+from pathlib import Path
+
+from entitysdk.client import Client
+from entitysdk.downloaders.simulation_result import (
+    download_spike_report_file,
+    download_voltage_report_files,
+)
+from entitysdk.models import Simulation, SimulationResult
+from entitysdk.staging.simulation import stage_simulation
+from entitysdk.types import StrOrPath
+from entitysdk.utils.filesystem import create_dir
+from entitysdk.utils.io import load_json
+
+L = logging.getLogger(__name__)
+
+
+DEFAULT_REPORTS_DIR_NAME = "output"
+DEFAULT_SPIKE_FILE_NAME = "output/spikes.h5"
+DEFAULT_SIMULATION_CONFIG_FILENAME = "simulation_config.json"
+
+
+def stage_simulation_result(
+    client: Client,
+    *,
+    model: SimulationResult,
+    output_dir: StrOrPath,
+    simulation_config_file: StrOrPath | None = None,
+) -> Path:
+    """Stage a SimulationResult entity."""
+    output_dir: Path = create_dir(output_dir)
+
+    if simulation_config_file is None:
+        L.info(
+            "Simulation will be staged from simulation result's simulation_id %s",
+            model.simulation_id,
+        )
+        simulation_config_file = stage_simulation(
+            client,
+            model=client.get_entity(entity_id=model.simulation_id, entity_type=Simulation),
+            output_dir=output_dir,
+        )
+    else:
+        L.info(
+            "External simulation config provided at %s. Outputs will be staged relative to it.",
+            simulation_config_file,
+        )
+
+    config: dict = load_json(simulation_config_file)
+    reports_dir, spikes_file = _get_output_paths(config, Path(simulation_config_file).parent)
+    create_dir(reports_dir)
+
+    download_spike_report_file(
+        client,
+        model=model,
+        output_path=spikes_file,
+    )
+    download_voltage_report_files(
+        client,
+        model=model,
+        output_dir=reports_dir,
+    )
+
+    return Path(simulation_config_file)
+
+
+def _get_output_paths(config: dict, output_dir: Path) -> tuple[Path, Path]:
+    reports_dir = output_dir / config["output"]["output_dir"]
+    spikes_file = output_dir / config["output"]["spikes_file"]
+    return reports_dir, spikes_file

{entitysdk-0.6.2 → entitysdk-0.7.0}/src/entitysdk/types.py

@@ -123,3 +123,9 @@ class SimulationExecutionStatus(StrEnum):
     running = auto()
     done = auto()
     error = auto()
+
+
+class ContentType(StrEnum):
+    """Content types."""
+
+    json = "application/json"

entitysdk-0.7.0/src/entitysdk/utils/io.py

@@ -0,0 +1,16 @@
+"""IO utilities."""
+
+import json
+from pathlib import Path
+
+from entitysdk.types import StrOrPath
+
+
+def write_json(data: dict, path: StrOrPath, **json_kwargs) -> None:
+    """Write dictionary to file as JSON."""
+    Path(path).write_text(json.dumps(data, **json_kwargs))
+
+
+def load_json(path: StrOrPath) -> dict:
+    """Load JSON file to dict."""
+    return json.loads(Path(path).read_bytes())

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: entitysdk
-Version: 0.6.2
+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>