feat(simulators): Add support for running simulations using revisions (#2339)

lpereiracgn · web-flow · commit dc7651306638 · 2025-10-08T12:02:04.000Z
diff --git a/cognite/client/_api/simulators/routines.py b/cognite/client/_api/simulators/routines.py
@@ -231,22 +231,58 @@ def list(
             filter=routines_filter.dump(),
         )
 
+    @overload
     def run(
         self,
+        *,
         routine_external_id: str,
         inputs: Sequence[SimulationInputOverride] | None = None,
         run_time: int | None = None,
         queue: bool | None = None,
         log_severity: Literal["Debug", "Information", "Warning", "Error"] | None = None,
         wait: bool = True,
         timeout: float = 60,
+    ) -> SimulationRun: ...
+
+    @overload
+    def run(
+        self,
+        *,
+        routine_revision_external_id: str,
+        model_revision_external_id: str,
+        inputs: Sequence[SimulationInputOverride] | None = None,
+        run_time: int | None = None,
+        queue: bool | None = None,
+        log_severity: Literal["Debug", "Information", "Warning", "Error"] | None = None,
+        wait: bool = True,
+        timeout: float = 60,
+    ) -> SimulationRun: ...
+
+    def run(
+        self,
+        routine_external_id: str | None = None,
+        routine_revision_external_id: str | None = None,
+        model_revision_external_id: str | None = None,
+        inputs: Sequence[SimulationInputOverride] | None = None,
+        run_time: int | None = None,
+        queue: bool | None = None,
+        log_severity: Literal["Debug", "Information", "Warning", "Error"] | None = None,
+        wait: bool = True,
+        timeout: float = 60,
     ) -> SimulationRun:
         """`Run a simulation <https://developer.cognite.com/api#tag/Simulation-Runs/operation/run_simulation_simulators_run_post>`_
 
-        Run a simulation for a given simulator routine.
+        Run a simulation for a given simulator routine. Supports two modes:
+        1. By routine external ID only
+        2. By routine revision external ID + model revision external ID
 
         Args:
-            routine_external_id (str): External id of the simulator routine to run
+            routine_external_id (str | None): External id of the simulator routine to run.
+                Cannot be specified together with routine_revision_external_id and model_revision_external_id.
+            routine_revision_external_id (str | None): External id of the simulator routine revision to run.
+                Must be specified together with model_revision_external_id.
+            model_revision_external_id (str | None): External id of the simulator model revision.
+                Must be specified together with routine_revision_external_id.
             inputs (Sequence[SimulationInputOverride] | None): List of input overrides
             run_time (int | None): Run time in milliseconds. Reference timestamp used for data pre-processing and data sampling.
             queue (bool | None): Queue the simulation run when connector is down.
@@ -258,17 +294,25 @@ def run(
             SimulationRun: Created simulation run
 
         Examples:
-            Create new simulation run:
+            Create new simulation run using routine external ID:
                 >>> from cognite.client import CogniteClient
                 >>> client = CogniteClient()
                 >>> run = client.simulators.routines.run(
                 ...     routine_external_id="routine1",
                 ...     log_severity="Debug"
                 ... )
+
+            Create new simulation run using routine and model revision external IDs:
+                >>> run = client.simulators.routines.run(
+                ...     routine_revision_external_id="routine_revision1",
+                ...     model_revision_external_id="model_revision1",
+                ... )
         """
         self._warning.warn()
         run_object = SimulationRunWrite(
             routine_external_id=routine_external_id,
+            routine_revision_external_id=routine_revision_external_id,
+            model_revision_external_id=model_revision_external_id,
             inputs=list(inputs) if inputs is not None else None,
             run_time=run_time,
             queue=queue,
diff --git a/cognite/client/data_classes/simulators/runs.py b/cognite/client/data_classes/simulators/runs.py
@@ -87,20 +87,33 @@ class SimulationRunCore(WriteableCogniteResource["SimulationRunWrite"]):
     def __init__(
         self,
         run_type: str | None,
-        routine_external_id: str,
+        routine_external_id: str | None,
         run_time: int | None = None,
+        routine_revision_external_id: str | None = None,
+        model_revision_external_id: str | None = None,
     ) -> None:
         self.run_type = run_type
         self.run_time = run_time
         self.routine_external_id = routine_external_id
+        self.routine_revision_external_id = routine_revision_external_id
+        self.model_revision_external_id = model_revision_external_id
 
 
 class SimulationRunWrite(SimulationRunCore):
     """
     Request to run a simulator routine asynchronously.
 
+    This class supports two modes of running simulations:
+    1. By routine external ID only
+    2. By routine revision external ID + model revision external ID
+
     Args:
-        routine_external_id (str): External id of the associated simulator routine
+        routine_external_id (str | None): External id of the associated simulator routine.
+            Cannot be specified together with routine_revision_external_id and model_revision_external_id.
+        routine_revision_external_id (str | None): External id of the associated simulator routine revision.
+            Must be specified together with model_revision_external_id.
+        model_revision_external_id (str | None): External id of the associated simulator model revision.
+            Must be specified together with routine_revision_external_id.
         run_type (str | None): The type of the simulation run
         run_time (int | None): Run time in milliseconds. Reference timestamp used for data pre-processing and data sampling.
         queue (bool | None): Queue the simulation run when connector is down.
@@ -110,15 +123,34 @@ class SimulationRunWrite(SimulationRunCore):
 
     def __init__(
         self,
-        routine_external_id: str,
+        routine_external_id: str | None = None,
+        routine_revision_external_id: str | None = None,
+        model_revision_external_id: str | None = None,
         run_type: str | None = None,
         run_time: int | None = None,
         queue: bool | None = None,
         log_severity: str | None = None,
         inputs: list[SimulationInputOverride] | None = None,
     ) -> None:
+        is_routine_mode = routine_external_id and not routine_revision_external_id and not model_revision_external_id
+        is_revision_mode = (
+            not routine_external_id
+            and routine_revision_external_id is not None
+            and model_revision_external_id is not None
+        )
+
+        # Validate that either routine_external_id OR (routine_revision_external_id + model_revision_external_id) is provided
+        if not (is_routine_mode or is_revision_mode):
+            param_error = ValueError(
+                "Must specify either 'routine_external_id' alone, or both "
+                "'routine_revision_external_id' and 'model_revision_external_id' together."
+            )
+            raise param_error
+
         super().__init__(
             routine_external_id=routine_external_id,
+            routine_revision_external_id=routine_revision_external_id,
+            model_revision_external_id=model_revision_external_id,
             run_type=run_type,
             run_time=run_time,
         )
@@ -128,20 +160,36 @@ def __init__(
 
     @classmethod
     def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> SimulationRunWrite:
-        inputs = resource.get("inputs", None)
+        inputs = resource.get("inputs")
+        routine_revision_external_id = resource.get("routineRevisionExternalId")
+        model_revision_external_id = resource.get("modelRevisionExternalId")
+        routine_external_id = resource.get("routineExternalId")
+
         return cls(
             run_type=resource.get("runType"),
-            routine_external_id=resource["routineExternalId"],
             run_time=resource.get("runTime"),
             queue=resource.get("queue"),
             log_severity=resource.get("logSeverity"),
             inputs=([SimulationInputOverride._load(_input) for _input in inputs] if inputs else None),
+            routine_external_id=routine_external_id,
+            routine_revision_external_id=routine_revision_external_id,
+            model_revision_external_id=model_revision_external_id,
         )
 
     def dump(self, camel_case: bool = True) -> dict[str, Any]:
         output = super().dump(camel_case=camel_case)
         if self.inputs is not None:
             output["inputs"] = [input_.dump(camel_case=camel_case) for input_ in self.inputs]
+
+        # Remove fields based on the mode we're in
+        if self.routine_external_id is not None:
+            # Routine-only mode: remove revision fields that might be None
+            output.pop("routineRevisionExternalId", None)
+            output.pop("modelRevisionExternalId", None)
+        else:
+            # Revision mode: remove routine_external_id
+            output.pop("routineExternalId", None)
+
         return output
 
     def as_write(self) -> SimulationRunWrite:
diff --git a/tests/tests_integration/test_api/test_simulators/test_routines.py b/tests/tests_integration/test_api/test_simulators/test_routines.py
@@ -4,9 +4,12 @@
 
 from cognite.client._cognite_client import CogniteClient
 from cognite.client.data_classes.simulators.filters import PropertySort
+from cognite.client.data_classes.simulators.routine_revisions import SimulatorRoutineRevision
 from cognite.client.data_classes.simulators.routines import SimulatorRoutine, SimulatorRoutineWrite
 from cognite.client.utils._text import random_string
-from tests.tests_integration.test_api.test_simulators.seed.data import ResourceNames
+from tests.tests_integration.test_api.test_simulators.seed.data import (
+    ResourceNames,
+)
 
 
 @pytest.mark.usefixtures(
@@ -55,6 +58,34 @@ def test_sort(
         for i in range(1, len(routines_asc)):
             assert routines_asc[i].created_time >= routines_asc[i - 1].created_time
 
+
+@pytest.mark.usefixtures(
+    "seed_resource_names",
+    "seed_simulator_routine_revisions",
+)
+class TestSimulatorRoutinesRunWithRevisions:
+    def test_run_with_revisions(
+        self,
+        cognite_client: CogniteClient,
+        seed_resource_names: ResourceNames,
+        seed_simulator_routine_revisions: tuple[SimulatorRoutineRevision, SimulatorRoutineRevision],
+    ) -> None:
+        """Test running a simulation using routine and model revision external IDs."""
+        routine_revision_external_id = seed_simulator_routine_revisions[0].external_id
+        model_revision_external_id = seed_resource_names.simulator_model_revision_external_id
+        simulator_integration_unique_external_id = seed_resource_names.simulator_integration_external_id
+
+        # Run simulation using revision external IDs
+        run = cognite_client.simulators.routines.run(
+            routine_revision_external_id=routine_revision_external_id,
+            model_revision_external_id=model_revision_external_id,
+            wait=False,  # Don't wait to avoid timeout in tests
+        )
+
+        assert run is not None
+        assert run.id is not None
+        assert run.routine_revision_external_id == routine_revision_external_id
+        assert run.model_revision_external_id == model_revision_external_id
         # Test iterator with sort
         routines_iter = list(
             cognite_client.simulators.routines(
diff --git a/tests/tests_integration/test_api/test_simulators/test_runs.py b/tests/tests_integration/test_api/test_simulators/test_runs.py
@@ -7,15 +7,14 @@
 from cognite.client._cognite_client import CogniteClient
 from cognite.client.data_classes import TimestampRange
 from cognite.client.data_classes.simulators.filters import SimulationRunsSort
-from cognite.client.data_classes.simulators.routine_revisions import SimulatorRoutineRevision
 from cognite.client.data_classes.simulators.runs import (
     SimulationInput,
     SimulationOutput,
     SimulationRun,
     SimulationRunWrite,
     SimulationValueUnitName,
 )
-from tests.tests_integration.test_api.test_simulators.seed.data import ResourceNames
+from tests.tests_integration.test_api.test_simulators.seed.data import RESOURCES, ResourceNames
 
 
 @pytest.mark.usefixtures("seed_resource_names", "seed_simulator_routine_revisions")
@@ -113,23 +112,42 @@ async def test_run_with_wait_and_retrieve(
         assert data_res is not None
         assert data_res.dump() == data_res2.dump()
 
+    @pytest.mark.parametrize(
+        "simulation_run_item",
+        [
+            (
+                SimulationRunWrite(
+                    routine_revision_external_id=f"{RESOURCES.simulator_routine_external_id}_v1",
+                    model_revision_external_id=RESOURCES.simulator_model_revision_external_id,
+                )
+            ),
+            (
+                SimulationRunWrite(
+                    routine_external_id=RESOURCES.simulator_routine_external_id,
+                )
+            ),
+        ],
+        ids=[
+            "with_routine_revision_and_model_revision",
+            "with_routine_only",
+        ],
+    )
     def test_create_run(
         self,
         cognite_client: CogniteClient,
-        seed_simulator_routine_revisions: tuple[SimulatorRoutineRevision, SimulatorRoutineRevision],
-        seed_resource_names: ResourceNames,
+        simulation_run_item: SimulationRunWrite,
     ) -> None:
-        routine_external_id = seed_resource_names.simulator_routine_external_id
-        created_runs = cognite_client.simulators.runs.create(
-            [
-                SimulationRunWrite(
-                    run_type="external",
-                    routine_external_id=routine_external_id,
-                )
-            ]
+        created_runs = cognite_client.simulators.runs.create([simulation_run_item])
+        is_run_by_revision = (
+            simulation_run_item.routine_revision_external_id is not None
+            and simulation_run_item.model_revision_external_id is not None
         )
+        if is_run_by_revision:
+            assert created_runs[0].routine_revision_external_id == simulation_run_item.routine_revision_external_id
+            assert created_runs[0].model_revision_external_id == simulation_run_item.model_revision_external_id
+        else:
+            assert created_runs[0].routine_external_id == simulation_run_item.routine_external_id
         assert len(created_runs) == 1
-        assert created_runs[0].routine_external_id == routine_external_id
         assert created_runs[0].id is not None
 
     def test_list_filtering_timestamp_ranges(
diff --git a/tests/tests_unit/test_data_classes/test_simulators.py b/tests/tests_unit/test_data_classes/test_simulators.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import re
+
 import pytest
 
 from cognite.client.data_classes.simulators.models import (
@@ -22,6 +24,7 @@
     SimulationOutput,
     SimulationRunDataItem,
     SimulationRunDataList,
+    SimulationRunWrite,
     SimulationValueUnitName,
 )
 from cognite.client.data_classes.simulators.simulators import (
@@ -464,3 +467,25 @@ def test_get_units_empty_units_list(self) -> None:
         units = simulator.get_units("empty")
         assert units == []
         assert len(units) == 0
+
+
+class TestSimulationRunWrite:
+    def test_error_handling(self) -> None:
+        error_message = "Must specify either 'routine_external_id' alone, or both 'routine_revision_external_id' and 'model_revision_external_id' together."
+        with pytest.raises(
+            ValueError,
+            match=re.escape(error_message),
+        ):
+            SimulationRunWrite(
+                routine_external_id="routine_external_id_1",
+                routine_revision_external_id="routine_revision_external_id_1",
+                model_revision_external_id="model_revision_external_id_1",
+            )
+
+        with pytest.raises(
+            ValueError,
+            match=re.escape(error_message),
+        ):
+            SimulationRunWrite(
+                routine_revision_external_id="routine_revision_external_id_1",
+            )
diff --git a/tests/utils.py b/tests/utils.py
@@ -53,6 +53,7 @@
 from cognite.client.data_classes.filters import Filter
 from cognite.client.data_classes.hosted_extractors.jobs import BodyLoad, NextUrlLoad, RestConfig
 from cognite.client.data_classes.simulators.routine_revisions import SimulatorRoutineStepArguments
+from cognite.client.data_classes.simulators.runs import SimulationRunWrite
 from cognite.client.data_classes.transformations.notifications import TransformationNotificationWrite
 from cognite.client.data_classes.transformations.schedules import TransformationScheduleWrite
 from cognite.client.data_classes.transformations.schema import TransformationSchemaUnknownType
@@ -485,6 +486,16 @@ def create_instance(self, resource_cls: type[T_Object], skip_defaulted_args: boo
                 keyword_arguments.pop("max_list_size", None)
         elif resource_cls is SimulatorRoutineStepArguments:
             keyword_arguments = {"data": {"reference_id": self._random_string(50), "arg2": self._random_string(50)}}
+        elif resource_cls is SimulationRunWrite:
+            # SimulationRunWrite requires either routine_external_id alone OR both routine_revision_external_id and model_revision_external_id
+            if self._random.choice([True, False]):
+                # Use routine_external_id only
+                keyword_arguments.pop("routine_revision_external_id", None)
+                keyword_arguments.pop("model_revision_external_id", None)
+                keyword_arguments["routine_external_id"] = self._random_string(50)
+            else:
+                # Use revision-based parameters
+                keyword_arguments.pop("routine_external_id", None)
 
         return resource_cls(*positional_arguments, **keyword_arguments)
 
