simstack 0.2.0__py3-none-any.whl

simstack/models/runner_model.py

@@ -0,0 +1,39 @@
+from datetime import datetime
+from enum import Enum
+from typing import Optional, List
+
+from odmantic import Model, ObjectId, Field
+from pydantic import ConfigDict
+
+from simstack.models.parameters import Resource
+
+
+class RunnerEventEnum(str, Enum):
+    SHUTDOWN = "shutdown"
+    RUNNER_STARTED = "runner_started"
+    NODE_STARTED = "node_started"
+    NODE_SUBMIT = "node_submit"
+    ALIVE = "alive"
+    RESTART = "restart"
+    CRONTAB_OK = "crontab_ok"
+    CRONTAB_GONE = "crontab_gone"
+
+
+class RunnerType(str, Enum):
+    NODE_RUNNER = "node_runner"
+    RESOURCE_RUNNER = "resource_runner"
+
+
+class RunnerEvent(Model):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    timestamp: datetime = Field(default_factory=datetime.now)
+    event: RunnerEventEnum
+    resource: Resource
+    runner_type: RunnerType
+    hostname: Optional[str] = None
+    user: Optional[str] = None
+    pid: Optional[int] = None
+    node_id: Optional[ObjectId] = None
+    message: Optional[str] = None
+    git_status: List[str] = Field(default_factory=list)

simstack/models/simple_table.py

@@ -0,0 +1,32 @@
+from typing import Dict, Any, List
+
+from odmantic import Model, Field
+
+from simstack.models import simstack_model
+
+
+@simstack_model
+class SimpleTable(Model):
+    """
+    A simple table model to display tabular data using ag-grid
+    """
+
+    name: str = Field(default="SimpleTable")
+    heading: List[str] = Field(default_factory=list)
+    row: List[Dict[str, Any]] = Field(default_factory=list)
+    type: List[str] = Field(default_factory=list)
+
+    def add_column(self, column_name: str, column_type: str):
+        if column_name not in self.heading:
+            self.heading.append(column_name)
+            self.type.append(column_type)
+
+    def add_row(self, row: Dict[str, Any]):
+        self.row.append(row)
+
+    @classmethod
+    def ui_schema(cls) -> dict:
+        ui_schema = {
+            "ui:field": "SimpleTableField",
+        }
+        return ui_schema

simstack/models/simstack_model.py

@@ -0,0 +1,97 @@
+from functools import wraps
+from typing import TypeVar, Type, Any, get_type_hints, overload
+
+from simstack.util.cleaned_json_schema import cleaned_json_schema
+from simstack.util.custom_model_dump import custom_model_dump
+from simstack.util.default_from_dict import default_from_dict, default_from_model
+from simstack.util.generate_ui_schema import generate_ui_schema
+from simstack.util.ui_tools import ui_make_title
+
+T = TypeVar("T")
+
+
+@overload
+def simstack_model(cls: Type[T]) -> Type[T]:
+    ...
+
+
+def simstack_model(cls: T) -> T:
+    """
+    Decorates a given class to equip it with default implementations of utility
+    methods for handling operations such as dictionary conversion, schema
+    generation, and UI schema generation.
+    """
+
+    # Function to create a properly typed wrapper that preserves docstrings
+    def create_typed_wrapper(func, first_param_name="this_class"):
+        # Get original type hints
+        original_hints = get_type_hints(func)
+
+        # Create a wrapper with the correct parameters
+        @wraps(func)  # This preserves metadata like docstrings
+        def wrapper(*args, **kwargs):
+            return func(*args, **kwargs)
+
+        # Add explicit docstring if not preserved by wraps
+        if not wrapper.__doc__ and func.__doc__:
+            wrapper.__doc__ = func.__doc__
+
+        # Copy the original function's annotations
+        wrapper.__annotations__ = {
+            first_param_name: Type[Any],  # the first param is now the class
+            **{
+                k: v
+                for k, v in original_hints.items()
+                if k != "return" and k != first_param_name
+            },
+        }
+
+        # Preserve the return annotation if it exists
+        if "return" in original_hints:
+            wrapper.__annotations__["return"] = original_hints["return"]
+
+        return wrapper
+
+    # Create typed wrappers for all functions with explicit docstrings
+    default_class_methods = {
+        "json_schema": create_typed_wrapper(cleaned_json_schema),
+        "ui_schema": create_typed_wrapper(generate_ui_schema),
+        "ui_make_title": create_typed_wrapper(ui_make_title),
+        "from_dict": create_typed_wrapper(default_from_dict),
+        "from_model": create_typed_wrapper(default_from_model),
+        #'make_column_defs': create_typed_wrapper(make_column_defs_helper),
+    }
+
+    # Add methods only if they don't exist
+    for method_name, default_implementation in default_class_methods.items():
+        if not hasattr(cls, method_name):
+            setattr(cls, method_name, classmethod(default_implementation))
+
+    default_methods = {
+        "custom_model_dump": custom_model_dump,
+        #'make_table_entries': make_table_entries_helper
+    }
+
+    # Add methods only if they don't exist
+    for method_name, default_implementation in default_methods.items():
+        if not hasattr(cls, method_name):
+            setattr(cls, method_name, default_implementation)
+
+    # Add a marker attribute to identify decorated classes
+    setattr(cls, "_is_simstack_model", True)
+
+    # Return the original class without casting - the overload handles the typing
+    return cls
+
+
+def is_simstack_model(cls: Type) -> bool:
+    """
+    Check if a class has been decorated with @simstack_model.
+
+    Args:
+        cls: The class to check
+
+    Returns:
+        bool: True if the class was decorated with @simstack_model, False otherwise
+    """
+    return getattr(cls, "_is_simstack_model", False)

simstack/models/slurm_info.py

@@ -0,0 +1,17 @@
+import datetime
+
+from odmantic import Model, ObjectId
+
+from simstack.models.parameters import Resource
+
+
+class SlurmInfo(Model):
+    node_registry: ObjectId
+    updated: datetime.datetime
+    resource: Resource
+    job_id: str
+    name: str
+    user: str
+    code: str
+    time: str
+    nodes: list[str]

simstack/models/table_artifact.py

@@ -0,0 +1,87 @@
+from typing import List, Optional, Dict, Any
+
+from odmantic import Model, Field, EmbeddedModel, ObjectId
+
+
+class AGGridColumnDef(EmbeddedModel):
+    """AG-Grid column definition with support for nested tables."""
+
+    # Core properties
+    field: str = Field(..., description="The field name in the row data to bind to")
+    headerName: str = Field(..., description="Display name for the column header")
+
+    # Sizing properties
+    width: Optional[int] = Field(default=None, description="Column width in pixels")
+    minWidth: Optional[int] = Field(
+        default=None, description="Minimum column width in pixels"
+    )
+    maxWidth: Optional[int] = Field(
+        default=None, description="Maximum column width in pixels"
+    )
+    flex: Optional[int] = Field(
+        default=None, description="Flex sizing for responsive columns"
+    )
+
+    # Display and interaction properties
+    hide: Optional[bool] = Field(default=None, description="Whether to hide the column")
+    # pinned: Optional[Literal["left", "right"]] = Field(default=None, description="Pin column to left or right")
+    sortable: Optional[bool] = Field(
+        default=None, description="Whether the column can be sorted"
+    )
+    resizable: Optional[bool] = Field(
+        default=None, description="Whether the column can be resized"
+    )
+    editable: Optional[bool] = Field(
+        default=None, description="Whether cells in this column are editable"
+    )
+
+    # # Filtering and rendering
+    # filter: Optional[Union[bool, str]] = Field(default=None, description="Filter type")
+    # cellRenderer: Optional[str] = Field(default=None, description="Custom cell renderer component")
+    # cellRendererParams: Optional[Dict[str, Any]] = Field(default=None, description="Parameters for cell renderer")
+
+    # # Master-Detail specific properties for nested tables
+    # cellRendererSelector: Optional[Dict[str, Any]] = Field(default=None, description="Function to select cell renderer based on data")
+    #
+    # # Nested table configuration
+    # is_nested_table: Optional[bool] = Field(default=False, description="Whether this column contains nested table data")
+    # nested_table_config: Optional[Dict[str, Any]] = Field(default=None, description="Configuration for nested table rendering")
+    #
+    # # Row grouping and detail rows
+    # rowGroup: Optional[bool] = Field(default=None, description="Whether this column is used for row grouping")
+    # showRowGroup: Optional[Union[bool, str]] = Field(default=None, description="Show row group column")
+    #
+    # # Additional properties
+    # custom_properties: Optional[Dict[str, Any]] = Field(default_factory=dict, description="Additional custom properties")
+
+
+class TableArtifactModel(Model):
+    parent_id: ObjectId = Field(default=None, description="ID of the node registry")
+    columns_defs: List[AGGridColumnDef] = Field(
+        default_factory=list, description="AG-Grid column definitions"
+    )
+    row_data: List[Dict[str, Any]] = Field(
+        default_factory=list, description="List of row objects"
+    )
+
+    # Master-Detail configuration
+    master_detail: Optional[bool] = Field(
+        default=False, description="Enable master-detail rows"
+    )
+    detail_cell_renderer: Optional[str] = Field(
+        default=None, description="Detail panel cell renderer"
+    )
+    detail_cell_renderer_params: Optional[Dict[str, Any]] = Field(
+        default=None, description="Parameters for detail renderer"
+    )
+    # TODO implement nested tables and rendering
+    # # Nested tables storage - simplified to avoid recursion
+    # nested_tables: Optional[Dict[str, str]] = Field(default_factory=dict, description="Nested table IDs by reference key")
+    #
+    # # Tree data configuration
+    # tree_data: Optional[bool] = Field(default=False, description="Enable tree data structure")
+    # get_data_path: Optional[str] = Field(default=None, description="Function to get data path for tree structure")
+    #
+    # # Row grouping
+    # auto_group_column_def: Optional[AGGridColumnDef] = Field(default=None, description="Auto group column definition")
+    # group_default_expanded: Optional[int] = Field(default=None, description="Default expansion level for groups")

simstack/util/async_zip_utils.py

@@ -0,0 +1,18 @@
+async def async_zip(async_gen1, async_gen2):
+    """Async version of zip for two async generators"""
+    async for item1, item2 in async_zip_multiple(async_gen1, async_gen2):
+        yield item1, item2
+
+
+async def async_zip_multiple(*async_generators):
+    """Async version of zip for multiple async generators"""
+    iterators = [aiter(ag) for ag in async_generators]
+
+    while True:
+        try:
+            values = []
+            for it in iterators:
+                values.append(await anext(it))
+            yield tuple(values)
+        except StopAsyncIteration:
+            break

simstack/util/cleaned_json_schema.py

@@ -0,0 +1,20 @@
+def cleaned_json_schema(cls):
+    """
+
+    Generates a JSON schema for the given class and its fields, but eliminates
+    all fields which are models, embedded models, or references to models.
+
+    :param cls:
+    :return:
+    """
+    schema = cls.model_json_schema()
+    schema["title"] = cls.__name__
+
+    # for field_name, field in cls.model_fields.items():
+    #     field_type = field.annotation
+    #     if hasattr(field_type, '_is_simstack_model'):
+    #         # If the field is a model or embedded model, remove its proprties
+    #         if field_name in schema['$defs'] and "properties" in schema['$defs']['properties']:
+    #             del schema['$defs']['properties']
+
+    return schema

simstack/util/config_reader.py

@@ -0,0 +1,273 @@
+import os
+import sys
+import tomllib
+from pathlib import Path
+from typing import List, Dict
+
+from simstack.util.project_root_finder import find_project_root
+
+
+class ConfigReader:
+    """
+    Handles the loading and processing of a TOML configuration file,
+    to retrieve relevant settings for the workflow environment.
+    Provides critical parameters such as database name and connection string are provided, and
+    allows resource-based configurations for flexibility.
+
+    :ivar config: The dictionary representation of the parsed configuration file.
+    :type config: dict
+    :ivar _resource: Specifies the resource type or scope used to determine configurations.
+    :type _resource: str
+    :ivar _db_name: The name of the database specified in the configuration.
+    :type _db_name: str
+    :ivar _connection_string: The connection string for the database specified in the configuration.
+    :type _connection_string: str
+    :ivar _workdir: The working directory based on the resource settings from the configuration.
+    :type _workdir: Path | None
+    :ivar _python_path: The Python executable path defined in the resource settings from the configuration.
+    :type _python_path: str | None
+    :ivar _environment_start: Environment-specific start command extracted from the resource settings.
+    :type _environment_start: str | None
+    """
+
+    def __init__(self, **kwargs):
+        config_path = (
+            kwargs["config_path"] if "config_path" in kwargs else find_project_root()
+        )
+        config_file = (
+            kwargs["config_file"] if "config_file" in kwargs else "simstack.toml"
+        )
+        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)
+            elif kwargs.get("is_test", False):
+                self.config = {
+                    "paths": {
+                        "tests": {"path": "tests", "drops": "", "use_pickle": False}
+                    }
+                }
+            else:
+                print(f"Config file {toml_file} does not exist. Aborting.")
+                sys.exit(-1)
+
+        except tomllib.TOMLDecodeError:
+            print("There was an error decoding the TOML file.")
+            sys.exit(-1)
+
+        self._resource = kwargs.get("resource", "local")
+        self._is_test = kwargs.get("is_test", False)
+        self._secret_key = None
+        # parameter overrides config file
+        self._db_name = kwargs.get("db_name", None)
+
+        # for tests we can use an in_memory db
+        if self._db_name is None:
+            # the package simstack.toml has no db_name and connections string
+            self._db_name = (
+                self.config.get("parameters", {})
+                .get("common", {})
+                .get("database", "NONE")
+            )
+            if not self._is_test and self._db_name == "NONE":
+                print("You must specify a database name in the config file")
+                sys.exit(-1)
+
+        self._connection_string = kwargs.get("connection_string", None)
+        if self._connection_string is None:
+            self._connection_string = (
+                self.config.get("parameters", {})
+                .get("common", {})
+                .get("connection_string", "NONE")
+            )
+        if not self._is_test and self._connection_string == "NONE":
+            print("You must specify a connection string in the config file")
+            sys.exit(-1)
+
+        self._docker = False
+        self._workdir = None
+        self._external_workdir = None
+
+        self._python_path = None
+        self._external_source_dir = None
+        self._environment_start = None
+        self._routes = []
+        self._git = []
+
+    def secondary_init(self, workdir):
+        """
+        Called to perform secondary initialization after the context has been set.
+
+        This method is responsible for configuring certain parameters such as the
+        working directory (`workdir`), Python path (`python_path`), and the environment
+        start command (`environment_start`). It fetches these setups from the provided
+        configuration. If critical settings such as the `workdir` are not provided,
+        an error is logged, and a `ValueError` is raised. Logging is used throughout
+        to provide information about the initialization process.
+
+        :raises ValueError: If the `workdir` configuration for the resource is not specified.
+        """
+
+        import logging
+
+        logger = logging.getLogger("ConfigReader")
+
+        logger.info(
+            f"Initializing ConfigReader with resource: {self._resource} on database {self._db_name}"
+        )
+
+        self._docker = (
+            self.config.get("parameters", {})
+            .get(self._resource, {})
+            .get("docker", False)
+        )
+        logger.info(f"docker: {self._docker}")
+        if workdir is None:
+            workdir = (
+                self.config.get("parameters", {})
+                .get(self._resource, {})
+                .get("workdir", "NONE")
+            )
+        if workdir == "NONE":
+            logger.error(
+                f"You must specify a working directory for resource: {self._resource} in the config file"
+            )
+            raise ValueError(
+                f"You must specify a working directory for resource: {self._resource} in the config file"
+            )
+        self._workdir = Path(workdir)
+        logger.info(f"workdir: {self._workdir}")
+
+        if self._docker:
+            self._external_workdir = workdir
+            self._workdir = Path("/home/appuser/simstack")
+            logger.info(f"external_workdir: {self._external_workdir}")
+
+        self._external_source_dir = Path(
+            self.config.get("parameters", {})
+            .get(self._resource, {})
+            .get("source_dir", "NONE")
+        )
+        logger.info(f"source directory: {self._external_source_dir}")
+        if self.docker and self._external_source_dir == Path("NONE"):
+            logger.error(
+                f"You must specify an external source directory for resource: {self._resource} in the config file"
+            )
+            raise ValueError(
+                f"You must specify an external source directory for resource: {self._resource} in the config file"
+            )
+
+        self._python_path = (
+            self.config.get("parameters", {})
+            .get(self._resource, {})
+            .get("python_path", "NONE")
+        )
+        if self._python_path == "NONE":
+            logger.error("PYTHON PATH IS MISSING")
+        logger.info(f"python_path: {self._python_path}")
+
+        self._environment_start = (
+            self.config.get("parameters", {})
+            .get(self._resource, {})
+            .get("environment_start", "")
+        )
+        logger.info(f"environment_start: {self._environment_start}")
+
+        self._resources = (
+            self.config.get("parameters", {}).get("common", {}).get("resources", [])
+        )
+        logger.info(f"Initialized resources to: {self._resources}")
+
+        self._routes = self.config.get("routes", [])
+        for route in self._routes:
+            if not isinstance(route, dict):
+                logger.error(f"Route {route} is not a dictionary.")
+                raise ValueError("Route {route} is not a dictionary.")
+            if not ("source" in route and "target" in route and "host" in route):
+                logger.error(
+                    f"Route {route} does not contain 'source', 'target', 'host' keys."
+                )
+                raise ValueError(
+                    f"Route {route} does not contain 'source', 'target', 'host' keys."
+                )
+
+        self._secret_key = self.config.get("server", {}).get("secret_key", "")
+        self._git = (
+            self.config.get("parameters", {}).get(self._resource, {}).get("git", [])
+        )
+
+    def get_route(self, source: str, target: str) -> List[Dict[str, str]]:
+        """
+        Retrieves the route configuration for a given source and target.
+
+        :param source: The source node.
+        :param target: The target node.
+        :return: A list of dictionaries representing the route configuration.
+        """
+        for route in self._routes:
+            if route.get("source") == source and route.get("target") == target:
+                return route
+        return []
+
+    @property
+    def allowed_resources(self) -> List[str]:
+        return self._allowed_resources
+
+    @property
+    def secret_key(self) -> str:
+        return self._secret_key
+
+    @property
+    def docker(self) -> bool:
+        return self._docker
+
+    @property
+    def environment_start(self) -> str:
+        return self._environment_start
+
+    @property
+    def python_path(self) -> str:
+        return self._python_path
+
+    @property
+    def workdir(self) -> Path:
+        return self._workdir
+
+    @property
+    def external_workdir(self) -> Path:
+        return self._external_workdir
+
+    @property
+    def external_source_dir(self) -> Path:
+        return self._external_source_dir
+
+    @property
+    def git(self) -> List[Dict]:
+        return self._git
+
+    @property
+    def connection_string(self) -> str:
+        return self._connection_string
+
+    @property
+    def database_name(self) -> str:
+        return self._db_name
+
+    @property
+    def resource(self) -> str:
+        return self._resource
+
+    @resource.setter
+    def resource(self, value: str):
+        self._resource = value
+
+    @property
+    def paths(self) -> Dict:
+        """
+        Get the path configuration from the TOML file.
+
+        Returns:
+            Dictionary containing path configurations or an empty dict if not found
+        """
+        return self.config.get("paths", {})

simstack/util/custom_model_dump.py

@@ -0,0 +1,75 @@
+from typing import Dict, Any
+
+
+async def custom_model_dump(self, **kwargs) -> Dict[str, Any]:
+    """
+    Custom model dump method to handle the conversion of model instances to dictionaries.
+    This method recursively traverses dictionaries and lists to convert any nested
+    model instances to their dictionary representation.
+
+    :param self: The model instance
+    :param kwargs: Additional keyword arguments
+    :return: A dictionary representation of the model instance
+    """
+    # First get the dictionary from the standard model_dump
+
+    filtered_kwargs = {k: v for k, v in kwargs.items() if k != "engine"}
+    data = self.model_dump(**filtered_kwargs)
+
+    # Process each field - the key is to process the original values from self
+    for field_name, field in self.model_fields.items():
+        field_value = getattr(self, field_name)
+
+        # Skip None values
+        if field_value is None:
+            continue
+
+        data[field_name] = await _recursive_process(field_value, **kwargs)
+
+    return data
+
+
+async def _recursive_process(obj, **kwargs):
+    """
+    Helper function to recursively process an object that might contain models
+    """
+    if hasattr(obj, "custom_model_dump"):
+        return await obj.custom_model_dump(**kwargs)
+    elif isinstance(obj, dict):
+        return {k: await _recursive_process(v, **kwargs) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [await _recursive_process(item, **kwargs) for item in obj]
+    else:
+        return obj
+
+
+async def _process_container(container, **kwargs):
+    """
+    Helper function to recursively process dictionaries and lists,
+    converting any model instances to their dictionary representation.
+
+    :param container: The container to process (dict or list)
+    :param kwargs: Additional keyword arguments
+    :return: Processed container with model instances converted to dictionaries
+    """
+    if isinstance(container, dict):
+        result = {}
+        for k, v in container.items():
+            if hasattr(v, "custom_model_dump"):
+                result[k] = await v.custom_model_dump(**kwargs)
+            elif isinstance(v, (dict, list)):
+                result[k] = await _process_container(v, **kwargs)
+            else:
+                result[k] = v
+        return result
+    elif isinstance(container, list):
+        result = []
+        for item in container:
+            if hasattr(item, "custom_model_dump"):
+                result.append(await item.custom_model_dump(**kwargs))
+            elif isinstance(item, (dict, list)):
+                result.append(await _process_container(item, **kwargs))
+            else:
+                result.append(item)
+        return result
+    return container