simstack 2.0.1__py3-none-any.whl

simstack/models/parameters.py

@@ -0,0 +1,402 @@
+from enum import Enum
+from typing import Optional, List, ClassVar, Dict, Any
+from odmantic import Field, EmbeddedModel
+from pydantic import field_validator, model_validator
+
+from simstack.core.resources import allowed_resources
+
+
+def _validate_value_on_read(v: str) -> str:
+    """
+    Validate the value against allowed resources when it is read.
+    """
+    if not allowed_resources.has_resource(v):
+        allowed_str = ", ".join(repr(val) for val in allowed_resources.get_resources())
+        raise ValueError(f"Invalid resource value: {v!r}. Allowed values are: {allowed_str}")
+    return v
+
+
+class Resource(EmbeddedModel):
+    """
+    Resource whose value is validated against the allowed resources
+    only when the value is *read*, not when it is *set* or constructed.
+    """
+
+    # Regular model field, no leading underscore
+    value: str
+
+    def __getattribute__(self, name: str):
+        """
+        Intercept reads of `value` and validate at access time.
+        Setting/constructing does not perform validation.
+        """
+        if name == "value":
+            # Get raw stored value without triggering this override again
+            raw_value = object.__getattribute__(self, "__dict__").get("value")
+            return _validate_value_on_read(raw_value)
+        # Avoid triggering Pydantic deprecation warnings on model_fields
+        if name == "model_fields":
+            return self.__class__.model_fields
+        return object.__getattribute__(self, name)
+
+    def __str__(self) -> str:
+        """String representation of the resource (validated on access)."""
+        return self.value
+
+    def __repr__(self) -> str:
+        """Debug representation of the resource (validated on access)."""
+        return f"{self.__class__.__name__}(value={self.value!r})"
+
+    def __eq__(self, other) -> bool:
+        """Equal comparison operator (uses validated value)."""
+        if isinstance(other, Resource):
+            return self.value == other.value
+        elif isinstance(other, str):
+            return self.value == other
+        return False
+
+    def __ne__(self, other) -> bool:
+        """Not equal comparison operator."""
+        return not self.__eq__(other)
+
+
+class Queue(str, Enum):
+    DEFAULT = "default"
+    SLURM_QUEUE = "slurm-queue"
+    DOCKER = "docker"
+    SLURM_DOCKER = "slurm-docker"
+
+
+# TODO Fix Slurm Parameters
+class SlurmParameters(EmbeddedModel):
+    # Essential Resource Allocation Parameters
+    nodes: Optional[int] = Field(default=1, ge=1, description="Number of compute nodes")
+    tasks: Optional[int] = Field(
+        default=1, ge=1, description="Total number of tasks/processes"
+    )
+    tasks_per_node: Optional[int] = Field(
+        default=1, ge=1, description="Number of tasks per node"
+    )
+    cpus_per_task: Optional[int] = Field(
+        default=1, ge=1, description="Number of CPU cores per task"
+    )
+
+    # Memory Parameters
+    mem: Optional[str] = Field(
+        default="1G", description="Memory per node (e.g., '32G', '1024M')"
+    )
+    mem_per_cpu: Optional[str] = Field(
+        default=None, description="Memory per CPU core (e.g., '4G')"
+    )
+
+    # Time Parameters
+    time: Optional[str] = Field(
+        default="1:00:00",
+        description="Maximum runtime (e.g., '24:00:00', '1-12:00:00')",
+    )
+    begin: Optional[str] = Field(
+        default=None, description="Defer job start until specified time"
+    )
+
+    # Queue/Partition Parameters
+    partition: Optional[str] = Field(
+        default=None, description="Queue/partition to submit to"
+    )
+    qos: Optional[str] = Field(default=None, description="Quality of Service level")
+
+    # Job Information Parameters
+    job_name: Optional[str] = Field(default="simstack", description="Job name")
+    output: Optional[str] = Field(default=None, description="Standard output file path")
+    error: Optional[str] = Field(default=None, description="Standard error file path")
+
+    # Notification Parameters
+    mail_type: Optional[str] = Field(
+        default=None, description="Email notification triggers (BEGIN,END,FAIL,ALL)"
+    )
+    mail_user: Optional[str] = Field(
+        default=None, description="Email address for notifications"
+    )
+
+    # GPU Parameters
+    gres: Optional[str] = Field(
+        default=None, description="Generic resources (e.g., 'gpu:2', 'gpu:v100:1')"
+    )
+
+    # Advanced Scheduling Parameters
+    account: Optional[str] = Field(default=None, description="Billing account")
+    priority: Optional[int] = Field(default=None, description="Job priority")
+    reservation: Optional[str] = Field(
+        default=None, description="Use specific reservation"
+    )
+    constraint: Optional[str] = Field(
+        default=None, description="Node feature constraints"
+    )
+    exclusive: Optional[bool] = Field(default=None, description="Exclusive node access")
+    nice: Optional[int] = Field(default=None, description="Adjust scheduling priority")
+
+    # Dependency and Array Parameters
+    dependency: Optional[str] = Field(
+        default=None, description="Job dependencies (e.g., 'afterok:jobid')"
+    )
+    array: Optional[str] = Field(
+        default=None, description="Job arrays (e.g., '1-100', '1-10:2')"
+    )
+
+    # Additional Commands
+    startup_commands: List[str] = Field(
+        default_factory=list, description="Commands to run before main job"
+    )
+
+    # Working Directory
+    chdir: Optional[str] = Field(
+        default=None, description="Working directory for the job"
+    )
+
+    # Export Environment
+    export: Optional[str] = Field(
+        default=None, description="Environment variables to export"
+    )
+
+    # Signal handling
+    signal: Optional[str] = Field(
+        default=None, description="Signal to send when time limit is reached"
+    )
+
+    # Requeue options
+    requeue: Optional[bool] = Field(
+        default=None, description="Allow job to be requeued"
+    )
+    no_requeue: Optional[bool] = Field(
+        default=None, description="Prevent job from being requeued"
+    )
+
+    model_config: ClassVar[Dict[str, Any]] = {
+        "extra": "forbid",
+        "json_schema_extra": {
+            "title": "SlurmParameters",
+            "description": "Comprehensive parameters for Slurm job submission",
+            "examples": [
+                {
+                    "nodes": 2,
+                    "ntasks_per_node": 8,
+                    "cpus_per_task": 4,
+                    "mem": "64G",
+                    "time": "12:00:00",
+                    "partition": "compute",
+                    "gres": "gpu:2",
+                    "job_name": "my_simulation",
+                    "output": "job_%j.out",
+                    "error": "job_%j.err",
+                    "mail_type": "END,FAIL",
+                    "mail_user": "user@institution.edu",
+                }
+            ],
+        },
+    }
+
+    def to_sbatch_args(self) -> List[str]:
+        """Convert parameters to SBATCH arguments list."""
+        args = []
+
+        # Validate memory parameters first
+        memory_options = []
+
+        if self.mem_per_cpu is not None:
+            memory_options.append("mem_per_cpu")
+        elif self.mem is not None:
+            memory_options.append("mem")
+        if hasattr(self, "mem_per_gpu") and getattr(self, "mem_per_gpu") is not None:
+            memory_options.append("mem_per_gpu")
+
+        if len(memory_options) > 1:
+            raise ValueError(
+                f"SLURM memory parameters are mutually exclusive. "
+                f"Found: {', '.join(memory_options)}. "
+                f"Please specify only one of: --mem, --mem-per-cpu, or --mem-per-gpu"
+            )
+
+        # Ensure task count options are mutually exclusive
+        # If both total tasks (--ntasks) and tasks per node (--ntasks-per-node)
+        # are provided, prefer --ntasks-per-node and omit --ntasks.
+        tasks_value = self.tasks
+        tasks_per_node_value = self.tasks_per_node
+
+        if tasks_value is not None and tasks_per_node_value is not None:
+            tasks_value = None
+
+        # Map field names to SBATCH parameters
+        field_mapping = {
+            "nodes": "--nodes",
+            # NOTE: tasks and tasks_per_node are handled specially below to enforce
+            # mutual exclusivity when converting to sbatch arguments.
+            "tasks": "--ntasks",
+            "tasks_per_node": "--ntasks-per-node",
+            "cpus_per_task": "--cpus-per-task",
+            "mem": "--mem",
+            "mem_per_cpu": "--mem-per-cpu",
+            "time": "--time",
+            "begin": "--begin",
+            "partition": "--partition",
+            "qos": "--qos",
+            "job_name": "--job-name",
+            "output": "--output",
+            "error": "--error",
+            "mail_type": "--mail-type",
+            "mail_user": "--mail-user",
+            "gres": "--gres",
+            "account": "--account",
+            "priority": "--priority",
+            "reservation": "--reservation",
+            "constraint": "--constraint",
+            "dependency": "--dependency",
+            "array": "--array",
+            "chdir": "--chdir",
+            "export": "--export",
+            "signal": "--signal",
+        }
+
+        # Add mem_per_gpu to field mapping if it exists
+        if hasattr(self, "mem_per_gpu"):
+            field_mapping["mem_per_gpu"] = "--mem-per-gpu"
+
+        # Add parameters with values
+        for field_name, sbatch_param in field_mapping.items():
+            # Handle mutually exclusive task options
+            if field_name == "tasks":
+                value = tasks_value
+            elif field_name == "tasks_per_node":
+                value = tasks_per_node_value
+            else:
+                value = getattr(self, field_name, None)
+
+            if value is not None:
+                args.append(f"{sbatch_param}={value}")
+
+        # Add boolean flags
+        if self.exclusive:
+            args.append("--exclusive")
+        if self.requeue:
+            args.append("--requeue")
+        if self.no_requeue:
+            args.append("--no-requeue")
+        if self.nice is not None:
+            args.append(f"--nice={self.nice}")
+
+        return args
+
+    def to_sbatch_header(self) -> str:
+        """Convert parameters to SBATCH header string for script files."""
+        lines = ["#!/bin/bash"]
+
+        for arg in self.to_sbatch_args():
+            lines.append(f"#SBATCH {arg}")
+
+        if self.startup_commands:
+            lines.append("")
+            lines.extend(self.startup_commands)
+
+        return "\n".join(lines)
+
+
+class Parameters(EmbeddedModel):
+    force_rerun: bool = False
+    resource: Resource = Field(default_factory=lambda: Resource(value="self"))
+    queue: str = Field(default="default")
+    recompute_artifacts: Optional[bool] = Field(
+        default=False, description="Recompute artifacts for this node"
+    )
+    docker_image: Optional[str] = Field(default=None, description="Docker image to use")
+    other_value: str = Field(default="other")
+    test_dict: Dict[str, Any] = Field(default_factory=lambda: {"test": "value"})
+
+    slurm_parameters: SlurmParameters = Field(default=None)
+    # slurm_parameters_data: Dict[str, Any] = Field(default_factory=dict)
+
+    model_config = {
+        "extra": "forbid",
+        "json_schema_extra": {
+            "title": "Parameters",
+            "description": "Parameters for running a simulation",
+            "examples": [
+                {
+                    "resource": "self",
+                    "queue": "default",
+                    "slurm_parameters": {
+                        "nodes": 2,
+                    },
+                }
+            ],
+        },
+    }
+
+    @model_validator(mode="before")
+    @classmethod
+    def migrate_slurm_parameters(cls, data):
+        if "slurm_parameters_data" in data and "slurm_parameters" not in data:
+            data["slurm_parameters"] = SlurmParameters(**data["slurm_parameters_data"])
+            del data["slurm_parameters_data"]
+        if "slurm_parameters" not in data or data["slurm_parameters"] is None:
+            data["slurm_parameters"] = SlurmParameters()
+        return data
+
+    @field_validator("resource", mode="before")
+    @classmethod
+    def validate_resource(cls, v):
+        """
+        Validate and convert resource input to a Resource object.
+
+        Accepts string, Resource objects, and dictionary representations.
+        If a string is provided, converts it to a Resource object.
+        If a dictionary is provided (e.g., during deserialization),
+        extracts the value and creates a Resource object.
+
+        Args:
+            v: The value to validate (str, Resource, or dict)
+
+        Returns:
+            Resource: The validated Resource object
+        """
+        if isinstance(v, str):
+            return Resource(value=v)
+        elif isinstance(v, Resource):
+            return v
+        elif isinstance(v, dict):
+            # Handle deserialization from the dictionary
+            if "value" in v:
+                return Resource(value=v["value"])
+            else:
+                raise ValueError(
+                    f"Dictionary representation of resource must contain 'value' key, got {v}"
+                )
+        else:
+            raise ValueError(
+                f"resource must be a string, Resource object, or dictionary, got {type(v)}"
+            )
+
+    # # property getters and setters to handle the slurm_parameters
+    # @property
+    # def slurm_parameters(self) -> Optional[SlurmParameters]:
+    #     if not self.slurm_parameters_data:
+    #         default_slurm = SlurmParameters()
+    #         self.slurm_parameters_data = default_slurm.model_dump()
+    #     return SlurmParameters(**self.slurm_parameters_data)
+    #
+    # @slurm_parameters.setter
+    # def slurm_parameters(self, value: Optional[SlurmParameters]) -> None:
+    #     if value is None:
+    #         self.slurm_parameters_data = {}
+    #     else:
+    #         self.slurm_parameters_data = value.model_dump()
+    #
+    # # Add any convenience methods to work with slurm parameters
+    # def set_slurm_config(self, **kwargs):
+    #     """Update slurm parameters with the given keyword arguments."""
+    #     current = self.slurm_parameters
+    #     if current is None:
+    #         current = SlurmParameters()
+    #
+    #     for key, value in kwargs.items():
+    #         if hasattr(current, key):
+    #             setattr(current, key, value)
+    #
+    #     self.slurm_parameters = current

simstack/models/pickle_models.py

@@ -0,0 +1,300 @@
+import logging
+import sys
+import types
+from typing import Optional, Type, Any, Callable
+
+import cloudpickle
+from odmantic import Model
+from simstack.util.b64mixin import BytesB64Mixin
+
+logger = logging.getLogger("PickleModels")
+
+
+class ClassPickle(BytesB64Mixin, Model):
+    """
+    Persist an *arbitrary Python class* in MongoDB.
+
+    Fields
+    ------
+    name         – class __name__ (for reference / debugging)
+    module_path  – original module path (dotted)
+    pickle_data  – base64-encoded pickled bytes of the class
+
+    Methods
+    -------
+    store_class(cls)
+        Serialise and save the given class into `pickle_data`.
+    load_class()
+        Reconstruct the class object from `pickle_data`.
+    """
+
+    name: str
+    module_path: str
+    pickle_data: Optional[bytes] = None
+
+    def store_class(self, cls: Type[Any]) -> None:
+        self.name = cls.__name__
+        self.module_path = cls.__module__
+        self.pickle_data = cloudpickle.dumps(cls)
+        logger.debug("Stored class %s (%s)", self.name, self.module_path)
+
+    def load_class(self) -> Type[Any]:
+        if self.pickle_data is None:
+            raise ValueError("pickle_data is empty")
+        return cloudpickle.loads(self.pickle_data)
+
+
+class FunctionPickle(BytesB64Mixin, Model):
+    """
+    Persist an *arbitrary Python function* in MongoDB.
+
+    Fields
+    ------
+    name         – function __name__
+    module_path  – original module path
+    pickle_data  – base64 pickled bytes of the function
+    """
+
+    name: str
+    module_path: str
+    pickle_data: Optional[bytes] = None
+
+    def _is_problematic_object(self, obj) -> bool:
+        """Check if an object is problematic for pickling."""
+        try:
+            # Quick type checks
+            # if isinstance(obj, (asyncio.Future, concurrent.futures.Future, asyncio.AbstractEventLoop)):
+            #     return True
+            #
+            # Check type name for Windows-specific objects
+            type_name = type(obj).__name__
+            problematic_types = {
+                "_OverlappedFuture",
+                "ProactorEventLoop",
+                "WindowsProactorEventLoopPolicy",
+                "_ProactorBasePipeTransport",
+                "_ProactorSocketTransport",
+                "Handle",
+                "_WindowsSelectorEventLoop",
+                "_ProactorReadPipeTransport",
+            }
+
+            if type_name in problematic_types:
+                return True
+
+            # Check module origin
+            if hasattr(obj, "__class__") and hasattr(obj.__class__, "__module__"):
+                module_name = obj.__class__.__module__ or ""
+                if any(x in module_name for x in ["_overlapped", "_winapi"]):
+                    return True
+
+            return False
+
+        except Exception:
+            return True  # If we can't inspect it safely, consider it problematic
+
+    def _clean_object_recursively(self, obj, seen=None):
+        """Recursively clean an object, removing problematic references."""
+        if seen is None:
+            seen = set()
+
+        # Avoid infinite recursion
+        obj_id = id(obj)
+        if obj_id in seen:
+            return obj
+        seen.add(obj_id)
+
+        try:
+            if self._is_problematic_object(obj):
+                return None
+
+            # Handle different types
+            if isinstance(obj, dict):
+                cleaned = {}
+                for key, value in obj.items():
+                    if not self._is_problematic_object(value):
+                        cleaned_value = self._clean_object_recursively(value, seen)
+                        if cleaned_value is not None:
+                            cleaned[key] = cleaned_value
+                return cleaned
+
+            elif isinstance(obj, (list, tuple)):
+                cleaned_items = []
+                for item in obj:
+                    if not self._is_problematic_object(item):
+                        cleaned_item = self._clean_object_recursively(item, seen)
+                        if cleaned_item is not None:
+                            cleaned_items.append(cleaned_item)
+                return type(obj)(cleaned_items)
+
+            # For other objects, try to pickle test
+            try:
+                cloudpickle.dumps(obj)
+                return obj
+            except Exception:
+                return None
+
+        except Exception:
+            return None
+
+    def _clean_globals(self, func_globals: dict) -> dict:
+        """Remove unpickleable objects from function globals."""
+        cleaned_globals = {}
+
+        # Essential builtins that we need to keep
+        essential_builtins = {
+            "__builtins__",
+            "__name__",
+            "__doc__",
+            "__package__",
+            "print",
+            "len",
+            "str",
+            "int",
+            "float",
+            "bool",
+            "dict",
+            "list",
+            "tuple",
+            "Exception",
+            "ValueError",
+            "TypeError",
+            "AttributeError",
+        }
+
+        for key, value in func_globals.items():
+            try:
+                # Always skip certain names
+                skip_names = {"_OverlappedFuture", "_overlapped", "_winapi"}
+
+                if any(skip_name in key.lower() for skip_name in skip_names):
+                    continue
+
+                # Handle builtins specially
+                if key == "__builtins__":
+                    if isinstance(value, dict):
+                        # Keep only essential builtins
+                        cleaned_builtins = {
+                            k: v
+                            for k, v in value.items()
+                            if k in essential_builtins
+                            and not self._is_problematic_object(v)
+                        }
+                        cleaned_globals[key] = cleaned_builtins
+                    continue
+
+                # Skip modules entirely to avoid complex dependencies
+                if isinstance(value, types.ModuleType):
+                    continue
+
+                # Deep clean the object
+                cleaned_value = self._clean_object_recursively(value)
+                if cleaned_value is not None:
+                    # Final pickle test
+                    try:
+                        cloudpickle.dumps(cleaned_value)
+                        cleaned_globals[key] = cleaned_value
+                    except Exception as e:
+                        logger.debug(
+                            f"Skipping '{key}' after cleaning failed pickle test: {e}"
+                        )
+                else:
+                    logger.debug(f"Skipping '{key}': cleaned to None")
+
+            except Exception as e:
+                logger.debug(f"Error processing global '{key}': {e}")
+                continue
+
+        return cleaned_globals
+
+    def _clean_closure(self, closure):
+        """Clean closure variables of problematic objects."""
+        if not closure:
+            return closure
+
+        cleaned_closure = []
+        for cell in closure:
+            try:
+                cell_contents = cell.cell_contents
+                if not self._is_problematic_object(cell_contents):
+                    cleaned_contents = self._clean_object_recursively(cell_contents)
+                    if cleaned_contents is not None:
+                        # Create new cell with cleaned contents
+                        new_cell = types.CellType(cleaned_contents)
+                        cleaned_closure.append(new_cell)
+                    else:
+                        # Create cell with None if we had to remove the contents
+                        cleaned_closure.append(types.CellType(None))
+                else:
+                    # Replace problematic closure variable with None
+                    cleaned_closure.append(types.CellType(None))
+            except (ValueError, AttributeError):
+                # Cell is empty or has issues, create empty cell
+                cleaned_closure.append(types.CellType(None))
+            except Exception as e:
+                logger.debug(f"Error cleaning closure cell: {e}")
+                cleaned_closure.append(types.CellType(None))
+
+        return tuple(cleaned_closure) if cleaned_closure else None
+
+    def store_function(self, func: Callable) -> None:
+        self.name = func.__name__
+        self.module_path = func.__module__
+
+        try:
+            # Clean the function's globals and closure
+            cleaned_globals = self._clean_globals(func.__globals__)
+            cleaned_closure = self._clean_closure(func.__closure__)
+
+            # Create a copy of the function with cleaned components
+            func_copy = types.FunctionType(
+                func.__code__,
+                cleaned_globals,
+                func.__name__,
+                func.__defaults__,
+                cleaned_closure,
+            )
+            func_copy.__module__ = "__pickled_function__"
+
+            # Final test - try to pickle the cleaned function
+            test_pickle = cloudpickle.dumps(func_copy)
+
+            self.pickle_data = test_pickle
+            logger.debug("Stored function %s (%s)", self.name, self.module_path)
+
+        except Exception as e:
+            logger.error(f"Failed to pickle function {self.name}: {e}")
+            self.pickle_data = None
+            logger.warning(
+                f"Function {self.name} will use regular import instead of pickle"
+            )
+
+    def load_function(self) -> Callable:
+        if self.pickle_data is None:
+            raise ValueError("pickle_data is empty")
+
+        # Create a dummy module to satisfy cloudpickle's import requirements
+        if "__pickled_function__" not in sys.modules:
+            dummy_module = types.ModuleType("__pickled_function__")
+            sys.modules["__pickled_function__"] = dummy_module
+
+        func = cloudpickle.loads(self.pickle_data)
+
+        # Ensure the function has the correct __name__ attribute
+        if not hasattr(func, "__name__") or func.__name__ is None:
+            func.__name__ = self.name
+            logger.debug(f"Restored missing __name__ attribute: {self.name}")
+
+        # Ensure other essential attributes are present
+        if not hasattr(func, "__module__") or func.__module__ is None:
+            func.__module__ = self.module_path
+            logger.debug(f"Restored missing __module__ attribute: {self.module_path}")
+
+        # Ensure __qualname__ exists (needed for some introspection)
+        if not hasattr(func, "__qualname__"):
+            func.__qualname__ = self.name
+
+        logger.debug(
+            "Loaded function %s (%s)", self.name, getattr(func, "__name__", "unnamed")
+        )
+        return func