simstack 2.0.1__py3-none-any.whl

simstack/core/context.py

@@ -0,0 +1,283 @@
+import logging
+import sys
+from pathlib import Path
+from typing import TYPE_CHECKING
+from urllib.parse import urlparse, urlunparse
+
+from simstack.util.database_information import DatabaseInformation
+from simstack.util.db import DBType
+from simstack.util.project_root_finder import find_project_root
+from simstack.util.toml_reader import TomlReader
+
+from simstack.util.resource_config import ResourceConfig
+from simstack.util.setup_logging import setup_logging
+from simstack.util.mappings import ModelMappingTable, NodeMappingTable
+
+if TYPE_CHECKING:
+    from simstack.util.db import Database
+
+
+def remove_password_from_connection_string(connection_string):
+    parsed_url = urlparse(connection_string)
+
+    # Extract username and rebuild netloc without password
+    netloc = parsed_url.hostname
+    if parsed_url.username:
+        netloc = f"{parsed_url.username}@{netloc}"
+    if parsed_url.port:
+        netloc += f":{parsed_url.port}"
+
+    clean_url = parsed_url._replace(netloc=netloc)
+
+    return urlunparse(clean_url)
+
+
+async def initialize_git_list(db: "Database", toml_reader: TomlReader | None):
+    from simstack.models.resource_definition import GitRepo
+    git_list = await db.find_all(GitRepo)
+    if git_list is None and toml_reader is not None:
+        git_list = toml_reader.get("parameters.common.git", [])
+    return git_list
+
+
+class GlobalState:
+    _instance = None
+    _initialized = False
+
+    def __new__(cls, *args, **kwargs):
+        if not cls._instance:
+            cls._instance = super().__new__(cls)
+            cls._instance._initialized = False
+            cls._instance._db = None
+            cls._instance._log_handler = None
+            cls._instance._path_manager = None
+            cls._instance._config = None
+            cls._instance._model_mappings = None
+            cls._instance._node_mappings = None
+            cls._instance._resource_config = None
+        return cls._instance
+
+    def __init__(self, **kwargs):
+        """Initialize GlobalState instance.
+
+        This method only runs once due to the singleton pattern.
+        Use the initialize() method to set up the instance with database settings.
+        """
+        # We don't call self.initialize(**kwargs) here because it is async
+        # and __init__ must be sync. The user should call initialize() explicitly.
+
+        self._db = None
+        self._log_handler = None
+        self._path_manager = None
+        self._config = None
+        self._model_mappings = None
+        self._node_mappings = None
+        self._resource_config = None
+
+    def __getattribute__(self, name):
+        # These special attributes should always be accessible
+        if name in (
+            "_initialized",
+            "initialize",
+            "initialized",
+        ):
+            return object.__getattribute__(self, name)
+
+        # For other attributes, check initialization
+        if not object.__getattribute__(self, "_initialized"):
+            raise RuntimeError(
+                "GlobalState must be initialized with database settings before use"
+            )
+
+        return object.__getattribute__(self, name)
+
+    async def initialize(self, **kwargs):
+        """
+        Initializes the global state with the given configuration parameters.
+
+        Args:
+            **kwargs: Arbitrary keyword arguments for configuration. The following keys are expected:
+                - project_root (str, optional): The project root directory. If not provided, it will be
+                    determined using `find_project_root`.
+                - db_name (str, optional): The name of the database. Required when not using a TOML
+                    configuration.
+                - connection_string (str, optional): The connection string for the database. Required
+                    when not using a TOML configuration.
+                - db_type (DBType, optional): The type of the database. Required when not using a TOML
+                    configuration.
+                - is_test (bool, optional): Indicates whether the initialization is for testing purposes.
+                    Defaults to `False`.
+                - log_level (str, optional): Specifies the logging level. Defaults to `"INFO"`.
+                - resource (str, optional): Specifies the resource identifier for the configuration reader.
+                    Defaults to `"self"`.
+                - config_file (str, optional): Specifies the path to the configuration file, defaults to simstack.toml
+
+        Logic:
+            if all values for the DB are provided in the kwargs, use the provided values
+            otherwise the TOML file in the project root is used
+
+            The project root is determined using `find_project_root`, which by default looks for a
+            set of marker files, starting from the directory of find_project root and traversing up the directory tree.
+            In normal runs, it should skip the project root of the simstack package
+            This fails in the tests of the simstack package. The trick is to determine
+            project_root manually outside the call to initialize and pass it as a kwarg.
+
+            The TOML file decides whether the database or the file is used to set the variables
+            for the specific resource
+
+        """
+        is_test = kwargs.get("is_test", False)
+        resource_str = kwargs.get("resource", "self")
+
+        if self._initialized:
+            # If already initialized, we just return if not in test mode,
+            if not is_test:
+                return
+        self._initialized = True  # required because some later functions use context
+
+        project_root = kwargs.get("project_root", find_project_root())
+        if project_root is None:  # maybe None was passed
+            project_root = find_project_root()
+
+
+        project_root = Path(project_root)
+        kwargs["project_root"] = project_root
+
+        db_name: str | None = kwargs.get("db_name", None)
+        connection_string: str | None = kwargs.get("connection_string", None)
+        db_type: DBType | None = kwargs.get("db_type", None)
+
+        toml_reader = None
+        if db_name is None or connection_string is None or db_type is None:
+            # use not all info in the kwargs
+            simstack_toml_path = project_root / "simstack.toml"
+            toml_reader = TomlReader(project_root, config_file=Path(simstack_toml_path).resolve())
+            db_info = DatabaseInformation.from_config(toml_reader.config, **kwargs)
+        else:
+            db_info = DatabaseInformation(db_name, connection_string, db_type)
+
+        # check that the database can be reached and set logging up
+        self.initialize_database(db_info, is_test)
+        self.initialize_logging(db_info.connection_string, db_info.db_name, is_test, kwargs.get("log_level", "INFO"))
+
+        logger = logging.getLogger("Context")
+        if db_info.connection_string is not None:
+            db_type = DBType.MONGODB.value
+            safe_connection_string = remove_password_from_connection_string(db_info.connection_string)
+            logger.info(f"Database connection to {db_type} {safe_connection_string}/{db_name}")
+        else:
+            logger.info(f"Database connection in_memory {db_type}")
+
+        if not kwargs.get("skip_config", False):
+            await self.initialize_configs(self._db, toml_reader, **kwargs)
+
+        # Initialize memory-loaded mappings
+        await self.refresh_mappings()
+
+        self._initialized = True
+
+    async def initialize_configs(self, db, toml_reader, **kwargs):
+        project_root = kwargs["project_root"]
+        resource_str = kwargs.get("resource", "self")
+        is_test = kwargs.get("is_test", False)
+        try:
+            from simstack.util.config_reader import ConfigReader
+            from simstack.util.resource_config import ResourceConfig
+            resource_config_file = kwargs.get("config_file")
+            if resource_config_file is None:
+                resource_config_file = project_root / "config.toml"
+            else:
+                resource_config_file = Path(resource_config_file)
+            self._config = await ConfigReader.create(resource_str, db, toml_reader, **kwargs)
+            self._resource_config = ResourceConfig(resource_config_file, resource_str)
+        except Exception as e:
+            if is_test:
+                logger = logging.getLogger("Context")
+                logger.warning(f"Failed to initialize ConfigReader in test mode: {e}")
+            else:
+                raise e
+
+    async def refresh_mappings(self, *, models: bool = True, nodes: bool = True):
+        if models:
+            self._model_mappings = await ModelMappingTable.load(self._db)
+        if nodes:
+            self._node_mappings = await NodeMappingTable.load(self._db)
+
+    def initialize_logging(self, connection_string: str, db_name: str, is_test: bool, log_level: str = "INFO"):
+        if is_test:
+            # For tests, use simple console logging without the database handler
+            # We use force=True to ensure it overrides any existing configuration (e.g. from pytest or PyCharm)
+            # We also ensure the stream is sys.stderr so PyCharm/pytest can capture it properly if configured
+            logging.basicConfig(
+                level=log_level,
+                format="%(asctime)s - %(name)-15s - %(levelname)-10s - %(filename)-20s:%(lineno)4d - %(message)s",
+                stream=sys.stderr,
+                force=True,
+            )
+            self._log_handler = logging.getLogger()
+            # Flush stdout and stderr to ensure logs are not buffered
+            sys.stdout.flush()
+            sys.stderr.flush()
+        else:
+            self._log_handler = setup_logging(
+                connection_string,
+                db_name,
+                log_level,
+            )
+
+    def initialize_database(self, db_info: DatabaseInformation, is_test: bool):
+        from simstack.util.db import Database
+
+        try:
+            self._db = Database.from_db_info(db_info)
+            if self._db is not None and db_info.db_type == DBType.MONGODB:
+                # Only ping real MongoDB connections
+                self._db.client.admin.command("ping")
+
+        except ConnectionError as e:
+            if not is_test:
+                print(f"Could not connect to the database: {e}")
+                sys.exit(-1)
+            else:
+                # For tests, continue ignoring the database connection failure
+                print(f"Warning: Database connection failed in test mode: {e}")
+
+    @property
+    def db(self):
+        return self._db
+
+    @property
+    def config(self):
+        return self._config
+
+    @config.setter
+    def config(self, value):
+        self._config = value
+
+    @property
+    def model_mappings(self):
+        return self._model_mappings
+
+    @model_mappings.setter
+    def model_mappings(self, value):
+        self._model_mappings = value
+
+    @property
+    def node_mappings(self):
+        return self._node_mappings
+
+    @node_mappings.setter
+    def node_mappings(self, value):
+        self._node_mappings = value
+
+    @property
+    def resource_config(self) -> "ResourceConfig":
+        return self._resource_config
+
+    @property
+    def initialized(self):
+        return self._initialized
+
+
+# Create the singleton instance, but it's not initialized yet
+context = GlobalState()

simstack/core/definitions.py

@@ -0,0 +1,22 @@
+from enum import Enum
+
+
+class TaskStatus(str, Enum):
+    SUBMITTED = "submitted"
+    RUNNING = "running"
+    SLURM_QUEUED = "slurm_queued"
+    SLURM_RUNNING = "slurm_running"
+    RETRIEVED = "retrieved"
+    TIME_OUT = "timeout"
+    FAILED = "failed"
+    COMPLETED = "completed"
+    RECOVERED = "recovered"
+
+
+# TODO: eliminate
+class DBType(str, Enum):
+    SQLITE = "sqlite"
+    MONGODB = "mongodb"
+    POSTGRES = "postgres"
+    IN_MEMORY = "in_memory"
+    WITH_PATH = "with_path"

simstack/core/engine.py

@@ -0,0 +1,115 @@
+from contextvars import ContextVar
+from typing import Optional, Any, Iterable
+from odmantic import AIOEngine
+
+# TDOD engine remove AIOEngineProxy
+class AIOEngineProxy(AIOEngine):
+    """
+    A proxy engine that inherits all behavior from AIOEngine but overrides `save`.
+
+    The overridden `save` will:
+    - If the model has a `save` member function, call it with the engine instance.
+    - Else, if any direct part/attribute of the model has a `save` member function, call those with the engine instance.
+    - Otherwise, fall back to the original AIOEngine.save implementation.
+
+    Notes:
+    - Supports single model instances or iterables of model instances (list/tuple/set).
+    - Avoids recursion by delegating to super().save for the fallback.
+    """
+
+    async def save(self, obj: Any, *args, **kwargs) -> Any:  # type: ignore[override]
+        # Handle collections of models
+        if isinstance(obj, (list, tuple, set)):
+            results = []
+            for item in obj:
+                results.append(await self._save_one(item, *args, **kwargs))
+            return results
+        # Single model
+        return await self._save_one(obj, *args, **kwargs)
+
+    async def save_unchecked(self, obj: Any, *args, **kwargs) -> Any:
+        """
+        Save without checking for custom save methods.
+        This is a direct call to the original AIOEngine.save.
+        """
+        return await super().save(obj, *args, **kwargs)
+
+    async def _save_one(self, model: Any, *args, **kwargs) -> Any:
+        # 1) Try model's own `save`
+        if await self._maybe_call_custom_save(model):
+            return None
+
+        # 2) Try parts' `save` (shallow scan of attributes and common containers)
+        parts_saved = await self._call_parts_saves(model)
+
+        # 3) If neither the model nor any part handled saving, fallback to AIOEngine
+        if not parts_saved:
+            return await super().save(model, *args, **kwargs)
+        return None
+
+    async def _maybe_call_custom_save(self, target: Any) -> bool:
+        """
+        If `target` has a callable `save`, invoke it with this engine.
+        Tries positional engine and keyword engine forms, awaits if coroutine.
+        Returns True if a save was invoked, False otherwise.
+        """
+        if not hasattr(target, "save"):
+            return False
+        save_attr = getattr(target, "save")
+        if not callable(save_attr):
+            raise AttributeError(f"Model {target} has no callable `save` method")
+        return await target.save(self)
+
+    async def _call_parts_saves(self, model: Any) -> bool:
+        """
+        Inspect immediate parts/attributes of the model and call their `save` if present.
+        Handles:
+          - Direct attributes from vars(model)
+          - Items inside lists/tuples/sets
+          - Values of dicts
+        Returns True if at least one part save was called.
+        """
+        any_saved = False
+
+        def iter_parts(root: Any) -> Iterable[Any]:
+            # Direct attributes
+            try:
+                for v in vars(root).values():
+                    yield v
+            except TypeError:
+                # Objects without __dict__
+                pass
+
+            # For convenience, also scan common container-typed attributes one level deep
+            try:
+                for v in vars(root).values():
+                    if isinstance(v, (list, tuple, set)):
+                        for item in v:
+                            yield item
+                    elif isinstance(v, dict):
+                        for item in v.values():
+                            yield item
+            except TypeError:
+                pass
+
+        seen_ids = set()
+        for part in iter_parts(model):
+            pid = id(part)
+            if pid in seen_ids:
+                continue
+            seen_ids.add(pid)
+
+            if await self._maybe_call_custom_save(part):
+                any_saved = True
+
+        return any_saved
+
+#
+# current_engine_context: ContextVar[Optional[AIOEngineProxy]] = ContextVar(
+#     "current_engine", default=None
+# )
+#
+#
+# def get_current_engine_from_context() -> Optional[AIOEngineProxy]:
+#     """Get the current engine from context"""
+#     return current_engine_context.get()

simstack/core/find_simstack_modules.py

@@ -0,0 +1,41 @@
+import importlib
+import pkgutil
+from importlib.metadata import entry_points
+from logging import getLogger
+
+logger = getLogger("find_simstack_modules")
+
+
+def walk_packages(package_name, all_modules):
+    """Walk through all packages and modules."""
+    try:
+        package = importlib.import_module(package_name)
+        package_path = package.__path__
+
+        logger.info(f"Package: {package_name}")
+        # print(f"Path: {package_path}")
+
+        for importer, modname, ispkg in pkgutil.walk_packages(
+                package_path, prefix=f"{package_name}."
+        ):
+            if ispkg:
+                logger.debug(f"  Subpackage: {modname}")
+            else:
+                logger.debug(f"  Module: {modname}")
+                # Split module name by periods
+            all_modules.append(modname)
+
+    except Exception as e:
+        logger.error(f"Error walking {package_name}: {e}")
+
+
+def find_simstack_modules():
+    """Find all packages and modules within the simstack package."""
+
+    all_modules = []
+
+    entry_point_list = entry_points(group="simstack.modules")
+    logger.info(f"Entry points for simstack modules: {entry_point_list}.")
+    for entry_point in entry_point_list:
+        walk_packages(entry_point.value, all_modules)
+    return all_modules

simstack/core/hash.py

@@ -0,0 +1,179 @@
+import hashlib
+import inspect
+from collections.abc import Callable, Iterable, Mapping
+from typing import Any, cast
+
+from odmantic import ObjectId
+
+max_iterable_hash_count = 10000
+
+hash_exclusions = [
+    "pandas",
+    "numpy",
+    "ase",
+    "indigo",
+    "openbabel",
+    "pybel",
+    "pymatgen",
+    "ase",
+    "numpy",
+    "sqlalchemy",
+]
+
+
+HashResult = str | int
+
+
+def is_iterable(obj: Any) -> bool:
+    try:
+        iter(obj)
+        return True
+    except TypeError:
+        return False
+
+
+def is_primitive_type(obj: Any) -> bool:
+    return isinstance(obj, (int, float, bytes, bool, bytearray, type(None)))
+
+
+def hash_value(value: Any) -> str:
+    return hashlib.sha256(str(value).encode()).hexdigest()
+
+
+def hash_class_def(cls: type[Any]) -> str:
+    try:
+        source_code = inspect.getsource(cls)
+        source_hash = hashlib.sha256(source_code.encode("utf-8")).hexdigest()
+        return source_hash
+    except Exception:
+        mro = cls.__mro__
+        for mro_class in mro:
+            try:
+                source_code = inspect.getsource(mro_class)
+                source_hash = hashlib.sha256(source_code.encode("utf-8")).hexdigest()
+                return source_hash
+            except Exception:
+                pass
+        return "no source code"
+
+
+def hash_iterable(iterable: Iterable[Any]) -> str:
+    hash_text = ""
+    for count, item in enumerate(iterable):
+        hash_text += str(complex_hash_function(item))
+        if count > max_iterable_hash_count:
+            break
+    return hashlib.sha256(hash_text.encode("utf-8")).hexdigest()
+
+
+def hash_function_body(func: Callable[..., Any]) -> str:
+    # Get the source code of the function
+    source_code = inspect.getsource(func)
+    # Compute the hash of the source code
+    hash_object = hashlib.sha256(source_code.encode("utf-8"))
+    hash_digest = hash_object.hexdigest()
+    return hash_digest
+
+
+def hash_non_callable_members(instance: Any) -> dict[str, str]:
+    hashed_values: dict[str, str] = {}
+    for attr_name, attr_value in vars(instance).items():
+        if not callable(attr_value):
+            hashed_values[attr_name] = hashlib.sha256(
+                str(attr_value).encode()
+            ).hexdigest()
+    return hashed_values
+
+
+class ComplexHash:
+    def __init__(self, obj: Any) -> None:
+        self.hash_history: list[Any] = []
+
+    def hash_dict(self, obj: Mapping[Any, Any]) -> str:
+        hashed_values: dict[Any, HashResult] = {}
+        for k, v in obj.items():
+            hashed_values[k] = self.complex_hash(v)
+        combined_hash = "".join(f"{k}:{v}" for k, v in sorted(hashed_values.items()))
+        return hashlib.sha256(combined_hash.encode("utf-8")).hexdigest()
+
+    def hash_class(self, cls_obj: Any) -> str:
+        # check is the class name starts with a name in the hash_exclusions list
+        class_type = cls_obj.__class__
+        class_name = class_type.__module__ + "." + class_type.__name__
+        if class_type.__module__ == "builtins":
+            return "builtin"
+        if any([class_name.startswith(exclusion) for exclusion in hash_exclusions]):
+            return "excluded"
+        # print("hashing class", cls_obj.__class__.__name__)
+        class_hash = hash_class_def(cls_obj.__class__)
+        # TODO is __dict__ better than vars ?
+        # dict_hash = self.hash_dict(vars(cls_obj))
+
+        obj_dict = cls_obj.__dict__.copy()
+        if hasattr(cls_obj, "__class__") and hasattr(cls_obj.__class__, "__bases__"):
+            for base in cls_obj.__class__.__bases__:
+                if base.__name__ == "Model" and "odmantic.model" in str(
+                    base.__module__
+                ):
+                    if "id" in obj_dict:
+                        del obj_dict["id"]
+
+        dict_hash = self.hash_dict(obj_dict)
+        if hasattr(cls_obj, "model_extra") and cls_obj.model_extra is not None:
+            dict_hash = dict_hash + self.hash_dict(cls_obj.model_extra)
+        combined_hash = class_hash + dict_hash
+        return hashlib.sha256(combined_hash.encode("utf-8")).hexdigest()
+
+    def complex_hash(self, obj: Any) -> HashResult:
+        # TODO: are the functions of a class hashed correctly?
+
+        if isinstance(obj, type):
+            return hash_class_def(obj)
+        elif is_primitive_type(obj):
+            return hash(obj)
+        elif isinstance(obj, ObjectId):
+            return str(obj)
+        elif isinstance(
+            obj, str
+        ):  # strings are iterable but we want to hash them directly
+            return hash_value(obj)
+        elif inspect.isfunction(obj):
+            return hash_function_body(obj)
+        elif inspect.ismethod(obj):
+            if hasattr(obj, "__self__"):
+                if obj in self.hash_history:
+                    return "recursive"
+                self.hash_history.append(obj)
+                if hasattr(obj, "complex_hash"):
+                    return cast(HashResult, obj.complex_hash())
+                return self.hash_class(obj.__self__)
+            else:
+                return hash_function_body(obj)
+        # this should be instantiated classes
+        elif hasattr(obj, "__dict__"):
+            if obj in self.hash_history:
+                return "recursive"
+            self.hash_history.append(obj)
+            if hasattr(obj, "complex_hash"):
+                return cast(HashResult, obj.complex_hash())
+            return self.hash_class(obj)
+        elif isinstance(obj, dict):
+            return self.hash_dict(obj)
+        elif is_iterable(obj):
+            hashed_values: dict[int, HashResult] = {}
+            count = 0
+            for i, item in enumerate(obj):
+                hashed_values[i] = self.complex_hash(item)
+                count += 1
+                if count > max_iterable_hash_count:
+                    break
+            combined_hash = "".join(
+                f"{k}:{v}" for k, v in sorted(hashed_values.items())
+            )
+            return hashlib.sha256(combined_hash.encode("utf-8")).hexdigest()
+        else:
+            return hash_value(obj)
+
+
+def complex_hash_function(obj: Any) -> HashResult:
+    return ComplexHash(obj).complex_hash(obj)