simstack 2.0.1__py3-none-any.whl

simstack/core/resource_assignment.py

@@ -0,0 +1,222 @@
+import logging
+from dataclasses import dataclass
+from typing import Optional
+
+from odmantic import AIOEngine
+
+from simstack.models import NodeRegistry
+from simstack.models.parameters import Parameters, Resource, SlurmParameters
+from simstack.models.resource_assignment import (
+    ResourceAssignmentRule,
+    SlurmParametersPatch,
+)
+from simstack.util.db import Database
+
+logger = logging.getLogger("resource_assignment")
+
+SLURM_QUEUE_NAME = "slurm-queue"
+
+
+@dataclass
+class ResourceAssignmentResolution:
+    parameters: Parameters
+    normalized_call_path: str
+    matched_rule: Optional[ResourceAssignmentRule] = None
+
+
+def normalize_call_path(call_path: Optional[str]) -> str:
+    return (call_path or "").strip().lstrip(".")
+
+
+def _normalize_queue(queue: Optional[str]) -> str:
+    normalized = (queue or "").strip()
+    return normalized or "default"
+
+
+def _is_slurm_queue(queue: Optional[str]) -> bool:
+    return _normalize_queue(queue).lower() == SLURM_QUEUE_NAME
+
+
+def _clone_parameters(parameters: Optional[Parameters]) -> Parameters:
+    if isinstance(parameters, Parameters):
+        return parameters.model_copy(deep=True)
+    return Parameters()
+
+
+def _merge_slurm_patch(
+    patch: Optional[SlurmParametersPatch],
+) -> SlurmParameters:
+    if patch is None:
+        return SlurmParameters()
+    return SlurmParameters(**patch.model_dump(exclude_none=True))
+
+
+def empty_slurm_parameters() -> SlurmParameters:
+    cleared_values: dict[str, object] = {}
+    for field_name, field_info in SlurmParameters.model_fields.items():
+        if field_info.default_factory is not None:
+            cleared_values[field_name] = field_info.default_factory()
+        else:
+            cleared_values[field_name] = None
+    return SlurmParameters.model_validate(cleared_values)
+
+
+def _apply_assignment_patch(
+    base_parameters: Parameters,
+    rule: Optional[ResourceAssignmentRule],
+) -> Parameters:
+    effective = _clone_parameters(base_parameters)
+    if rule is None:
+        return effective
+
+    if rule.resource_str:
+        effective.resource = Resource(value=rule.resource_str)
+    if rule.queue is not None:
+        effective.queue = _normalize_queue(rule.queue)
+    if rule.slurm_parameters_patch:
+        effective.slurm_parameters = _merge_slurm_patch(
+            SlurmParametersPatch.model_validate(rule.slurm_parameters_patch),
+        )
+
+    return effective
+
+
+def normalize_and_validate_effective_parameters(
+    parameters: Optional[Parameters],
+) -> None:
+    if parameters is None:
+        return
+
+    if not _is_slurm_queue(getattr(parameters, "queue", None)):
+        parameters.slurm_parameters = empty_slurm_parameters()
+        return
+
+    slurm_parameters = getattr(parameters, "slurm_parameters", None)
+    if slurm_parameters is None:
+        raise ValueError('Slurm queue requires "slurm_parameters".')
+
+    fields_set: set[str] = getattr(slurm_parameters, "model_fields_set", set())
+    has_nodes = "nodes" in fields_set and slurm_parameters.nodes is not None
+    has_tasks = "tasks" in fields_set and slurm_parameters.tasks is not None
+    has_tasks_per_node = (
+        "tasks_per_node" in fields_set and slurm_parameters.tasks_per_node is not None
+    )
+    uses_default_nodes = (
+        not has_nodes
+        and not has_tasks
+        and not has_tasks_per_node
+        and slurm_parameters.nodes is not None
+    )
+
+    if uses_default_nodes:
+        has_nodes = True
+    elif not has_nodes:
+        slurm_parameters.nodes = None
+    if not has_tasks:
+        slurm_parameters.tasks = None
+    if not has_tasks_per_node:
+        slurm_parameters.tasks_per_node = None
+
+    if has_tasks and has_tasks_per_node:
+        raise ValueError(
+            'Slurm parameters conflict: use either "tasks" or "tasks_per_node".'
+        )
+    if not has_nodes and not has_tasks:
+        raise ValueError('Slurm requires at least one of "nodes" or "tasks".')
+
+
+def _select_matching_rule(
+    normalized_call_path: str,
+    rules: list[ResourceAssignmentRule],
+) -> Optional[ResourceAssignmentRule]:
+    enabled_rules = [rule for rule in rules if getattr(rule, "enabled", True)]
+    matching_rules = [
+        rule
+        for rule in enabled_rules
+        if ResourceAssignmentRule.matches_call_path(
+            rule.regex_pattern, normalized_call_path
+        )
+    ]
+    if not matching_rules:
+        return None
+
+    highest_score = max(
+        ResourceAssignmentRule.pattern_specificity_score(rule.regex_pattern)
+        for rule in matching_rules
+    )
+    top_rules = [
+        rule
+        for rule in matching_rules
+        if ResourceAssignmentRule.pattern_specificity_score(rule.regex_pattern)
+        == highest_score
+    ]
+    if len(top_rules) > 1:
+        rule_names = ", ".join(sorted(rule.name for rule in top_rules))
+        raise ValueError(
+            "Ambiguous resource assignment: "
+            f"multiple equally specific rules matched call_path '{normalized_call_path}': {rule_names}"
+        )
+    return top_rules[0]
+
+
+async def resolve_resource_assignment(
+    db: Database,
+    *,
+    call_path: Optional[str],
+    base_parameters: Optional[Parameters],
+    parent_parameters: Optional[Parameters] = None,
+) -> ResourceAssignmentResolution:
+    normalized_call_path = normalize_call_path(call_path)
+    effective_base = _clone_parameters(base_parameters)
+
+    if not normalized_call_path:
+        normalize_and_validate_effective_parameters(effective_base)
+        return ResourceAssignmentResolution(
+            parameters=effective_base,
+            normalized_call_path=normalized_call_path,
+            matched_rule=None,
+        )
+
+    rules = await db.find(ResourceAssignmentRule)
+    matched_rule = _select_matching_rule(normalized_call_path, list(rules))
+    effective_parameters = _apply_assignment_patch(effective_base, matched_rule)
+    normalize_and_validate_effective_parameters(effective_parameters)
+
+    if matched_rule is not None:
+        logger.info(
+            "Applied resource assignment rule '%s' to call_path '%s'",
+            matched_rule.name,
+            normalized_call_path,
+        )
+
+    return ResourceAssignmentResolution(
+        parameters=effective_parameters,
+        normalized_call_path=normalized_call_path,
+        matched_rule=matched_rule,
+    )
+
+
+async def apply_resource_assignment_to_node_registry(
+    db: Database,
+    node_registry: NodeRegistry,
+    *,
+    parent_parameters: Optional[Parameters] = None,
+) -> ResourceAssignmentResolution:
+    resolution = await resolve_resource_assignment(
+        db,
+        call_path=getattr(node_registry, "call_path", None),
+        base_parameters=getattr(node_registry, "parameters", None),
+        parent_parameters=parent_parameters,
+    )
+
+    node_registry.parameters = resolution.parameters
+    if resolution.matched_rule is None:
+        node_registry.assignment_rule_id = None
+        node_registry.assignment_rule_name = None
+        node_registry.assignment_pattern = None
+    else:
+        node_registry.assignment_rule_id = str(resolution.matched_rule.id)
+        node_registry.assignment_rule_name = resolution.matched_rule.name
+        node_registry.assignment_pattern = resolution.matched_rule.regex_pattern
+
+    return resolution

simstack/core/resources.py

@@ -0,0 +1,90 @@
+import logging
+from typing import List
+logger = logging.getLogger("resources")
+
+
+class AllowedResources:
+    """
+    Singleton class that holds a list of allowed resource strings.
+
+    This class is problematic because the @node decorator is called before the config is read.
+    @node may contain Parameters which set default values for resources.
+    The solution is to create the AllowedResources singleton class here, and set the resources from
+    context.initialize. Before context.initialize is called, any resource is allowed.
+
+    However, resources are validated on read, so before any node is executed, the resources must be set.
+    AllowedResources can be set only once.
+    """
+
+    _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
+        self._initialized = False
+
+    def set_resources(self, resources: List[str]) -> None:
+        """Set the list of allowed resources."""
+        if self._initialized and resources != self._resources:
+            raise RuntimeError("Resources can only be set once")
+        self._initialized = True
+        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()
+        self._initialized = False
+
+    def has_resource(self, resource: str) -> bool:
+        """Check if a resource exists in the list."""
+        if hasattr(self, "_initialized") and self._initialized:
+            return resource in self._resources
+        else:
+            return True # before this is initialized, any resource is allowed
+
+    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})"
+
+    @property
+    def initialized(self):
+        return hasattr(self, "_initialized")
+
+allowed_resources = AllowedResources()

simstack/core/route_table.py

@@ -0,0 +1,24 @@
+from typing import List, Dict
+
+
+class RouteTable:
+    _instance = None
+
+    def __init__(self):
+        if RouteTable._instance is not None:
+            raise RuntimeError("Use get_instance() instead")
+        self.targets = {}
+
+    @classmethod
+    def get_instance(cls) -> 'RouteTable':
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._instance
+
+    def add_route_set(self, source: str, targets: List[str]) -> None:
+        self.targets[source] = targets
+
+    def clear_routes(self) -> None:
+        self.targets.clear()
+
+route_table = RouteTable.get_instance()

simstack/core/run_docker.py

@@ -0,0 +1,109 @@
+from simstack.models import NodeRegistry
+from simstack.core.definitions import TaskStatus
+from simstack.core.context import context
+import asyncio
+import platform
+import subprocess
+import logging
+import os
+import locale
+
+logger = logging.getLogger("DockerRunner")
+
+async def run_docker(registry_entry: NodeRegistry) -> bool:
+    resource = context.config.resource
+    parameters = registry_entry.parameters
+
+
+    program_config = context.resource_config.get_program(registry_entry.name)
+    image= program_config.get("docker_image", None)
+    if image is None:
+        logger.error(f"Docker image not found for task_id={registry_entry.id}")
+        registry_entry.status = TaskStatus.FAILED
+        await context.db.save(registry_entry)
+        return False
+    docker_cmd = program_config.get("docker_cmd", None)
+    if docker_cmd is None:
+        docker_cmd = "docker"
+
+    if docker_cmd == "apptainer":
+        if not image.endswith(".sif") and not image.startswith("docker://"):
+             image += ".sif"
+
+    workdir = context.config.workdir
+    host_simstack_toml = context.config.project_root / "simstack.toml"
+
+    if docker_cmd == "docker":
+        cmd = [
+            "docker", "run",
+            "-e", f"SIMSTACK_DB_DATABASE={context.config.db_name}",
+            "-e", f"SIMSTACK_DB_TEST_DATABASE={context.config.db_name}",
+            "-e", f"SIMSTACK_DB_CONNECTION_STRING={context.config.connection_string}",
+            "-v", f"{workdir}:/root/simstack",
+            "-v", f"{host_simstack_toml}:/app/simstack.toml",
+            image,
+            "--node-id", str(registry_entry.id), "--resource", str(resource), "--project-root", "/app"
+        ]
+    elif docker_cmd == "apptainer":
+        cmd = [
+            "apptainer", "run",
+            "--env", f"SIMSTACK_DB_DATABASE={context.config.db_name}",
+            "--env", f"SIMSTACK_DB_TEST_DATABASE={context.config.db_name}",
+            "--env", f"SIMSTACK_DB_CONNECTION_STRING={context.config.connection_string}",
+            "--bind", f"{workdir}:/root/simstack",
+            "--bind", f"{host_simstack_toml}:/app/simstack.toml",
+            image,
+            "--node-id", str(registry_entry.id),
+            "--resource", str(resource),
+            "--project-root", "/app",
+        ]
+    else:
+        logger.error(f"Unsupported command {docker_cmd} for task_id={registry_entry.id}")
+        registry_entry.status = TaskStatus.FAILED
+        await context.db.save(registry_entry)
+        return False
+
+    # Use platform specific flags to ensure the process survives if runner is killed
+    creationflags = 0
+    if platform.system() == "Windows":
+        creationflags = subprocess.CREATE_NEW_PROCESS_GROUP | subprocess.DETACHED_PROCESS
+
+    try:
+        process = await asyncio.create_subprocess_exec(
+            *cmd,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            creationflags=creationflags,
+            start_new_session=True if platform.system() != "Windows" else False
+        )
+
+        stdout_b, stderr_b = await process.communicate()
+
+        enc = locale.getpreferredencoding(False) or "utf-8"
+        stdout = (stdout_b or b"").decode(enc, errors="replace").strip()
+        stderr = (stderr_b or b"").decode(enc, errors="replace").strip()
+
+        if process.returncode != 0:
+            logger.error(
+                "docker run failed for task_id=%s rc=%s stderr=%s stdout=%s cmd=%s",
+                registry_entry.id, process.returncode, stderr, stdout, cmd
+            )
+            registry_entry.status = TaskStatus.FAILED
+            await context.db.save(registry_entry)
+            return False
+
+        if stdout:
+            # For `docker run` without -d, stdout is the actual output of the command
+            logger.info("Docker container output for task_id=%s: %s", registry_entry.id, stdout)
+        else:
+            logger.info("Docker container spawned for task_id=%s", registry_entry.id)
+        if stderr:
+            # Some Docker setups warn on stderr even on success
+            logger.warning("docker run stderr for task_id=%s: %s", registry_entry.id, stderr)
+
+        return True
+    except Exception as e:
+        logger.exception(f"fatal error in running docker task_id: {registry_entry.id} {str(e)}")
+        registry_entry.status = TaskStatus.FAILED
+        await context.db.save(registry_entry)
+        return False

simstack/core/run_node.py

@@ -0,0 +1,71 @@
+import argparse
+import asyncio
+import logging
+import sys
+import os
+
+from simstack.core.context import context
+from simstack.core.definitions import TaskStatus
+from simstack.core.node import node_from_database
+from simstack.core.services.node_execution_service import run_node_from_registry
+
+logger = logging.getLogger("Run Node")
+
+async def run_node_from_id(node_id: str, resource_str: str, project_root: str = None):
+    """Run a single node by its ID from the database"""
+    logger.info(f"Initializing node run: node_id={node_id}, resource={resource_str}, project_root={project_root}")
+    try:
+        await context.initialize(resource=resource_str, project_root=project_root)
+    except Exception as e:
+        logger.error(f"Failed to initialize context: {str(e)}", exc_info=True)
+        return False
+    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
+        return await run_node_from_registry(registry_entry)
+    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():
+    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', stream=sys.stdout)
+    logger.info("Starting 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",
+    )
+
+    parser.add_argument(
+        "--project-root",
+        default=None,
+        nargs="?",
+        type=str,
+        help="project root directory",
+    )
+
+    args = parser.parse_args()
+
+
+    if args.node_id:
+        # Run a specific node once
+        asyncio.run(run_node_from_id(args.node_id, args.resource, args.project_root))
+
+
+if __name__ == "__main__":
+    run_node_main()

simstack/core/runner.py

@@ -0,0 +1,169 @@
+import argparse
+import asyncio
+import logging
+import os
+import sys
+
+from simstack.core.definitions import DBType
+
+if sys.version_info >= (3, 11):
+    import tomllib  # Python 3.11+
+else:
+    import tomli as tomllib  # type: ignore[import-not-found]
+
+from simstack.core.context import context
+from simstack.models.resource_definition import ResourceDefinition
+from simstack.tables.node_table import make_node_table
+from simstack.tables.model_table import make_model_table
+from simstack.core.services.runner_manager import RunnerManager
+
+logger = logging.getLogger("NodeRunner")
+
+
+async def initialize_default_resource() -> ResourceDefinition | None:
+    """
+    Checks if the current resource is the default one.
+    If so, syncs the node and model tables based on config.toml.
+    """
+    resource_def = await context.db.find_one(
+        ResourceDefinition,
+        ResourceDefinition.resource_str == str(context.config.resource),
+    )
+
+    if resource_def is None:
+        logger.warning(
+            "No ResourceDefinition found for '%s'; skipping default-resource initialization.",
+            str(context.config.resource),
+        )
+        return None
+
+    if resource_def.is_default:
+        config_path = context.config.project_root / "config.toml"
+        if not config_path.exists():
+            logger.warning(f"Default resource detected, but {config_path} not found.")
+            return resource_def
+
+        try:
+            with open(config_path, "rb") as f:
+                config_data = tomllib.load(f)
+
+            active_dirs = config_data.get("active_dirs", [])
+            if not active_dirs:
+                logger.info("No active_dirs found in config.toml.")
+                return resource_def
+
+            logger.info(f"Default resource: initializing tables for {active_dirs}")
+            await make_model_table(context.db, dirs=active_dirs)
+            await make_node_table(context.db, dirs=active_dirs)
+
+        except Exception as e:
+            logger.error(f"Failed to initialize default resource tables: {e}")
+
+    return resource_def
+
+
+async def async_main(args: argparse.Namespace) -> None:
+    """Async entry point"""
+    if args.connection_string == "none" or args.db_name == "none":
+        await context.initialize(resource=args.resource, config_file=args.config)
+    else:
+        await context.initialize(
+            resource=args.resource,
+            db_name=args.db_name,
+            connection_string=args.connection_string,
+            db_type=DBType.MONGODB,
+            config_file=args.config,
+        )
+
+    # Initialize tables if this is the default resource
+    resource_def = await initialize_default_resource()
+    is_default_resource = bool(resource_def and resource_def.is_default)
+
+    if args.resource:
+        logger.info(f"Setting resource for runner to {args.resource}")
+        runner_manager = RunnerManager(
+            context.config.resource,
+            detach=args.detach,
+            no_pull=args.no_pull,
+            is_default=is_default_resource,
+            with_file_transfer=args.with_file_transfer,
+        )
+        await runner_manager.run_nodes_for_resource(
+            args.polling_interval, 10, timeout=args.timeout
+        )
+
+
+def runner_main() -> None:
+    parser = argparse.ArgumentParser(description="Run nodes for a specific resource")
+    parser.add_argument(
+        "--config",
+        type=str,
+        # default="config.toml",
+        help="Path to the configuration file",
+    )
+
+    parser.add_argument(
+        "--resource",
+        type=str,
+        default="local",
+        help="Resource name to process tasks for",
+    )
+
+    parser.add_argument(
+        "--db-name",
+        type=str,
+        default="none",
+        help="Specify a non-standard database",
+    )
+
+    parser.add_argument(
+        "--connection-string",
+        type=str,
+        default="none",
+        help="Specify a non-standard connection string",
+    )
+
+    parser.add_argument(
+        "--polling-interval",
+        type=int,
+        default=20,
+        help="Interval in seconds between polling for new tasks",
+    )
+
+    parser.add_argument(
+        "--detach",
+        type=lambda x: (str(x).lower() not in ["false", "0", "no"]),
+        default=True,
+        help="If true (default), run nodes in an external process. Set to 'false' to run inline.",
+    )
+
+    parser.add_argument(
+        "--timeout",
+        type=int,
+        default=None,
+        help="Timeout in minutes after which the runner will terminate",
+    )
+
+    parser.add_argument(
+        "--no-pull",
+        action="store_true",
+        default=False,
+        help="If true, do not pull from git (GitUvUpdateService will not be started)",
+    )
+
+    parser.add_argument(
+        "--with-file-transfer",
+        type=lambda x: (str(x).lower() not in ["false", "0", "no"]),
+        default=True,
+        help="If true (default), start the FileTransferService.",
+    )
+
+    args = parser.parse_args()
+    # Run the async main function
+    asyncio.run(async_main(args))
+    pid = os.getpid()
+    logger.info(f"runner with pid {pid} shutting down normally")
+
+
+if __name__ == "__main__":
+    runner_main()

simstack/core/services/__init__.py

@@ -0,0 +1,11 @@
+from simstack.core.services.base_service import BaseService, RestartService
+from simstack.core.services.git_uv_update_service import GitUvUpdateService
+from simstack.core.services.node_execution_service import NodeExecutionService
+from simstack.core.services.runner_status_service import RunnerStatusService
+from simstack.core.services.resource_branch_monitor_service import ResourceBranchMonitorService
+from simstack.core.services.runner_cleanup_service import RunnerCleanupService
+from simstack.core.services.slurm_status_service import SlurmStatusService
+from simstack.core.services.git_restart_service import GitRestartService
+from simstack.core.services.stop_check_service import StopCheckService
+from simstack.core.services.timeout_restart_service import TimeoutRestartService
+from simstack.core.services.runner_manager import RunnerManager