simstack 0.2.0__py3-none-any.whl

simstack/core/recompute_artifacts.py

@@ -0,0 +1,120 @@
+import logging
+
+from simstack.core.artifacts import ArtifactArguments, create_artifacts
+from simstack.core.definitions import TaskStatus
+from simstack.core.engine import current_engine_context
+from simstack.core.node import node_from_database
+from simstack.models import NodeRegistry, ArtifactModel
+from simstack.models.charts_artifact import ChartArtifactModel
+from simstack.models.table_artifact import TableArtifactModel
+from simstack.util.importer import import_class
+from simstack.util.importer import import_function
+
+logger = logging.getLogger("recompute_artifacts")
+
+
+async def recompute_artifacts(node_registry: NodeRegistry):
+    """
+    Recomputes artifacts for a node and all its children recursively.
+
+    First processes all children nodes recursively, then removes and recomputes
+    the current node's artifacts.
+
+    :param node_registry: The registry entry for the node to process
+    :type node_registry: NodeRegistry
+    """
+    # Create Node from the registry
+
+    engine = current_engine_context.get()
+    node = await node_from_database(node_registry)
+    if node is None:
+        logger.error(f"Failed to create node from registry task_id: {node_registry.id}")
+        return
+
+    if node.status != TaskStatus.COMPLETED:
+        logger.error(f"Cannot recompute artifacts for task_id: {node_registry.id} name: {node.name} status: {node.status}")
+        return
+
+    # Find all children of this node
+    children = await engine.find(NodeRegistry, NodeRegistry.parent_ids.in_([node_registry.id]))
+
+    # Recursively recompute artifacts for all children first
+    for child_registry in children:
+        await recompute_artifacts(child_registry)
+
+    # Remove current node's artifacts
+    if node_registry.artifact_ids:
+        logger.info(f"Removing {len(node_registry.artifact_ids)} artifacts for node {node_registry.id}")
+
+        table_artifacts = await engine.find(
+            TableArtifactModel, TableArtifactModel.parent_id == node_registry.id
+        )
+        for table_artifact in table_artifacts:
+            await engine.delete(table_artifact)
+
+        chart_artifacts = await engine.find(
+            ChartArtifactModel, ChartArtifactModel.parent_id == node_registry.id
+        )
+        for chart_artifact in chart_artifacts:
+            await engine.delete(chart_artifact)
+        # Delete artifacts from the database
+        for artifact_id in node_registry.artifact_ids:
+            instance = await engine.find_one(
+                ArtifactModel, ArtifactModel.id == artifact_id
+            )
+            if instance:
+                await engine.delete(instance)
+            else:
+                logger.warning(
+                    f"task_id: {node_registry.id} Failed to delete artifact {artifact_id} from database"
+                )
+        node_registry.artifact_ids = []
+
+    # Recompute artifacts for this node
+    if node_registry.status == TaskStatus.COMPLETED:
+        logger.info(f"Recomputing artifacts for node {node_registry.id}")
+
+        # Load the result to create new artifacts
+        result = await node.load_results()
+        if result is not None:
+            artifact_arguments = ArtifactArguments(result, node_registry.id)
+            # Reconstruct the function arguments for artifact creation
+            args = []
+            for table, table_id in zip(
+                node_registry.input_tables, node_registry.input_ids
+            ):
+                model = await import_class(table)
+                arg = await engine.find_one(model, model.id == table_id)
+                if arg:
+                    args.append(arg)
+
+            # Get the function for artifact creation
+            wrapped_func = await import_function(node_registry.func_mapping)
+            func = (
+                wrapped_func
+                if not hasattr(wrapped_func, "_inner")
+                else wrapped_func._inner
+            )
+
+            # Create node kwargs similar to run_local
+            node_kwargs = {
+                "parent_id": node_registry.id,
+                "task_id": node_registry.id,
+                "call_path": node_registry.call_path,
+                "parent_parameters": node_registry.parameters,
+            }
+
+            artifact_arguments.add_attributes(func, *args, **node_kwargs)
+            node_registry.artifact_ids = await create_artifacts(
+                artifact_arguments, node_registry
+            )
+
+            # Save the updated registry
+            await engine.save(node_registry)
+            logger.info(
+                f"Recomputed {len(node_registry.artifact_ids)} artifacts for node {node_registry.id}"
+            )
+    else:
+        logger.warning(
+            f"Node {node_registry.id} is not completed, cannot recompute artifacts"
+        )

simstack/core/reset_db.py

@@ -0,0 +1,12 @@
+from simstack.core.context import context
+
+
+async def main():
+    context.initialize()
+    await context.db.reset_database()
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    asyncio.run(main())

simstack/core/resources.py

@@ -0,0 +1,101 @@
+import logging
+import os
+import sys
+import tomllib
+from typing import List
+
+from simstack.core.config_file import get_config_file
+from simstack.util.project_root_finder import find_project_root
+
+logger = logging.getLogger("resources")
+
+
+class AllowedResources:
+    """
+    Singleton class that holds a list of allowed resource strings.
+    # This class is independent of context because its initialized to an empty list.
+    Creating any resource will fail if the config has not been read
+    """
+
+    _instance = None
+    _resources: List[str] = []
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super(AllowedResources, cls).__new__(cls)
+        return cls._instance
+
+    def __init__(self):
+        # Only initialize once
+        if not hasattr(self, "_initialized"):
+            config_file = get_config_file().config_file_path
+            config_path = find_project_root()
+
+            try:
+                toml_file = os.path.join(config_path, config_file)
+                if os.path.exists(toml_file):
+                    with open(toml_file, "rb") as f:
+                        self.config = tomllib.load(f)
+                else:
+                    print("simstack.toml file not found in the project root.")
+                    sys.exit(-1)
+            except tomllib.TOMLDecodeError:
+                print("There was an error decoding the TOML file.")
+                sys.exit(-1)
+
+            self._resources = (
+                self.config.get("parameters", {}).get("common", {}).get("resources", [])
+            )
+            logger.info(f"Initialized resources to: {self._resources}")
+            self._initialized = True
+
+    def set_resources(self, resources: List[str]) -> None:
+        """Set the list of allowed resources."""
+        self._resources = resources.copy() if resources else []
+
+    def get_resources(self) -> List[str]:
+        """Get the list of allowed resources."""
+        return self._resources.copy()
+
+    def add_resource(self, resource: str) -> None:
+        """Add a single resource to the list."""
+        if resource not in self._resources:
+            self._resources.append(resource)
+
+    def remove_resource(self, resource: str) -> None:
+        if resource == "self":
+            raise ValueError("Cannot remove 'self' resource")
+        """Remove a resource from the list."""
+        if resource in self._resources:
+            self._resources.remove(resource)
+
+    def clear_resources(self) -> None:
+        """Clear all resources from the list."""
+        self._resources.clear()
+
+    def has_resource(self, resource: str) -> bool:
+        """Check if a resource exists in the list."""
+        return resource in self._resources
+
+    def __len__(self) -> int:
+        """Return the number of resources."""
+        return len(self._resources)
+
+    def __iter__(self):
+        """Make the class iterable."""
+        return iter(self._resources)
+
+    def __contains__(self, resource: str) -> bool:
+        """Support 'in' operator."""
+        return resource in self._resources
+
+    def __str__(self) -> str:
+        """String representation."""
+        return f"AllowedResources({self._resources})"
+
+    def __repr__(self) -> str:
+        """String representation."""
+        return f"AllowedResources({self._resources!r})"
+
+
+allowed_resources = AllowedResources()

simstack/core/run_node.py

@@ -0,0 +1,69 @@
+import argparse
+import asyncio
+import logging
+
+from simstack.core.context import context
+from simstack.core.definitions import TaskStatus
+from simstack.core.node import node_from_database
+
+logger = logging.getLogger("SlurmRunner")
+
+
+async def run_node(node_id: str, **kwargs):
+    """Run a single node by its ID from the database"""
+    registry_entry = None
+    try:
+        registry_entry = await context.db.load_task_by_id(node_id)
+
+        if not registry_entry:
+            logger.error(f"Node with ID {node_id} not found in the database")
+            return False
+        # Create node from the registry entry
+        node = await node_from_database(registry_entry)
+        if not node:
+            logger.error(
+                f"Failed to create node from registry entry task_id: {registry_entry.id}"
+            )
+            registry_entry.status = TaskStatus.FAILED
+            await context.db.save(registry_entry)
+            return False
+        registry_entry = node.registry_entry # it may have changed
+        # if the node was recovered we do not have to run it again
+        if node.status == TaskStatus.SUBMITTED or node.status == TaskStatus.SLURM_QUEUED or node.status == TaskStatus.SLURM_QUEUED:
+            await node.execute_node_locally()
+        else:
+            logger.info(f"task_id: {registry_entry.id} skipping task: {registry_entry.name} with status {registry_entry.status}")
+        return node.status == TaskStatus.COMPLETED
+    except Exception as e:
+        logger.exception(f"Error running node task_id: {node_id}: {str(e)}")
+        if registry_entry:
+            registry_entry.status = TaskStatus.FAILED
+            await context.db.save(registry_entry)
+        return False
+
+def run_node_main():
+    parser = argparse.ArgumentParser(description="Run nodes for a specific resource")
+    parser.add_argument(
+        "--node-id",
+        type=str,
+        help="Specific node ID to run (overrides resource-based polling)",
+    )
+
+    parser.add_argument(
+        "--resource",
+        default="local",
+        nargs="?",
+        type=str,
+        help="resource to load",
+    )
+
+    args = parser.parse_args()
+    context.initialize(resource=args.resource)
+
+    if args.node_id:
+        # Run a specific node once
+        asyncio.run(run_node(args.node_id))
+
+
+if __name__ == "__main__":
+    run_node_main()