simstack 2.0.1__py3-none-any.whl

simstack/util/safe_code_executor.py

@@ -0,0 +1,194 @@
+import threading
+from typing import Dict
+
+from simstack.core.artifacts import ArtifactArguments
+
+
+def safe_code_executor(
+    code_string: str, artifact_arguments: ArtifactArguments, timeout: int = 30
+) -> Dict:
+    """
+    Safely executes Python code from a string with a controlled environment.
+
+    Args:
+        code_string (str): The Python code to execute
+        artifact_arguments (ArtifactArguments): Artifact arguments to be used in the execution environment
+        timeout (int, optional): Maximum execution time in seconds before timeout
+
+    Returns:
+        Dict: A dictionary containing:
+            'success' (bool): Whether execution was successful
+            'result': The return value if successful
+            'error': Error message if unsuccessful
+            'error_type': Type of error if unsuccessful
+
+    """
+    import ast
+    import traceback
+
+    # Default result structure
+    result = {"success": False, "result": None, "error": None, "error_type": None}
+
+    # Verify code is not empty
+    if not code_string or not code_string.strip():
+        result["error"] = "Empty code string provided"
+        result["error_type"] = "ValueError"
+        return result
+
+    # Create a safe globals dictionary
+    safe_globals = {"__builtins__": {}}
+
+    # Add only whitelisted builtins that are considered safe
+    for name in [
+        "abs",
+        "all",
+        "any",
+        "bool",
+        "dict",
+        "enumerate",
+        "filter",
+        "float",
+        "frozenset",
+        "int",
+        "isinstance",
+        "issubclass",
+        "len",
+        "list",
+        "map",
+        "max",
+        "min",
+        "print",
+        "range",
+        "reversed",
+        "round",
+        "set",
+        "slice",
+        "sorted",
+        "str",
+        "sum",
+        "tuple",
+        "zip",
+    ]:
+        safe_globals["__builtins__"][name] = __builtins__[name]
+
+    # Add artifact models to safe globals
+    from simstack.models.charts_artifact import (
+        ChartArtifactModel,
+        AGLineSeriesConfig,
+        AGChartAxisConfig,
+        AGChartTitleConfig,
+        AGChartLegendConfig,
+    )
+    from simstack.models.table_artifact import TableArtifactModel
+    from examples.science.electronic_structure.spectra.plot_spectra import (
+        make_multi_line_chart,
+    )
+    from simstack.models import ArtifactModel
+
+    safe_globals.update(
+        {
+            "ChartArtifactModel": ChartArtifactModel,
+            "TableArtifactModel": TableArtifactModel,
+            "AGLineSeriesConfig": AGLineSeriesConfig,
+            "AGChartAxisConfig": AGChartAxisConfig,
+            "AGChartTitleConfig": AGChartTitleConfig,
+            "AGChartLegendConfig": AGChartLegendConfig,
+            "make_multi_line_chart": make_multi_line_chart,
+            "ArtifactModel": ArtifactModel,
+            "ArtifactArguments": ArtifactArguments,
+        }
+    )
+
+    # Add artifact_arguments attributes to safe globals
+    # Extract all attributes from the artifact_arguments instance
+    for attr_name in dir(artifact_arguments):
+        if not attr_name.startswith("_"):  # Skip private/magic methods
+            attr_value = getattr(artifact_arguments, attr_name)
+            if not callable(attr_value):  # Skip methods, only include data attributes
+                safe_globals[attr_name] = attr_value
+
+    # Also add the full artifact_arguments object itself for backward compatibility
+    safe_globals["arg"] = artifact_arguments
+
+    # Verify code is syntactically correct
+    try:
+        ast.parse(code_string)
+    except SyntaxError as e:
+        result["error"] = str(e)
+        result["error_type"] = "SyntaxError"
+        return result
+
+    # Cross-platform timeout implementation using threading
+    execution_result = {
+        "completed": False,
+        "exception": None,
+        "local_vars": {},
+        "function_result": None,
+    }
+
+    def execute_code():
+        try:
+            local_vars = {}
+            exec(code_string, safe_globals, local_vars)
+            execution_result["local_vars"] = local_vars
+
+            # Look for a function in local_vars and call it with artifact_arguments
+            function_result = None
+            for var_name, var_value in local_vars.items():
+                if callable(var_value) and not var_name.startswith("_"):
+                    # Found a function, call it with artifact_arguments
+                    try:
+                        function_result = var_value(artifact_arguments)
+                        execution_result["function_result"] = function_result
+                        break
+                    except Exception as func_e:
+                        execution_result["exception"] = func_e
+                        execution_result["completed"] = True
+                        return
+
+            execution_result["completed"] = True
+        except Exception as e:
+            execution_result["exception"] = e
+            execution_result["completed"] = True
+
+    try:
+        # Start execution in a separate thread
+        thread = threading.Thread(target=execute_code, daemon=True)
+        thread.start()
+        thread.join(timeout=timeout)
+
+        if thread.is_alive():
+            # Timeout occurred
+            result["error"] = f"Code execution timed out after {timeout} seconds"
+            result["error_type"] = "TimeoutError"
+            return result
+
+        if execution_result["exception"]:
+            # Exception occurred during execution
+            e = execution_result["exception"]
+            result["error"] = str(e)
+            result["error_type"] = e.__class__.__name__
+            result["traceback"] = traceback.format_exc()
+            return result
+
+        # Successful execution
+        # Return the function result if available, otherwise return local vars
+        if execution_result["function_result"] is not None:
+            result["result"] = execution_result["function_result"]
+        else:
+            local_vars = execution_result["local_vars"]
+            if "result" in local_vars:
+                result["result"] = local_vars["result"]
+            else:
+                result["result"] = {
+                    k: v for k, v in local_vars.items() if not k.startswith("_")
+                }
+
+        result["success"] = True
+
+    except Exception as e:
+        result["error"] = str(e)
+        result["error_type"] = e.__class__.__name__
+        result["traceback"] = traceback.format_exc()
+
+    return result

simstack/util/sample_config.py

@@ -0,0 +1,98 @@
+_SAMPLE_SIMSTACK_TOML_DICT = {
+    "parameters": {
+        "common": {
+            "resources": [
+                "local",
+                "int-nano",
+                "horeka",
+                "justus",
+                "self",
+                "exchange",
+                "uploads",
+            ],
+            "database": "wolfgang_data",
+            "test_database": "wolfgang_test",
+            "connection_string": "mongodb://localhost:27017/",
+        },
+        "local": {
+            "ssh-key": "C:\\Users\\bj7610\\Documents\\etc\\.ssh\\surface11_openssh",
+            "resource": "local",
+            "workdir": "C:\\Users\\bj7610\\simstack",
+            "python_path": [
+                "C:\\Users\\bj7610\\PyCharmProjects\\simstack-model",
+                "C:\\Users\\bj7610\\PyCharmProjects\\simstack-model\\src",
+            ],
+        },
+        "self": {
+            "ssh-key": "C:\\Users\\bj7610\\Documents\\etc\\.ssh\\surface11_openssh",
+            "resource": "local",
+            "workdir": "C:\\Users\\bj7610\\simstack",
+            "python_path": [
+                "C:\\Users\\bj7610\\PyCharmProjects\\simstack-model",
+                "C:\\Users\\bj7610\\PyCharmProjects\\simstack-model\\src",
+            ],
+        },
+        "uploads": {
+            "ssh-key": "C:\\Users\\bj7610\\Documents\\etc\\.ssh\\surface11_openssh",
+            "resource": "self",
+            "workdir": "C:\\Users\\bj7610\\simstack",
+            "python_path": [
+                "C:\\Users\\bj7610\\PyCharmProjects\\simstack-model",
+                "C:\\Users\\bj7610\\PyCharmProjects\\simstack-model\\src",
+            ],
+        },
+        "int-nano": {
+            "ssh-key": "/home/ws/bj7610/.ssh/idrsa",
+            "workdir": "/home/ws/bj7610/simstack",
+            "python_path": [
+                "/home/ws/bj7610/projects/simstack-model",
+                "/home/ws/bj7610/projects/simstack-model/src",
+            ],
+            "environment_start": "conda activate simstack-env",
+        },
+    },
+    "server": {
+        "port": 8000,
+        "secret_key": "secret_key",
+        "upload_dir": "C:\\Users\\bj7610\\simstack\\uploads",
+    },
+    "hosts": {
+        "local": "localhost",
+        "int-nano": "int-nano.int.kit.edu",
+        "justus": "justus.int.kit.edu",
+        "horeka": "horeka.int.kit.edu",
+    },
+    "routes": [
+        {"source": "local", "target": "int-nano", "host": "local"},
+        {"source": "int-nano", "target": "local", "host": "local"},
+        {"source": "horeka", "target": "local", "host": "horeka"},
+        {"source": "local", "target": "horeka", "host": "horeka"},
+        {"source": "justus", "target": "local", "host": "justus"},
+        {"source": "local", "target": "justus", "host": "justus"},
+    ],
+    "paths": {
+        "models": {"path": "src/simstack/models", "drops": "src"},
+        "methods": {"path": "src/simstack/methods", "drops": "src"},
+        "tests": {"path": "tests", "drops": "", "use_pickle": True},
+    },
+}
+
+
+def sample_config_writer_entrypoint() -> None:
+    """
+    Entry point to write a sample simstack.toml configuration file.
+    """
+    import os
+    import toml
+
+    config_path = os.path.join(os.getcwd(), "simstack.toml")
+    if os.path.exists(config_path):
+        print(
+            f"simstack.toml already exists at {config_path}. Aborting to avoid overwrite."
+        )
+        return
+
+    with open(config_path, "w") as config_file:
+        toml.dump(_SAMPLE_SIMSTACK_TOML_DICT, config_file)
+
+    print(f"Sample simstack.toml written to {config_path}")

simstack/util/sensitive_content_filter.py

@@ -0,0 +1,48 @@
+from typing import Dict, Any
+
+
+def filter_file_content(data: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Recursively searches through nested data structures to filter out sensitive content.
+    Specifically looks for 'content' fields in dictionaries, especially within 'files' structures.
+
+    Args:
+        data: Any Python data structure (dict, list, etc.)
+
+    Returns:
+        The filtered data structure with sensitive content removed
+    """
+    if isinstance(data, dict):
+        # Create a new dict to avoid modifying the original during iteration
+        result = {}
+
+        # Check if this is a FileStack-like dict with content
+        if "content" in data:
+            # Make a copy of the dict without the content field
+            result = {
+                k: filter_file_content(v) for k, v in data.items() if k != "content"
+            }
+        else:
+            # Process all items recursively
+            for key, value in data.items():
+                result[key] = filter_file_content(value)
+
+        return result
+
+    elif isinstance(data, list):
+        # Process each item in the list recursively
+        return [filter_file_content(item) for item in data]
+
+    elif hasattr(data, "__dict__"):
+        # Handle object instances by converting to dict and processing
+        obj_dict = data.__dict__.copy()
+        filtered_dict = filter_file_content(obj_dict)
+
+        # This is tricky - we'd ideally create a new instance with the filtered data
+        # but that's complex without knowing the class structure
+        # So we'll return the filtered dict instead
+        return filtered_dict
+
+    else:
+        # For primitive types, return as is
+        return data

simstack/util/setup_logging.py

@@ -0,0 +1,49 @@
+import logging
+
+from simstack.util.db_logger import DBLogHandler
+
+
+def setup_logging(
+    connection_string: str,
+    db_name: str,
+    log_level=logging.INFO,
+    console=True,
+    log_format: str = "%(asctime)s - %(name)-15s - %(levelname)-10s - %(filename)-20s:%(lineno)4d - %(message)s",
+):
+    """
+    Configures the root logger with specified logging level, handlers for both database
+    and optionally the console.
+
+    :param log_format:
+    :param connection_string: A string representing the database connection.
+    :param db_name: A name of the database to log into.
+    :param log_level: Logging verbosity level, defaulting to logging.INFO.
+    :param console: A boolean indicating whether to log messages to the console.
+    :return: Configured root logger instance.
+    """
+    # Configure the root logger
+    root_logger = logging.getLogger()
+    root_logger.setLevel(log_level)
+
+    # Remove any existing handlers to avoid duplicates during reloads
+    for handler in root_logger.handlers[:]:
+        root_logger.removeHandler(handler)
+
+    # Create formatter
+    formatter = logging.Formatter(log_format, datefmt="%Y-%m-%d %H:%M:%S")
+
+    db_handler = DBLogHandler(connection_string, db_name, "logs")
+
+    # Add database handler
+    db_handler.setLevel(log_level)
+    db_handler.setFormatter(formatter)
+    root_logger.addHandler(db_handler)
+
+    # Optionally add console handler
+    if console:
+        console_handler = logging.StreamHandler()
+        console_handler.setFormatter(formatter)
+        console_handler.setLevel(log_level)
+        root_logger.addHandler(console_handler)
+
+    return root_logger

simstack/util/submit_to_watchdog.py

@@ -0,0 +1,106 @@
+# python
+import os
+import time
+from pathlib import Path
+from uuid import uuid4
+
+from pydantic import BaseModel
+
+POLL_INTERVAL = 1.0
+TIMEOUT_SEC = 60 * 60 * 72  # adjust
+
+
+class WatchdogResult(BaseModel):
+    job_id: str
+    status: str
+    returncode: int | None
+    stdout: str
+    stderr: str
+    paths: dict[str, str]
+
+
+def submit_to_watchdog(
+    command: str,
+    job_id: str,
+    queue_dir: Path = Path.cwd() / "queue",
+    env: dict | None = None,
+    cwd: str | None = None,
+) -> WatchdogResult:
+    """
+    Submit a command to be executed by the file-queue watcher.
+    Returns a WatchdogResult with status, exit_code, stdout, stderr, and paths.
+    """
+    base = queue_dir / job_id
+    cmd_file_tmp = base.with_suffix(".cmd.tmp")
+    cmd_file = base.with_suffix(".cmd")
+    ready_file = base.with_suffix(".ready")
+    done_file = base.with_suffix(".done")
+    fail_file = base.with_suffix(".fail")
+    exit_file = base.with_suffix(".exit")
+    out_file = base.with_suffix(".stdout")
+    err_file = base.with_suffix(".stderr")
+
+    spec_lines = []
+    if cwd:
+        spec_lines.append(f"CWD={cwd}")
+    if env:
+        for k, v in env.items():
+            spec_lines.append(f"ENV[{k}]={v}")
+    spec_lines.append("CMD=" + command)
+    data = "\n".join(spec_lines) + "\n"
+
+    queue_dir.mkdir(parents=True, exist_ok=True)
+    print(f"[watcher] Submitting {data} to {cmd_file}")
+    with open(cmd_file_tmp, "w", encoding="utf-8") as f:
+        f.write(data)
+        f.flush()
+        os.fsync(f.fileno())
+
+    os.replace(cmd_file_tmp, cmd_file)  # atomic
+    ready_file.touch()  # signal
+
+    deadline = time.time() + TIMEOUT_SEC
+    while time.time() < deadline:
+        if done_file.exists() or fail_file.exists():
+            break
+        time.sleep(POLL_INTERVAL)
+
+    result = WatchdogResult(
+        job_id=job_id,
+        status="timeout",
+        returncode=None,
+        stdout="",
+        stderr="",
+        paths={
+            "cmd": str(cmd_file),
+            "ready": str(ready_file),
+            "done": str(done_file),
+            "fail": str(fail_file),
+            "exit": str(exit_file),
+            "stdout": str(out_file),
+            "stderr": str(err_file),
+        },
+    )
+
+    if done_file.exists() or fail_file.exists():
+        exit_code: int | None = None
+        if exit_file.exists():
+            try:
+                exit_code = int(exit_file.read_text().strip())
+            except Exception:
+                exit_code = None
+        out = out_file.read_text(errors="replace") if out_file.exists() else ""
+        err = err_file.read_text(errors="replace") if err_file.exists() else ""
+        ok = done_file.exists() and (exit_code == 0)
+
+        result.status = "ok" if ok else "failed"
+        result.returncode = exit_code
+        result.stdout = out
+        result.stderr = err
+
+    return result
+
+
+if __name__ == "__main__":
+    res = submit_to_watchdog("dir", uuid4().hex)
+    print(res)

simstack/util/toml_reader.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+import logging
+import socket
+import sys
+import tomllib
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from simstack.core.resources import allowed_resources
+from simstack.core.route_table import route_table
+
+from simstack.util.path_manager import path_manager
+from simstack.util.transform_file_name import transform_file_name
+
+if TYPE_CHECKING:
+    from simstack.models.resource_definition import ResourceDefinition
+
+logger = logging.getLogger(__name__)
+
+
+class TomlReader:
+    def __init__(self, config_path: Path, config_file: Path = Path("simstack.toml")):
+        try:
+            if config_file.is_absolute():
+                toml_file = config_file
+            else:
+                toml_file = config_path / config_file
+            if toml_file.exists():
+                with open(toml_file, "rb") as f:
+                    self._config = tomllib.load(f)
+            else:
+                print(f"Config file {toml_file} does not exist. Aborting.")
+                sys.exit(-1)
+            self._config_path = config_path
+        except tomllib.TOMLDecodeError:
+            print("There was an error decoding the TOML file.")
+            sys.exit(-1)
+
+    @property
+    def config(self):
+        return self._config
+
+    def get(self, key: str, default=None):
+        keys = key.split(".")
+        value = self._config
+        for k in keys:
+            if isinstance(value, dict) and k in value:
+                value = value[k]
+            else:
+                return default
+        return value
+
+    def get_git_list(self):
+        return self.get("parameters.general.git", [])
+
+    def use_db(self):
+        return self.get("parameters.general.use_db", False)
+
+    def get_allowed_resources(self):
+        allowed_resources_list = self.get("resources.allowed_resources", None)
+        if allowed_resources_list is None:
+            raise ValueError("Allowed resources must be specified in the config file")
+        return allowed_resources_list
+
+    def get_resource_definition(self, resource_str) -> "ResourceDefinition":
+        """
+        Retrieves the resource definition for the given resource string from the TOML configuration.
+        Raises ValueError if the resource is not allowed or if the resource definition is not found.
+        Sets the allowed resources list from the config file.
+        """
+        from simstack.models.resource_definition import ResourceDefinition
+
+        if not allowed_resources.has_resource(resource_str):
+            raise ValueError(
+                f"Illegal resource {resource_str}. Allowed resources are: {allowed_resources.get_resources()}."
+            )
+
+        resource_definition = self.get(f"resources.{resource_str}", None)
+        if resource_definition is None:
+            raise ValueError(f"Resource definition for {resource_str} not found.")
+
+        resource_definition["resource_str"] = resource_str
+
+        if "workdir" not in resource_definition:
+            raise ValueError(f"No workdir specified for resource {resource_str}.")
+        else:
+            resource_definition["workdir"] = transform_file_name(
+                resource_definition["workdir"]
+            )
+
+        if "python_paths" not in resource_definition:
+            logger.warning(f"No python paths specified for resource {resource_str}.")
+        else:
+            resource_definition["python_paths"] = [
+                transform_file_name(p) for p in resource_definition["python_paths"]
+            ]
+
+        if "ssh_key" not in resource_definition:
+            logger.warning(f"No ssh key path specified for resource {resource_str}.")
+        else:
+            resource_definition["ssh_key"] = transform_file_name(
+                resource_definition["ssh_key"]
+            )
+
+        if "environment_start" not in resource_definition:
+            logger.warning(
+                f"No environment start command specified for resource {resource_str}."
+            )
+        if "routes" not in resource_definition:
+            logger.warning(f"No routes specified for resource {resource_str}.")
+            resource_definition["routes"] = []
+        if "hostname" not in resource_definition:
+            raise ValueError(f"No hostname specified for resource {resource_str}.")
+        elif (
+            resource_definition["hostname"] == "test_hostname"
+            or resource_definition["hostname"] == "localhost"
+        ):
+            resource_definition["hostname"] = socket.gethostname()
+            logger.info(
+                f"Overriding hostname for tests: {resource_definition['hostname']} "
+            )
+
+        return ResourceDefinition.model_validate(resource_definition)
+
+    def get_routes(self, resource_str: str):
+        return self.get(f"routes.{resource_str}", [])
+
+    def initialize_path_manager(self):
+        paths = self.get("paths", {})
+        for name, path_info in paths.items():
+            if isinstance(path_info, dict) and "path" in path_info:
+                path = Path(path_info["path"])
+                drops = path_info.get("drops", "")
+                path_use_pickle = path_info.get("use_pickle", False)
+                path_manager.add_path(name, path, drops, path_use_pickle)
+
+    def build_routes(self):
+        route_table.clear_routes()
+        for resource_str in allowed_resources.get_resources():
+            routes = self.get_routes(resource_str)
+            route_table.add_route_set(resource_str, routes)

simstack/util/transform_file_name.py

@@ -0,0 +1,41 @@
+import os
+from pathlib import Path
+from typing import Optional
+
+from simstack.util.project_root_finder import find_project_root
+
+def transform_file_name(path_str: str | Path, project_root: Path | None = None) -> Path:
+    """
+    Transform a string path by substituting environment variables.
+
+    Args:
+        path_str: String containing path with optional environment variables
+                ($HOME, $PROJECT, $TEMP)
+        project_root: Optional project root directory. If not provided, it will be automatically detected.
+
+    Returns:
+        Path object with resolved environment variables
+
+    Raises:
+        FileNotFoundError: If the resolved path does not exist
+    """
+
+    if isinstance(path_str, Path):
+        path_str = str(path_str)
+
+    if project_root is None:
+        project_root = find_project_root()
+
+    replacements = {
+        "$HOME": os.environ.get("HOME", str(Path.home())),
+        "$TEMP": os.environ.get("TEMP", os.environ.get("TMP", "/tmp")),
+        "$PROJECT": os.environ.get("PROJECT", str(project_root))
+    }
+
+    for var, value in replacements.items():
+        path_str = path_str.replace(var, value)
+
+    path = Path(path_str)
+    if not path.exists():
+        raise FileNotFoundError(f"Path does not exist: {path}")
+    return path