gwsim 0.1.0__py3-none-any.whl

gwsim/cli/utils/checkpoint.py

@@ -0,0 +1,178 @@
+"""Checkpoint management for simulation recovery."""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+from typing import Any
+
+from gwsim.data.serialize.decoder import Decoder
+from gwsim.data.serialize.encoder import Encoder
+
+logger = logging.getLogger("gwsim")
+
+
+class CheckpointManager:
+    """Manages checkpoint files for simulation recovery.
+
+    A checkpoint is created after each successfully completed batch,
+    allowing resumption from that point if the simulation is interrupted.
+
+    Checkpoint file format:
+    {
+        "completed_batch_indices": [0, 1, 2, ...],
+        "last_simulator_name": "signal",
+        "last_completed_batch_index": 2,
+        "last_simulator_state": {...}
+    }
+
+    The checkpoint is written atomically:
+    1. Write to .tmp file
+    2. Backup existing checkpoint to .bak
+    3. Rename .tmp to checkpoint file
+    This ensures we never have a corrupted checkpoint.
+    """
+
+    def __init__(self, checkpoint_directory: Path):
+        """Initialize checkpoint manager.
+
+        Args:
+            checkpoint_directory: Directory to store checkpoint files
+        """
+        self.checkpoint_directory = Path(checkpoint_directory)
+        self.checkpoint_directory.mkdir(parents=True, exist_ok=True)
+        self.checkpoint_file = self.checkpoint_directory / "simulation.checkpoint.json"
+        self.checkpoint_tmp = self.checkpoint_directory / "simulation.checkpoint.json.tmp"
+        self.checkpoint_backup = self.checkpoint_directory / "simulation.checkpoint.json.bak"
+
+    def load_checkpoint(self) -> dict[str, Any] | None:
+        """Load checkpoint from file if it exists.
+
+        Returns:
+            Checkpoint dict with keys:
+            - completed_batch_indices: List of completed batch indices
+            - last_simulator_name: Name of last simulator
+            - last_completed_batch_index: Index of last completed batch
+            - last_simulator_state: State dict of last simulator
+            None if no checkpoint exists or checkpoint is corrupted
+        """
+        # Try to restore from backup if checkpoint doesn't exist but backup does
+        if not self.checkpoint_file.exists() and self.checkpoint_backup.exists():
+            logger.warning("Checkpoint file missing but backup exists. Restoring from backup...")
+            try:
+                self.checkpoint_backup.rename(self.checkpoint_file)
+                logger.info("Checkpoint restored from backup")
+            except OSError as e:
+                logger.error("Failed to restore checkpoint from backup: %s", e)
+                return None
+
+        if not self.checkpoint_file.exists():
+            logger.debug("No checkpoint file found")
+            return None
+
+        try:
+            with self.checkpoint_file.open("r") as f:
+                checkpoint = json.load(f, cls=Decoder)
+            logger.debug(
+                "Loaded checkpoint: last_batch=%s, completed=%d batches",
+                checkpoint.get("last_completed_batch_index"),
+                len(checkpoint.get("completed_batch_indices", [])),
+            )
+            return checkpoint
+        except (OSError, json.JSONDecodeError) as e:
+            logger.error("Failed to load checkpoint: %s", e)
+            return None
+
+    def save_checkpoint(
+        self,
+        completed_batch_indices: list[int],
+        last_simulator_name: str,
+        last_completed_batch_index: int,
+        last_simulator_state: dict[str, Any],
+    ) -> None:
+        """Save checkpoint after completing a batch.
+
+        Args:
+            completed_batch_indices: List of all completed batch indices so far
+            last_simulator_name: Name of the simulator that completed the batch
+            last_completed_batch_index: Index of the batch that just completed
+            last_simulator_state: State dict of the simulator after completion
+
+        Raises:
+            OSError: If checkpoint cannot be written
+        """
+        checkpoint = {
+            "completed_batch_indices": completed_batch_indices,
+            "last_simulator_name": last_simulator_name,
+            "last_completed_batch_index": last_completed_batch_index,
+            "last_simulator_state": last_simulator_state,
+        }
+
+        # Write to temp file first (atomic write pattern)
+        try:
+            with self.checkpoint_tmp.open("w") as f:
+                json.dump(checkpoint, f, indent=2, cls=Encoder)
+
+            # Backup existing checkpoint if it exists
+            if self.checkpoint_file.exists():
+                try:
+                    # Remove old backup if it exists (to avoid rename conflicts)
+                    if self.checkpoint_backup.exists():
+                        self.checkpoint_backup.unlink()
+                    self.checkpoint_file.rename(self.checkpoint_backup)
+                except OSError as e:
+                    logger.warning("Failed to backup previous checkpoint: %s", e)
+
+            # Move temp to final checkpoint
+            self.checkpoint_tmp.rename(self.checkpoint_file)
+
+            logger.debug(
+                "Checkpoint saved: batch_index=%d, completed=%d batches",
+                last_completed_batch_index,
+                len(completed_batch_indices),
+            )
+        except OSError as e:
+            logger.error("Failed to save checkpoint: %s", e)
+            # Clean up temp file if it exists
+            if self.checkpoint_tmp.exists():
+                try:
+                    self.checkpoint_tmp.unlink()
+                except OSError:
+                    pass
+            raise
+
+    def cleanup(self) -> None:
+        """Clean up checkpoint files after successful completion."""
+        # Remove both checkpoint and backup after successful completion
+        try:
+            if self.checkpoint_file.exists():
+                self.checkpoint_file.unlink()
+                logger.debug("Cleaned up checkpoint file")
+            if self.checkpoint_backup.exists():
+                self.checkpoint_backup.unlink()
+                logger.debug("Cleaned up checkpoint backup file")
+        except OSError as e:
+            logger.warning("Failed to clean up checkpoint files: %s", e)
+
+    def get_completed_batch_indices(self) -> set[int]:
+        """Get set of completed batch indices from checkpoint.
+
+        Returns:
+            Set of batch indices that have already been completed
+        """
+        checkpoint = self.load_checkpoint()
+        if checkpoint is None:
+            return set()
+        return set(checkpoint.get("completed_batch_indices", []))
+
+    def should_skip_batch(self, batch_index: int) -> bool:
+        """Check if a batch has already been completed.
+
+        Args:
+            batch_index: Index of batch to check
+
+        Returns:
+            True if batch was already completed, False otherwise
+        """
+        return batch_index in self.get_completed_batch_indices()

gwsim/cli/utils/config.py

@@ -0,0 +1,347 @@
+"""
+Utility functions to load and save configuration files.
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Any
+
+import yaml
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+logger = logging.getLogger("gwsim")
+
+
+class SimulatorOutputConfig(BaseModel):
+    """Configuration for simulator output handling."""
+
+    file_name: str = Field(..., description="Output file name template (supports {{ variable }} placeholders)")
+    arguments: dict[str, Any] = Field(
+        default_factory=dict, description="Output-specific arguments (e.g., channel name)"
+    )
+    output_directory: str | None = Field(
+        default=None, description="Optional directory override for this simulator's output"
+    )
+    metadata_directory: str | None = Field(
+        default=None, description="Optional directory override for this simulator's metadata"
+    )
+
+    # Allow unknown fields
+    model_config = ConfigDict(extra="allow")
+
+
+class SimulatorConfig(BaseModel):
+    """Configuration for a single simulator."""
+
+    class_: str = Field(alias="class", description="Simulator class name or full import path")
+    arguments: dict[str, Any] = Field(default_factory=dict, description="Arguments passed to simulator constructor")
+    output: SimulatorOutputConfig = Field(
+        default_factory=lambda: SimulatorOutputConfig(file_name="output-{{counter}}.hdf5"),
+        description="Output configuration for this simulator",
+    )
+
+    model_config = ConfigDict(extra="allow", populate_by_name=True)
+
+    @field_validator("class_", mode="before")
+    @classmethod
+    def validate_class_name(cls, v: str) -> str:
+        """Validate class specification is non-empty."""
+        if not isinstance(v, str) or not v.strip():
+            raise ValueError("'class' must be a non-empty string")
+        return v
+
+
+class GlobalsConfig(BaseModel):
+    """Global configuration applying to all simulators.
+
+    This configuration provides universal directory settings and fallback arguments
+    for simulators and output handlers. The simulator_arguments and output_arguments
+    are agnostic to simulator type, supporting both time-series and population simulators.
+    """
+
+    working_directory: str = Field(
+        default=".", alias="working-directory", description="Base working directory for all output"
+    )
+    output_directory: str | None = Field(
+        default=None, alias="output-directory", description="Default output directory (can be overridden per simulator)"
+    )
+    metadata_directory: str | None = Field(
+        default=None,
+        alias="metadata-directory",
+        description="Default metadata directory (can be overridden per simulator)",
+    )
+    simulator_arguments: dict[str, Any] = Field(
+        default_factory=dict,
+        alias="simulator-arguments",
+        description="Global default arguments for simulators (e.g., sampling-frequency, duration, seed). "
+        "Simulator-specific arguments override these.",
+    )
+    output_arguments: dict[str, Any] = Field(
+        default_factory=dict,
+        alias="output-arguments",
+        description="Global default arguments for output handlers (e.g., channel names). "
+        "Simulator-specific output arguments override these.",
+    )
+
+    model_config = ConfigDict(extra="allow", populate_by_name=True)
+
+
+class Config(BaseModel):
+    """Top-level configuration model."""
+
+    globals: GlobalsConfig = Field(default_factory=GlobalsConfig, description="Global configuration")
+    simulators: dict[str, SimulatorConfig] = Field(..., description="Dictionary of simulators")
+
+    model_config = ConfigDict(extra="allow", populate_by_name=True)
+
+    @field_validator("simulators", mode="before")
+    @classmethod
+    def validate_simulators_not_empty(cls, v: dict[str, Any]) -> dict[str, Any]:
+        """Ensure simulators section is not empty."""
+        if not v:
+            raise ValueError("'simulators' section cannot be empty")
+        return v
+
+
+def load_config(file_name: Path, encoding: str = "utf-8") -> Config:
+    """Load configuration file with validation.
+
+    Args:
+        file_name (Path): File name.
+        encoding (str, optional): File encoding. Defaults to "utf-8".
+
+    Returns:
+        Config: Validated configuration dataclass.
+
+    Raises:
+        FileNotFoundError: If the configuration file does not exist.
+        ValueError: If the configuration is invalid or cannot be parsed.
+    """
+    if not file_name.exists():
+        raise FileNotFoundError(f"Configuration file not found: {file_name}")
+    try:
+        with file_name.open(encoding=encoding) as f:
+            raw_config = yaml.safe_load(f)
+    except yaml.YAMLError as e:
+        raise ValueError(f"Failed to parse YAML configuration: {e}") from e
+
+    if not isinstance(raw_config, dict):
+        raise ValueError("Configuration must be a YAML dictionary")
+
+    # Validate and convert to Config dataclass
+    try:
+        config = Config(**raw_config)
+        logger.info("Configuration loaded and validated: %s simulators", len(config.simulators))
+        return config
+    except ValueError as e:
+        raise ValueError(f"Configuration validation failed: {e}") from e
+
+
+def save_config(
+    file_name: Path, config: Config, overwrite: bool = False, encoding: str = "utf-8", backup: bool = True
+) -> None:
+    """Save configuration to YAML file safely.
+
+    Args:
+        file_name: Path to save configuration to
+        config: Config dataclass instance
+        overwrite: If True, overwrite existing file
+        encoding: File encoding (default: utf-8)
+        backup: If True and overwriting, create backup
+
+    Raises:
+        FileExistsError: If file exists and overwrite=False
+    """
+    if file_name.exists() and not overwrite:
+        raise FileExistsError(f"File already exists: {file_name}. Use overwrite=True to overwrite.")
+
+    # Create backup if needed
+    if file_name.exists() and overwrite and backup:
+        backup_path = file_name.with_suffix(f"{file_name.suffix}.backup")
+        logger.info("Creating backup: %s", backup_path)
+        backup_path.write_text(file_name.read_text(encoding=encoding), encoding=encoding)
+
+    # Atomic write
+    temp_file = file_name.with_suffix(f"{file_name.suffix}.tmp")
+    try:
+        # Convert to dict, excluding internal fields
+        config_dict = config.model_dump(by_alias=True, exclude_none=False)
+
+        with temp_file.open("w", encoding=encoding) as f:
+            yaml.safe_dump(config_dict, f, default_flow_style=False, sort_keys=False)
+
+        temp_file.replace(file_name)
+        logger.info("Configuration saved to: %s", file_name)
+
+    except Exception as e:
+        if temp_file.exists():
+            temp_file.unlink()
+        raise ValueError(f"Failed to save configuration: {e}") from e
+
+
+def validate_config(config: dict) -> None:
+    """Validate configuration structure and provide helpful error messages.
+
+    Args:
+        config (dict): Configuration dictionary to validate
+
+    Raises:
+        ValueError: If configuration is invalid with detailed error message
+    """
+    # Check for required top-level structure
+    if "simulators" not in config:
+        raise ValueError("Invalid configuration: Must contain 'simulators' section with simulator definitions")
+
+    simulators = config["simulators"]
+
+    if not isinstance(simulators, dict):
+        raise ValueError("'simulators' must be a dictionary")
+
+    if not simulators:
+        raise ValueError("'simulators' section cannot be empty")
+
+    for name, sim_config in simulators.items():
+        if not isinstance(sim_config, dict):
+            raise ValueError(f"Simulator '{name}' configuration must be a dictionary")
+
+        # Check required fields
+        if "class" not in sim_config:
+            raise ValueError(f"Simulator '{name}' missing required 'class' field")
+
+        # Validate class specification
+        class_spec = sim_config["class"]
+        if not isinstance(class_spec, str) or not class_spec.strip():
+            raise ValueError(f"Simulator '{name}' 'class' must be a non-empty string")
+
+        # Validate arguments if present
+        if "arguments" in sim_config and not isinstance(sim_config["arguments"], dict):
+            raise ValueError(f"Simulator '{name}' 'arguments' must be a dictionary")
+
+        # Validate output configuration if present
+        if "output" in sim_config:
+            output_config = sim_config["output"]
+            if not isinstance(output_config, dict):
+                raise ValueError(f"Simulator '{name}' 'output' must be a dictionary")
+
+    # Validate globals section if present
+    if "globals" in config:
+        globals_config = config["globals"]
+        if not isinstance(globals_config, dict):
+            raise ValueError("'globals' must be a dictionary")
+
+    logger.info("Configuration validation passed")
+
+
+def resolve_class_path(class_spec: str, section_name: str | None) -> str:
+    """Resolve class specification to full module path.
+
+    Args:
+        class_spec: Either 'ClassName' or 'third_party.module.ClassName'
+        section_name: Section name (e.g., 'noise', 'signal', 'glitch')
+
+    Returns:
+        Full path like 'gwsim.noise.ClassName' or 'third_party.module.ClassName'
+
+    Examples:
+        resolve_class_path("WhiteNoise", "noise") -> "gwsim.noise.WhiteNoise"
+        resolve_class_path("numpy.random.Generator", "noise") -> "numpy.random.Generator"
+    """
+    if "." not in class_spec and section_name:
+        # Just a class name - use section_name as submodule, class imported in __init__.py
+        return f"gwsim.{section_name}.{class_spec}"
+    # Contains dots - assume it's a third-party package, use as-is
+    return class_spec
+
+
+def merge_parameters(globals_config: GlobalsConfig, simulator_args: dict[str, Any]) -> dict[str, Any]:
+    """Merge global and simulator-specific parameters.
+
+    Flattens simulator_arguments from globals into the result, then applies
+    simulator-specific overrides.
+
+    Args:
+        globals_config: GlobalsConfig dataclass instance
+        simulator_args: Simulator-specific arguments dict
+
+    Returns:
+        Merged parameters with simulator args taking precedence
+
+    Note:
+        Simulator_arguments from globals_config are flattened into the result.
+        Directory settings (working-directory, output-directory, metadata-directory)
+        are included. Output_arguments are not included (handled separately).
+    """
+    # Start with directory settings from globals
+    merged = {}
+    if globals_config.working_directory:
+        merged["working-directory"] = globals_config.working_directory
+    if globals_config.output_directory:
+        merged["output-directory"] = globals_config.output_directory
+    if globals_config.metadata_directory:
+        merged["metadata-directory"] = globals_config.metadata_directory
+
+    # Flatten simulator_arguments from globals
+    merged.update(globals_config.simulator_arguments)
+
+    # Override with simulator-specific arguments (takes precedence)
+    merged.update(simulator_args)
+
+    return merged
+
+
+def get_output_directories(
+    globals_config: GlobalsConfig,
+    simulator_config: SimulatorConfig,
+    simulator_name: str,
+    working_directory: Path | None = None,
+) -> tuple[Path, Path]:
+    """Get output and metadata directories for a simulator.
+
+    Args:
+        globals_config: Global configuration
+        simulator_config: Simulator-specific configuration
+        simulator_name: Name of the simulator
+        working_directory: Override working directory (for testing)
+
+    Returns:
+        Tuple of (output_directory, metadata_directory)
+
+    Priority (highest to lowest):
+        1. Simulator output.output_directory / output.metadata_directory
+        2. Global output-directory / metadata-directory
+        3. working-directory / output / {simulator_name}
+
+    Examples:
+        >>> globals_cfg = GlobalsConfig(working_directory="/data")
+        >>> sim_cfg = SimulatorConfig(class_="Noise")
+        >>> get_output_directories(globals_cfg, sim_cfg, "noise")
+        (Path("/data/output/noise"), Path("/data/output/noise"))
+    """
+    working_dir = working_directory or Path(globals_config.working_directory)
+
+    # Simulator-specific overrides
+    if simulator_config.output.output_directory:
+        output_path = Path(simulator_config.output.output_directory)
+        # Prepend working_dir if path is relative
+        output_directory = output_path if output_path.is_absolute() else working_dir / output_path
+    elif globals_config.output_directory:
+        output_path = Path(globals_config.output_directory)
+        # Prepend working_dir if path is relative
+        output_directory = output_path if output_path.is_absolute() else working_dir / output_path
+    else:
+        output_directory = working_dir / "output" / simulator_name
+
+    if simulator_config.output.metadata_directory:
+        metadata_path = Path(simulator_config.output.metadata_directory)
+        # Prepend working_dir if path is relative
+        metadata_directory = metadata_path if metadata_path.is_absolute() else working_dir / metadata_path
+    elif globals_config.metadata_directory:
+        metadata_path = Path(globals_config.metadata_directory)
+        # Prepend working_dir if path is relative
+        metadata_directory = metadata_path if metadata_path.is_absolute() else working_dir / metadata_path
+    else:
+        metadata_directory = working_dir / "metadata" / simulator_name
+
+    return output_directory, metadata_directory

gwsim/cli/utils/hash.py

@@ -0,0 +1,23 @@
+"""Contains utility functions for hashing operations."""
+
+from __future__ import annotations
+
+import hashlib
+from pathlib import Path
+
+
+def compute_file_hash(file_path: str | Path, algorithm: str = "sha256") -> str:
+    """Compute the hash of a file using the specified algorithm.
+
+    Args:
+        file_path: Path to the file.
+        algorithm: Hashing algorithm to use (default is 'sha256').
+
+    Returns:
+        The computed hash as a hexadecimal string.
+    """
+    hash_func = hashlib.new(algorithm)
+    with open(file_path, "rb") as f:
+        for chunk in iter(lambda: f.read(8192), b""):
+            hash_func.update(chunk)
+    return f"{algorithm}:{hash_func.hexdigest()}"

gwsim/cli/utils/retry.py

@@ -0,0 +1,62 @@
+"""Utility functions for retrying operations with exponential backoff."""
+
+from __future__ import annotations
+
+import logging
+import time
+from typing import Any
+
+logger = logging.getLogger("gwsim")
+
+
+class RetryManager:
+    """Manages retry logic with exponential backoff."""
+
+    def __init__(
+        self,
+        max_retries: int = 3,
+        base_delay: float = 1.0,
+        retryable_exceptions: tuple[type[Exception], ...] = (
+            OSError,
+            PermissionError,
+            FileNotFoundError,
+            RuntimeError,
+            ValueError,
+        ),
+    ):
+        """Initialize the RetryManager.
+
+        Args:
+            max_retries: Maximum number of retries.
+            base_delay: Base delay in seconds for exponential backoff.
+            retryable_exceptions: Tuple of exception types that are considered retryable.
+        """
+        self.max_retries = max_retries
+        self.base_delay = base_delay
+        self.retryable_exception = retryable_exceptions
+
+    def retry_with_backoff(self, operation, *args, **kwargs) -> Any | None:
+        """Retry operation with exponential backoff."""
+        for attempt in range(self.max_retries + 1):
+            try:
+                return operation(*args, **kwargs)
+            except self.retryable_exception as e:
+                if attempt == self.max_retries:
+                    logger.error("Operation failed after %s retries: %s", self.max_retries, e)
+                    raise
+
+                delay = self.base_delay * (2**attempt)
+                logger.warning("Attempt %s failed: %s. Retrying in %ss...", attempt + 1, e, delay)
+                time.sleep(delay)
+        return None
+
+    def is_retryable_exception(self, exception: Exception) -> bool:
+        """Check if an exception is retryable.
+
+        Args:
+            exception: The exception to check
+
+        Returns:
+            bool: True if the exception is retryable, False otherwise
+        """
+        return isinstance(exception, self.retryable_exception)