Skip to content

Commit

Permalink
add support for simulation runs
Browse files Browse the repository at this point in the history
  • Loading branch information
lpereiracgn committed Feb 28, 2025
1 parent 01e65f4 commit 9e087ca
Show file tree
Hide file tree
Showing 6 changed files with 452 additions and 0 deletions.
2 changes: 2 additions & 0 deletions cognite/client/_api/simulators/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@

from cognite.client._api.simulators.integrations import SimulatorIntegrationsAPI
from cognite.client._api.simulators.models import SimulatorModelsAPI
from cognite.client._api.simulators.runs import SimulatorRunsAPI
from cognite.client._api_client import APIClient
from cognite.client._constants import DEFAULT_LIMIT_READ
from cognite.client.data_classes.simulators.simulators import Simulator, SimulatorList
Expand All @@ -22,6 +23,7 @@ def __init__(self, config: ClientConfig, api_version: str | None, cognite_client
super().__init__(config, api_version, cognite_client)
self.integrations = SimulatorIntegrationsAPI(config, api_version, cognite_client)
self.models = SimulatorModelsAPI(config, api_version, cognite_client)
self.runs = SimulatorRunsAPI(config, api_version, cognite_client)
self._warning = FeaturePreviewWarning(
api_maturity="General Availability", sdk_maturity="alpha", feature_name="Simulators"
)
Expand Down
217 changes: 217 additions & 0 deletions cognite/client/_api/simulators/runs.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,217 @@
from __future__ import annotations

from collections.abc import Iterator, Sequence
from typing import TYPE_CHECKING, overload

from cognite.client._api_client import APIClient
from cognite.client._constants import DEFAULT_LIMIT_READ
from cognite.client.data_classes.simulators.filters import SimulatorRunsFilter
from cognite.client.data_classes.simulators.runs import SimulationRun, SimulationRunWrite, SimulatorRunsList
from cognite.client.utils._experimental import FeaturePreviewWarning
from cognite.client.utils._validation import assert_type

if TYPE_CHECKING:
from cognite.client import CogniteClient
from cognite.client.config import ClientConfig


class SimulatorRunsAPI(APIClient):
_RESOURCE_PATH = "/simulators/runs"
_RESOURCE_PATH_RUN = "/simulators/run"

def __init__(self, config: ClientConfig, api_version: str | None, cognite_client: CogniteClient) -> None:
super().__init__(config, api_version, cognite_client)
self._warning = FeaturePreviewWarning(
api_maturity="General Availability", sdk_maturity="alpha", feature_name="Simulators"
)

def __iter__(self) -> Iterator[SimulationRun]:
"""Iterate over simulation runs
Fetches simulation runs as they are iterated over, so you keep a limited number of simulation runs in memory.
Returns:
Iterator[SimulationRun]: yields simulation runs one by one.
"""
return self()

@overload
def __call__(
self,
chunk_size: int,
limit: int | None = None,
status: str | None = None,
run_type: str | None = None,
model_external_ids: Sequence[str] | None = None,
simulator_integration_external_ids: Sequence[str] | None = None,
simulator_external_ids: Sequence[str] | None = None,
routine_external_ids: Sequence[str] | None = None,
routine_revision_external_ids: Sequence[str] | None = None,
model_revision_external_ids: Sequence[str] | None = None,
) -> Iterator[SimulationRun]: ...

@overload
def __call__(
self,
chunk_size: None = None,
limit: int | None = None,
status: str | None = None,
run_type: str | None = None,
model_external_ids: Sequence[str] | None = None,
simulator_integration_external_ids: Sequence[str] | None = None,
simulator_external_ids: Sequence[str] | None = None,
routine_external_ids: Sequence[str] | None = None,
routine_revision_external_ids: Sequence[str] | None = None,
model_revision_external_ids: Sequence[str] | None = None,
) -> Iterator[SimulationRun]: ...

def __call__(
self,
chunk_size: int | None = None,
limit: int | None = None,
status: str | None = None,
run_type: str | None = None,
model_external_ids: Sequence[str] | None = None,
simulator_integration_external_ids: Sequence[str] | None = None,
simulator_external_ids: Sequence[str] | None = None,
routine_external_ids: Sequence[str] | None = None,
routine_revision_external_ids: Sequence[str] | None = None,
model_revision_external_ids: Sequence[str] | None = None,
) -> Iterator[SimulationRun] | Iterator[SimulatorRunsList]:
"""Iterate over simulation runs
Fetches simulation runs as they are iterated over, so you keep a limited number of simulation runs in memory.
Args:
chunk_size (int | None): Number of simulation runs to return in each chunk. Defaults to yielding one simulation run a time.
limit (int | None): The maximum number of simulation runs to return, pass None to return all.
status (str | None): Filter by simulation run status
run_type (str | None): Filter by simulation run type
model_external_ids (Sequence[str] | None): Filter by simulator model external ids
simulator_integration_external_ids (Sequence[str] | None): Filter by simulator integration external ids
simulator_external_ids (Sequence[str] | None): Filter by simulator external ids
routine_external_ids (Sequence[str] | None): Filter by routine external ids
routine_revision_external_ids (Sequence[str] | None): Filter by routine revision external ids
model_revision_external_ids (Sequence[str] | None): Filter by model revision external ids
Returns:
Iterator[SimulationRun] | Iterator[SimulatorRunsList]: yields Simulation Run one by one if chunk is not specified, else SimulatorRunsList objects.
"""

filter_runs = SimulatorRunsFilter(
status=status,
run_type=run_type,
model_external_ids=model_external_ids,
simulator_integration_external_ids=simulator_integration_external_ids,
simulator_external_ids=simulator_external_ids,
routine_external_ids=routine_external_ids,
routine_revision_external_ids=routine_revision_external_ids,
model_revision_external_ids=model_revision_external_ids,
)

return self._list_generator(
list_cls=SimulatorRunsList,
resource_cls=SimulationRun,
method="POST",
filter=filter_runs.dump(),
chunk_size=chunk_size,
limit=limit,
)

def list(
self,
limit: int | None = DEFAULT_LIMIT_READ,
status: str | None = None,
run_type: str | None = None,
model_external_ids: Sequence[str] | None = None,
simulator_integration_external_ids: Sequence[str] | None = None,
simulator_external_ids: Sequence[str] | None = None,
routine_external_ids: Sequence[str] | None = None,
routine_revision_external_ids: Sequence[str] | None = None,
model_revision_external_ids: Sequence[str] | None = None,
) -> SimulatorRunsList:
"""`Filter simulation runs <https://developer.cognite.com/api#tag/Simulation-Runs/operation/filter_simulation_runs_simulators_runs_list_post>`_
Retrieves a list of simulation runs that match the given criteria
Args:
limit (int | None): The maximum number of simulation runs to return, pass None to return all.
status (str | None): Filter by simulation run status
run_type (str | None): Filter by simulation run type
model_external_ids (Sequence[str] | None): Filter by simulator model external ids
simulator_integration_external_ids (Sequence[str] | None): Filter by simulator integration external ids
simulator_external_ids (Sequence[str] | None): Filter by simulator external ids
routine_external_ids (Sequence[str] | None): Filter by routine external ids
routine_revision_external_ids (Sequence[str] | None): Filter by routine revision external ids
model_revision_external_ids (Sequence[str] | None): Filter by model revision external ids
Returns:
SimulatorRunsList: List of simulation runs
Examples:
List simulation runs:
>>> from cognite.client import CogniteClient
>>> client = CogniteClient()
>>> res = client.simulators.runs.list()
Filter runs by status and simulator external ids:
>>> res = client.simulators.runs.list(
... simulator_external_ids=["PROSPER", "DWSIM"],
... status="success"
... )
"""

filter_runs = SimulatorRunsFilter(
status=status,
run_type=run_type,
model_external_ids=model_external_ids,
simulator_integration_external_ids=simulator_integration_external_ids,
simulator_external_ids=simulator_external_ids,
routine_external_ids=routine_external_ids,
routine_revision_external_ids=routine_revision_external_ids,
model_revision_external_ids=model_revision_external_ids,
)
self._warning.warn()
return self._list(
method="POST",
limit=limit,
resource_cls=SimulationRun,
list_cls=SimulatorRunsList,
filter=filter_runs.dump(),
)

@overload
def create(self, run: SimulationRunWrite) -> SimulationRun: ...

@overload
def create(self, run: Sequence[SimulationRunWrite]) -> SimulatorRunsList: ...

def create(self, run: SimulationRunWrite | Sequence[SimulationRunWrite]) -> SimulationRun | SimulatorRunsList:
"""`Create simulation runs <https://developer.cognite.com/api#tag/Simulation-Runs/operation/filter_simulation_runs_simulators_runs_list_post>`_
Args:
run (SimulationRunWrite | Sequence[SimulationRunWrite]): The simulation run(s) to execute.
Returns:
SimulationRun | SimulatorRunsList: Created simulation run(s)
Examples:
Create new simulation run:
>>> from cognite.client import CogniteClient
>>> from cognite.client.data_classes.simulators.runs import SimulationRunWrite
>>> client = CogniteClient()
>>> run = [
... SimulationRunWrite(
... routine_external_id="routine1",
... log_severity="Debug",
... run_type="external",
... ),
... ]
>>> res = client.simulators.runs.create(run)
"""
assert_type(run, "simulation_run", [SimulationRunWrite, Sequence])

return self._create_multiple(
list_cls=SimulatorRunsList,
resource_cls=SimulationRun,
items=run,
input_resource_cls=SimulationRunWrite,
resource_path=self._RESOURCE_PATH_RUN,
)
10 changes: 10 additions & 0 deletions cognite/client/data_classes/simulators/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,12 @@
SimulatorModelRevisionWrite,
SimulatorModelWrite,
)
from cognite.client.data_classes.simulators.runs import (
SimulationRun,
SimulationRunWrite,
SimulationRunWriteList,
SimulatorRunsList,
)
from cognite.client.data_classes.simulators.simulators import (
Simulator,
SimulatorIntegration,
Expand All @@ -29,6 +35,9 @@
__all__ = [
"CreatedTimeSort",
"PropertySort",
"SimulationRun",
"SimulationRunWrite",
"SimulationRunWriteList",
"Simulator",
"SimulatorIntegration",
"SimulatorIntegrationFilter",
Expand All @@ -42,6 +51,7 @@
"SimulatorModelRevisionsFilter",
"SimulatorModelWrite",
"SimulatorModelsFilter",
"SimulatorRunsList",
"SimulatorStep",
"SimulatorStepField",
"SimulatorStepOption",
Expand Down
38 changes: 38 additions & 0 deletions cognite/client/data_classes/simulators/filters.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,3 +34,41 @@ def __init__(
all_versions: bool | None = None,
) -> None:
self.model_external_ids = [model_external_ids] if isinstance(model_external_ids, str) else model_external_ids


class SimulatorRunsFilter(CogniteFilter):
def __init__(
self,
status: str | None = None,
run_type: str | None = None,
model_external_ids: str | Sequence[str] | None = None,
simulator_integration_external_ids: str | Sequence[str] | None = None,
simulator_external_ids: str | Sequence[str] | None = None,
routine_external_ids: str | Sequence[str] | None = None,
routine_revision_external_ids: str | Sequence[str] | None = None,
model_revision_external_ids: str | Sequence[str] | None = None,
) -> None:
self.model_external_ids = [model_external_ids] if isinstance(model_external_ids, str) else model_external_ids
self.simulator_integration_external_ids = (
[simulator_integration_external_ids]
if isinstance(simulator_integration_external_ids, str)
else simulator_integration_external_ids
)
self.simulator_external_ids = (
[simulator_external_ids] if isinstance(simulator_external_ids, str) else simulator_external_ids
)
self.routine_external_ids = (
[routine_external_ids] if isinstance(routine_external_ids, str) else routine_external_ids
)
self.routine_revision_external_ids = (
[routine_revision_external_ids]
if isinstance(routine_revision_external_ids, str)
else routine_revision_external_ids
)
self.model_revision_external_ids = (
[model_revision_external_ids]
if isinstance(model_revision_external_ids, str)
else model_revision_external_ids
)
self.status = status
self.run_type = run_type
Loading

0 comments on commit 9e087ca

Please sign in to comment.