simstack 2.0.1__py3-none-any.whl

simstack/methods/generate_test.py

@@ -0,0 +1,125 @@
+import argparse
+import asyncio
+import json
+import logging
+import shutil
+from pathlib import Path
+from typing import Any, Dict, List
+
+from odmantic import Model, ObjectId
+from pydantic import BaseModel
+
+from simstack.core.context import context
+from simstack.models import NodeRegistry, NamedDataReference
+from simstack.util.importer import import_class
+
+logger = logging.getLogger("generate_test")
+
+async def load_models(references: List[NamedDataReference]) -> Dict[str, Any]:
+    """Load models from the database based on references."""
+    result = {}
+    db = context.db
+    for ref in references:
+        try:
+            model_cls = await import_class(ref.variable_mapping, db)
+            if not model_cls:
+                raise RuntimeError(f"Could not import class {ref.variable_mapping}")
+
+            obj = await db.find_one(model_cls, model_cls.id == ref.reference)
+            if obj is None:
+                raise RuntimeError(
+                    f"Object {ref.reference} not found in {ref.variable_mapping}"
+                )
+
+            # Preserve the argument/result name captured at execution time.
+            key = ref.variable_name
+            if not key:
+                from simstack.models import ModelMapping
+
+                model_mapping = await db.find_one(
+                    ModelMapping,
+                    ModelMapping.mapping == ref.variable_mapping,
+                )
+                key = model_mapping.name if model_mapping else ref.variable_mapping
+
+            if key in result:
+                key = f"{key}_{str(ref.reference)}"
+            result[key] = obj
+        except Exception as e:
+            logger.exception(f"Error loading {ref.variable_mapping} {ref.reference}: {e}")
+            raise RuntimeError(
+                f"Failed to load {ref.variable_mapping} {ref.reference}"
+            ) from e
+    return result
+
+def serialize_models(models: Dict[str, Any]) -> str:
+    """Serialize models to a JSON string."""
+    data = {}
+    for key, obj in models.items():
+        if isinstance(obj, (Model, BaseModel)):
+            data[key] = obj.model_dump()
+        else:
+            data[key] = str(obj)
+    return json.dumps(data, indent=4, default=str)
+
+async def generate_test(node_id: str, target_base: Path):
+    """Generate a test case for a given node ID."""
+    db = context.db
+    registry_entry = await db.find_one(NodeRegistry, NodeRegistry.id == ObjectId(node_id))
+    if not registry_entry:
+        raise ValueError(f"NodeRegistry entry with ID {node_id} not found.")
+
+    node_name = registry_entry.name
+    arg_hash = registry_entry.arg_hash
+    
+    # Target directory: target/node_name/arg_hash
+    target_dir = target_base / node_name / arg_hash
+    target_dir.mkdir(parents=True, exist_ok=True)
+    
+    # Source directory: workdir/node_name/id
+    workdir = context.config.workdir
+    source_dir = workdir / node_name / str(registry_entry.id)
+    
+    if source_dir.exists() and source_dir.is_dir():
+        print(f"Copying files from {source_dir} to {target_dir}")
+        # Copy all files from source to target
+        for item in source_dir.iterdir():
+            if item.is_file():
+                shutil.copy2(item, target_dir / item.name)
+            elif item.is_dir():
+                shutil.copytree(item, target_dir / item.name, dirs_exist_ok=True)
+    else:
+        print(f"Source directory {source_dir} does not exist. Skipping file copy.")
+
+    # Load and serialize inputs
+    print(f"Serializing inputs for {node_id}")
+    inputs = await load_models(registry_entry.input_references)
+    with open(target_dir / "inputs.json", "w") as f:
+        f.write(serialize_models(inputs))
+
+    # Load and serialize outputs
+    print(f"Serializing outputs for {node_id}")
+    outputs = await load_models(registry_entry.results_references)
+    with open(target_dir / "outputs.json", "w") as f:
+        f.write(serialize_models(outputs))
+
+    print(f"Test generation for node {node_id} completed at {target_dir}")
+
+async def async_main():
+    parser = argparse.ArgumentParser(description="Generate a test case from a node execution.")
+    parser.add_argument("--id", required=True, help="ID of the NodeRegistry entry")
+    parser.add_argument("--target", default=str(Path.cwd() / "tests"), help="Target base directory for tests")
+    parser.add_argument("--resource", default="self", help="Resource to use for database connection")
+    
+    args = parser.parse_args()
+    
+    await context.initialize(resource=args.resource)
+    
+    await generate_test(args.id, Path(args.target))
+
+def generate_test_main():
+    logging.basicConfig(level=logging.INFO)
+    asyncio.run(async_main())
+
+if __name__ == "__main__":
+    generate_test_main()

simstack/methods/get_file.py

@@ -0,0 +1,26 @@
+from pathlib import Path
+
+from simstack.core.node import node
+from simstack.models.files import FileStack, FileGetterArgs
+from simstack.models.parameters import Resource
+
+@node
+def get_file(file_getter_args: FileGetterArgs) -> FileStack:
+    """
+    Retrieves a file from a given source and saves it to the specified target path.
+
+    This function is responsible for handling the retrieval of files from a
+    FileStack source and storing them to a local resource's target location.
+    It will update the FileStack
+
+    Args:
+        file_getter_args (FileGetterArgs): The arguments for the file retrieval operation.
+
+    Returns:
+        FileStack: The modified FileStack with the local resource.
+
+    Raises:
+        NotImplementedError: If the function is not yet implemented.
+    """
+    raise NotImplementedError("get_file is not implemented yet.")
+    return False

simstack/methods/mass_runner.py

@@ -0,0 +1,154 @@
+import asyncio
+from typing import List, Callable, Any, Optional
+
+from odmantic import Model
+
+from simstack.core.context import context
+from simstack.core.definitions import TaskStatus
+from simstack.core.hash import complex_hash_function
+from simstack.core.simstack_result import SimstackResult
+from simstack.core.node_runner import NodeRunner
+from simstack.core.process_results import process_result_helper
+from simstack.models import DataSetMetadata, DataSet, DataSetSection, BooleanData, StringData
+
+class MassRunner(NodeRunner):
+    def __init__(self, node: Callable[..., Any],max_concurrency: Optional[int] = None, **kwargs):
+        super().__init__(kwargs["node_runner"].name, kwargs["node_runner"].task_id, kwargs["node_runner"].logger )
+        node_runner = kwargs["node_runner"]
+        arg_hash = kwargs["arg_hash"]
+        self._kwargs = kwargs.copy()
+
+
+        database_metadata = DataSetMetadata(
+            field_name=node_runner.name,
+            data={
+                "arg_hash": str(arg_hash),
+                "task_id": str(node_runner.task_id),
+                "call_path": getattr(node_runner, "call_path", "NA"),
+            }
+        )
+
+        self._max_concurrency = max_concurrency
+
+        self.dataset = DataSet(
+            field_name=f"{node_runner.name}.{self.task_id}",
+            metadata=database_metadata
+        )
+        self._existing_dataset = None
+        self._failure = False
+        self.dataset["tasks"] = DataSetSection()
+        self._tasks = []
+        self._node = node
+        self._semaphore = asyncio.Semaphore(max_concurrency) if max_concurrency else None
+
+
+    async def recover_orphaned_datasets(self):
+        previous_node_name = self._kwargs.get("previous_node_name", None)
+        previous_task_id = self._kwargs.get("previous_task_id", None)
+        if previous_node_name and previous_task_id:
+            db = context.db
+            previous_dataset_field_name = f"{previous_node_name}.{previous_task_id}"
+            existing_dataset = await db.find_one(DataSet, DataSet.field_name == previous_dataset_field_name)
+            if existing_dataset:
+                self._existing_dataset = existing_dataset
+                self.info(f"Found existing dataset: {previous_dataset_field_name}")
+                return
+
+    async def _run_node(self, args: List[Model]):
+        await self.recover_orphaned_datasets()
+
+        import inspect
+        sig = inspect.signature(self._node)
+        param_names = list(sig.parameters.keys())
+
+        arg_hashes = [complex_hash_function(arg) for arg in args]
+        combined_arg_hash = complex_hash_function(arg_hashes)
+        
+        if self._existing_dataset and combined_arg_hash in self._existing_dataset["tasks"].data:
+            row = self._existing_dataset["tasks"].get_item(combined_arg_hash)
+            success = row.get("success", None)
+            if success is not None:
+                success = success.value
+            if success:
+                self.info(f"Skipping node with arg_hash: {combined_arg_hash}")
+                self.dataset["tasks"].add_row(row, combined_arg_hash)
+                return None
+        
+        self.info(f"Running node with arg_hash: {combined_arg_hash}")
+
+        task_dict = {}
+
+        for i, m in enumerate(args):
+            if i < len(param_names):
+                task_dict[f"arg_{param_names[i]}"] = m
+            else:
+                task_dict[f"arg_{i}"] = m
+
+        try:
+            if self._semaphore:
+                async with self._semaphore:
+                    if asyncio.iscoroutinefunction(self._node):
+                        result = await self._node(*args, **self._kwargs)
+                    else:
+                        result = self._node(*args, **self._kwargs)
+            else:
+                if asyncio.iscoroutinefunction(self._node):
+                    result = await self._node(*args, **self._kwargs)
+                else:
+                    result = self._node(*args, **self._kwargs)
+        except Exception as e:
+            self._failure = True
+            self.error(f"Error running node: {e}")
+            task_dict["success"] = BooleanData(value=False)
+            task_dict["error"] = StringData(field_name="error_message",value=str(e))
+            self.dataset["tasks"].add_row(task_dict, name=combined_arg_hash)
+            return None
+
+        if result is None:
+            task_dict["success"] = BooleanData(value=False)
+        elif isinstance(result, bool):
+            task_dict["success"] = BooleanData(value=result)
+        elif isinstance(result, (SimstackResult, Model)):
+            references, models = await process_result_helper(result, self.task_id)
+            if isinstance(result, SimstackResult):
+                if result.status != TaskStatus.COMPLETED:
+                    self._failure = True
+                task_dict["success"] = BooleanData(value=result.status == TaskStatus.COMPLETED)
+            else: 
+                task_dict["success"] = BooleanData(value=True)
+            if references:
+                for reference, model in zip(references, models):
+                    task_dict[f"result_{reference.variable_name}"] = model
+        elif isinstance(result, (list, tuple)) and all(isinstance(m, Model) for m in result):
+            task_dict["success"] = BooleanData(value=True)
+            for i, m in enumerate(result):
+                task_dict[f"result_{i}"] = m
+        
+        self.dataset["tasks"].add_row(task_dict, name=combined_arg_hash)
+        return result
+
+    def create_tasks(self, *args: Model):
+
+        task = asyncio.create_task(self._run_node(list(args)))
+        self._tasks.append(task)
+        return task
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if self._tasks:
+            await asyncio.gather(*self._tasks)
+        
+        db = context.db
+        await self.dataset.save(db)
+        self.info("MassRunner saving dataset")
+
+        #
+        # # Check if any task failed
+        # for idx in range(len(self.dataset['tasks'].data)):
+        #     item = await self.dataset['tasks'].get_item(idx)
+        #     if 'success' in item and hasattr(item['success'], 'value') and not item['success'].value:
+        #         return self.fail("One or more tasks failed")
+
+        return self.succeed()

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, context.db, 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/methods/switch_git.py

@@ -0,0 +1,16 @@
+from simstack.core.node import node
+from simstack.models import StringData
+from simstack.core.context import context
+
+@node
+def switch_git(branch: StringData, **kwargs) -> bool:
+    node_runner = kwargs["node_runner"]
+
+    project_root = context.project_root.resolve()
+
+    cmd = f"git stash && git pull --recurse-submodules && git checkout {branch.value}"
+
+    node_runner.run_command("git_switch", cmd, str(project_root))
+
+    return node_runner.success
+