simstack 0.2.0__py3-none-any.whl

simstack/methods/multiple_executor.py

@@ -0,0 +1,307 @@
+import asyncio
+import logging
+from typing import List, Callable, TypeVar, Generic, Dict, Any, AsyncGenerator, Union
+
+from odmantic import Model, Field, ObjectId, Reference
+
+from simstack.core.asnyc_helper import async_helper
+from simstack.core.context import context
+from simstack.core.node import node
+from simstack.models import (
+    simstack_model,
+    ModelMapping,
+    BinaryOperationInput,
+    FloatData,
+)
+from simstack.util.async_zip_utils import async_zip
+from simstack.util.importer import import_function, import_class
+from simstack.util.make_table import make_table_entries
+
+# Define your prefix
+PREFIX = "[ParallelExecutor]"
+
+
+@node
+def multiple_executor_adder(
+    arg1: BinaryOperationInput, arg2: BinaryOperationInput
+) -> BinaryOperationInput:
+    return BinaryOperationInput(arg1=arg1.arg1 + arg2.arg1, arg2=arg1.arg2 + arg2.arg2)
+
+
+class PrefixFilter(logging.Filter):
+    def filter(self, record):
+        record.msg = f"{PREFIX} {record.msg}"
+        return True
+
+
+# Set up the logger
+logger = logging.getLogger(__name__)
+logger.addFilter(PrefixFilter())
+
+
+@async_helper
+async def get_mappings(func: Callable[[Model], Model]):
+    """
+    Decorator to map a function to a model.
+    """
+    node_model = await context.db.load_node_model_by_name(func.__name__)
+    if node_model is None:
+        raise ValueError(f"Node model {func.__name__} not found in database.")
+
+    # Get return type annotation from function
+    return_type = func.__annotations__.get("return")
+    if return_type is None:
+        raise ValueError(f"Function {func.__name__} must have return type annotation")
+
+    result_model = await context.db.find_one(
+        ModelMapping, ModelMapping.name == return_type.__name__
+    )
+    if result_model is None:
+        raise ValueError(f"Result model {return_type.__name__} not found in database.")
+
+    return node_model.function_mapping, node_model.input_mappings, result_model.mapping
+
+
+T = TypeVar("T", bound=Model)
+R = TypeVar("R", bound=Model)
+
+
+def _prefix_dict_keys(original_dict, prefix):
+    """Helper method to prefix dictionary keys."""
+    return {f"{prefix}{key}": value for key, value in original_dict.items()}
+
+
+async def _get_formatted_column_defs(mapping, prefix):
+    """Get column definitions for a model with proper field prefixes."""
+    model_class = await import_class(mapping)
+    column_defs = model_class.make_column_defs()
+
+    # Add prefix to field names
+    for col in column_defs:
+        if "field" in col:
+            col["field"] = f"{prefix}{col['field']}"
+
+    # Filter out id fields
+    return [col for col in column_defs if col.get("field") != f"{prefix}id"]
+
+
+@simstack_model
+class MultipleExecutorInput(Model, Generic[T]):
+    """
+    Represents a MultipleExecutor which handles multiple calculations in parallel and aggregates the results.
+    # TODO store the inputs as ids not in memory
+    """
+
+    function_mapping: str
+    input_mappings: List[str]
+    result_mapping: str
+    timeout: int = 600
+    inputs_data: List[List[ObjectId]] = Field(default_factory=list)
+
+    @classmethod
+    def from_function(cls, func: Callable[[T], R]):
+        function_mapping, input_mappings, result_mapping = get_mappings(func)
+        print(
+            f"Function mapping: {function_mapping}, Input mappings: {input_mappings}, Result mapping: {result_mapping}"
+        )
+        return cls(
+            function_mapping=function_mapping,
+            input_mappings=input_mappings,
+            result_mapping=result_mapping,
+        )
+
+    async def input_generator(self) -> AsyncGenerator[List[Model], None]:
+        """Deserialize the input models"""
+        input_model_classes = [
+            await import_class(input_mapping) for input_mapping in self.input_mappings
+        ]
+        for input_item in self.inputs_data:
+            input_models = [
+                await context.db.find_one(
+                    input_model_class, input_model_class.id == model_id
+                )
+                for input_model_class, model_id in zip(input_model_classes, input_item)
+            ]
+            # if one models is None: fail
+            if any(model is None for model in input_models):
+                raise ValueError(
+                    f"One or more input models not found in database for input IDs: {input_item}"
+                )
+            yield input_models
+
+    async def append(self, input_models: Union[T, List[T]], with_save=True):
+        """Add the inputs for a single function call"""
+        if isinstance(input_models, Model):
+            input_models = [input_models]
+        if with_save:
+            for input_model in input_models:
+                await context.db.save(input_model)
+        input_ids = [input_model.id for input_model in input_models]
+        self.inputs_data.append(input_ids)
+
+    async def execute(self, **kwargs) -> "MultipleExecutorOutput":
+        task_id = kwargs.get("task_id", "NA")
+        global PREFIX
+        if task_id != "NA":
+            PREFIX = f"[ParallelExecutor task_id: {task_id}]"
+
+        # Prepare all jobs to run concurrently
+        tasks = []
+        result_ids = [None] * len(self.inputs_data)
+        kwargs["parent_id"] = task_id
+        try:
+            func = await import_function(self.function_mapping, task_id)
+            i = 0
+            async for task_input in self.input_generator():
+                # Create an async task that calls our async wrapper
+                task = asyncio.create_task(func(*task_input, **kwargs))
+                # Store index for later reference
+                task.index = i
+                tasks.append(task)
+                i += 1
+        except Exception as e:
+            logger.error(f"Error preparing tasks: {str(e)}")
+
+        # Wait for all tasks with timeout
+        done, pending = await asyncio.wait(tasks, timeout=self.timeout)
+
+        # Process completed tasks
+        for task in done:
+            try:
+                task_result = task.result()
+                result_ids[task.index] = task_result.id
+                logger.info(f"Successfully processed task {task.index}")
+            except Exception as e:
+                logger.error(f"Error processing task {task.index}: {str(e)}")
+
+        # Log information about pending tasks
+        if pending:
+            logger.error(f"{len(pending)} tasks didn't complete within the timeout")
+            for task in pending:
+                logger.error(
+                    f"Task at index {task.index} didn't complete within timeout"
+                )
+
+        output = MultipleExecutorOutput(
+            executor_input=self,
+            result_mapping=self.result_mapping,
+            result_ids=result_ids,
+        )
+        logger.info("finished")
+        return output
+
+    @classmethod
+    def ui_base_schema(cls):
+        return {"inputs_data": {"ui:widget": "hidden"}}
+
+
+@simstack_model
+class MultipleExecutorOutput(Model, Generic[T, R]):
+    """
+    Represents the output of a MultipleExecutor.
+    """
+
+    executor_input: MultipleExecutorInput = Reference()
+    result_mapping: str  # there is one result: either a Model or a SimstackResult
+    result_ids: List[Union[ObjectId, None]] = Field(default_factory=list)
+
+    @classmethod
+    def json_schema(cls):
+        return {
+            "type": "object",
+            "properties": {"results": {"type": "array", "items": {"type": "object"}}},
+            "required": ["result_mapping", "results"],
+        }
+
+    @async_helper
+    async def custom_model_dump(self) -> Dict[str, Any]:
+        """
+        Retrieve all results and return them as a list of dictionaries.
+        """
+        results = []
+        async for result in self.result_generator():
+            if result is not None:
+                results.append(result.model_dump())
+            else:
+                results.append(None)
+        print(f"{PREFIX} custom_model_dump: {results}")
+        return {"result": results}
+
+    async def result_generator(self) -> AsyncGenerator[Union[R, None], None]:
+        """Deserialize the result models"""
+        result_model_class = await import_class(self.result_mapping)
+        for result_id in self.result_ids:
+            if result_id is not None:
+                result_model = await context.db.find_one(
+                    result_model_class, result_model_class.id == result_id
+                )
+                if result_model is not None:
+                    yield result_model
+                else:
+                    yield None
+            else:
+                yield None
+
+    @async_helper
+    async def make_table(self):
+        """
+        Create a table representation of the results with formatted columns.
+        Returns a dictionary with table data and column definitions.
+        """
+        # Get inputs and results data
+        inputs = await self.get_inputs()
+        results = await self.get_results()
+
+        # Create tables for inputs and results
+        inputs_table = make_table_entries(inputs)
+        results_table = make_table_entries(results)
+
+        # Merge data based on finished indices
+        merged_entries = []
+        for i, input_entry in enumerate(inputs_table["rowData"]):
+            entry = {**input_entry}
+            if i in self.finished_indices:
+                result_idx = self.finished_indices.index(i)
+                entry.update(results_table["rowData"][result_idx])
+            # TODO find all artifacts to this result and add the data to the entry
+            merged_entries.append(entry)
+
+        # Combine column definitions
+        merged_columns = inputs_table["columnDefs"] + results_table["columnDefs"]
+
+        return {"rowData": merged_entries, "columnDefs": merged_columns}
+
+
+@node
+async def multiple_executor(
+    executor: MultipleExecutorInput, **kwargs
+) -> MultipleExecutorOutput:
+    return await executor.execute(**kwargs)
+
+
+@node
+async def async_adder(arg: BinaryOperationInput, **kwargs) -> FloatData:
+    return FloatData(value=arg.arg1 + arg.arg2)
+
+
+async def main():
+    context.initialize()
+    executor_input = MultipleExecutorInput.from_function(async_adder)
+    # Create a list of input models
+    inputs = [BinaryOperationInput(arg1=2, arg2=value) for value in range(1, 3)]
+
+    # Append the input models to the executor
+    for input_model in inputs:
+        await executor_input.append(input_model)
+
+    executor_output = await multiple_executor(executor_input)
+
+    print("Results: ", executor_output)
+    async for item1, item2 in async_zip(
+        executor_input.input_generator(), executor_output.result_generator()
+    ):
+        print(f"Input: {item1}, Result: {item2}")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

simstack/models/__init__.py

@@ -0,0 +1,37 @@
+from .simstack_model import simstack_model
+from .array_list import ArrayList
+from .artifact_models import ArtifactMapping, ArtifactModel
+from .base_types import (
+    IntData,
+    FloatData,
+    StringData,
+    BooleanData,
+    BinaryOperationInput,
+    IteratorInput,
+)
+from .file_list import FileListModel
+from .models import (
+    ModelMapping,
+    NodeModel,
+)
+from .parameters import Parameters
+from .node_registry import NodeRegistry
+
+
+__all__ = [
+    "ModelMapping",
+    "NodeModel",
+    "IntData",
+    "FloatData",
+    "StringData",
+    "BooleanData",
+    "ArrayList",
+    "BinaryOperationInput",
+    "IteratorInput",
+    "ArtifactMapping",
+    "ArtifactModel",
+    "FileListModel",
+    "Parameters",
+    "NodeRegistry",
+    "simstack_model",
+]

simstack/models/array_list.py

@@ -0,0 +1,48 @@
+from typing import List
+
+from odmantic import Model, Field
+
+from simstack.models import simstack_model
+from simstack.models.array_storage import ArrayStorage
+
+
+@simstack_model
+class ArrayList(Model):
+    array_list: List[ArrayStorage] = Field(default_factory=list)
+
+    def append(self, array_storage: ArrayStorage):
+        """Add a new ArrayStorage to the list"""
+        self.array_list.append(array_storage)
+
+    def get(self, index: int) -> ArrayStorage:
+        """Get ArrayStorage at specified index"""
+        return self.array_list[index]
+
+    def remove(self, index: int):
+        """Remove ArrayStorage at specified index"""
+        del self.array_list[index]
+
+    @property
+    def length(self) -> int:
+        """Get number of items in the list"""
+        return len(self.array_list)
+
+    def clear(self):
+        """Remove all items from the list"""
+        self.array_list.clear()
+
+    def __iter__(self):
+        """Enable iteration over array_list"""
+        return iter(self.array_list)
+
+    def __getitem__(self, index: int) -> ArrayStorage:
+        """Enable access using square brackets"""
+        return self.array_list[index]
+
+    def __setitem__(self, index: int, value: ArrayStorage):
+        """Enable item assignment using square brackets"""
+        self.array_list[index] = value
+
+    def __len__(self) -> int:
+        """Enable len() function"""
+        return len(self.array_list)

simstack/models/array_storage.py

@@ -0,0 +1,54 @@
+import json
+from typing import Optional
+
+from odmantic import Model
+
+from simstack.models.simstack_model import simstack_model
+from simstack.util.ui_tools import ui_hide_fields
+
+
+@simstack_model
+class ArrayStorage(Model):
+    name: str
+    shape: Optional[str] = None  # Store array shape as string like "3,3"
+    data_json: Optional[str] = None  # Store flattened array data as JSON
+
+    def set_array(self, array):
+        """Store a numpy array"""
+        self.shape = ",".join(str(dim) for dim in array.shape)
+        self.data_json = json.dumps(array.flatten().tolist())
+
+    def get_array(self):
+        """Retrieve the numpy array"""
+        import numpy as np
+
+        shape = tuple(int(dim) for dim in self.shape.split(",")) if self.shape else ()
+        flat_array = (
+            np.array(json.loads(self.data_json)) if self.data_json else np.array([])
+        )
+        return flat_array.reshape(shape)
+
+    def make_table_entries(
+        self,
+        max_recursion_level=1,
+        drop_id=True,
+        current_level=0,
+        visited=None,
+        field_prefix="",
+    ):
+        return {"name": self.name}
+
+    def make_column_defs_instance(
+        self,
+        table_name=None,
+        max_recursion_level=1,
+        drop_id=True,
+        current_level=0,
+        visited=None,
+        field_prefix="",
+    ):
+        return [{"field": self.name, "headerName": "Array"}]
+
+    @classmethod
+    def ui_schema(cls, **kwargs) -> dict:
+        return ui_hide_fields({}, ["shape", "data_json"])

simstack/models/artifact_models.py

@@ -0,0 +1,127 @@
+from typing import Optional, Dict, Any
+
+from odmantic import Model, Field
+
+from simstack.models.pickle_models import FunctionPickle
+from simstack.models.simstack_model import simstack_model
+
+
+# class ArtifactList(List):
+#     """A list that knows how to update its parent Artifact when modified."""
+#
+#     def __init__(self, parent_artifact, key, initial_values=None):
+#         super().__init__(initial_values or [])
+#         self.parent_artifact = parent_artifact
+#         self.key = key
+#
+#     def append(self, value):
+#         super().append(value)
+#         self.parent_artifact[self.key] = self  # Update the parent
+#
+#     def extend(self, values):
+#         super().extend(values)
+#         self.parent_artifact[self.key] = self  # Update the parent
+#
+#     # Add other list methods that modify the list
+#     def __setitem__(self, index, value):
+#         super().__setitem__(index, value)
+#         self.parent_artifact[self.key] = self
+#
+#     def __delitem__(self, index):
+#         super().__delitem__(index)
+#         self.parent_artifact[self.key] = self
+#
+#     def insert(self, index, value):
+#         super().insert(index, value)
+#         self.parent_artifact[self.key] = self
+#
+#     def pop(self, index=-1):
+#         value = super().pop(index)
+#         self.parent_artifact[self.key] = self
+#         return value
+#
+#     def remove(self, value):
+#         super().remove(value)
+#         self.parent_artifact[self.key] = self
+#
+#     def clear(self):
+#         super().clear()
+#         self.parent_artifact[self.key] = self
+#
+#     def sort(self, *args, **kwargs):
+#         super().sort(*args, **kwargs)
+#         self.parent_artifact[self.key] = self
+#
+#     def reverse(self):
+#         super().reverse()
+#         self.parent_artifact[self.key] = self
+
+
+@simstack_model
+class ArtifactModel(Model):
+    name: str
+    description: Optional[str] = None
+    data: Dict[str, Any] = Field(default_factory=dict)
+    path: Optional[str] = None
+
+    model_config = {
+        "collection": "artifacts"  # All subclasses will use the same collection
+    }
+
+
+# class IntArtifactModel(ArtifactModel):
+#     type: Literal["int"] = "int"  # Fixed value for this subclass
+#     value: int
+#
+# class FloatArtifactModel(ArtifactModel):
+#     type: Literal["float"] = "float"  # Fixed value for this subclass
+#     value: float
+#
+# class StringArtifactModel(ArtifactModel):
+#     type: Literal["string"] = "string"  # Fixed value for this subclass
+#     value: str
+#
+# class ImageArtifactModel(ArtifactModel):
+#     type: Literal["image"] = "image"  # Fixed value for this subclass
+#     value: str  # the image is stored as a base64 string or a file path
+#
+# # Forward reference for type hints
+# ArtifactType = Union[ArtifactModel, IntArtifactModel, FloatArtifactModel, StringArtifactModel, 'ArtifactModelList', ImageArtifactModel]
+#
+# # List subclass - can contain other items
+# class ArtifactModelList(ArtifactModel):
+#     type: Literal["list"] = "list"
+#     items: List[ArtifactType] = Field(default_factory=list)  # List of Item objects
+
+
+class ArtifactMapping(Model):
+    """
+    ArtifactsMapper is a mapping between the artifact and a node registry-path.
+    The workflow executor passes a path of the type
+
+    node1.node2.node4. ... .nodeN
+
+    where node is the function name of the node
+
+    Regex can maps this to the target path of the ArtifactsMapping, e.g. a path
+
+    *.parent1.node
+
+    it would map on all nodes with name node that have been directly called by a node with the name parent1.
+
+
+    """
+
+    name: str = Field(default="artifact", unique=True)
+    regex_pattern: str = Field(default="")
+    function_mapping: str = Field(default="")
+    function_code: str = Field(default="")
+    pickle_function: Optional[FunctionPickle] = None
+
+    def set_values(self, other: "ArtifactMapping") -> "ArtifactMapping":
+        self.name = other.name
+        self.regex_pattern = other.regex_pattern
+        self.function_mapping = other.function_mapping
+        self.function_code = other.function_code
+        self.pickle_function = other.pickle_function
+        return self