simstack 2.0.1__py3-none-any.whl

simstack/util/path_manager.py

@@ -0,0 +1,223 @@
+from pathlib import Path
+from typing import List, Dict, Optional, Iterator, Any, Union
+
+from simstack.util.directory_iterator import DirectoryPath
+from simstack.util.project_root_finder import find_project_root
+
+
+class PathManager:
+    """
+    Manages paths for the SimStack application, providing mechanisms to find Python files
+    for nodes and models. Will read only .py files.
+    By default travers all directories below the project root.
+    Uses DirectoryPath for efficient directory traversal.
+    """
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super(PathManager, cls).__new__(cls)
+            cls._instance._initialized = False
+        return cls._instance
+
+    def default_excludes(self):
+        return [
+            "__pycache__",
+            "*.pyc",
+            ".git",
+            ".venv",
+            "venv",
+        ]
+
+    def reset(self):
+        self._initialized = False
+        self.paths.clear()
+        self._excluded_patterns = self.default_excludes()
+
+    def __init__(self, use_pickle: bool = False, include_project_root: bool = False):
+        """
+        Initialize the PathManager.
+
+        Args:
+            use_pickle: Boolean flag indicating whether to use pickle for serialization
+        """
+        if self._initialized:
+            return
+
+        self._initialized = True
+        self.use_pickle = use_pickle
+        self.root_dir = find_project_root()
+        self.paths: Dict[str, Dict[str, Union[Path, str, bool]]] = {}
+
+        if include_project_root:
+            self.add_path("project_root", self.root_dir)
+
+        self._excluded_patterns = self.default_excludes()
+
+    def add_path(self, name: str, path: Path, drops: str = "", use_pickle: bool = False) -> None:
+        """
+        Add a path to the manager.
+
+        Args:
+            name: Name identifier for the path
+            path: The directory path relative to the project root
+            drops: Prefix to drop from module names (for import paths)
+            use_pickle: Whether to use pickle for this path
+        """
+        # Convert relative paths to absolute paths
+        if not path.is_absolute():
+            path = self.root_dir / path
+
+        if not path.is_dir():
+            raise ValueError(f"'{path}' is not a valid directory")
+
+        self.paths[name] = {"path": path, "drops": drops, "use_pickle": use_pickle}
+
+    def get_path(self, name: str) -> Dict[str, str]:
+        """
+        Get a path by name.
+
+        Args:
+            name: Name of the path to retrieve
+
+        Returns:
+            Dictionary containing path information
+        """
+        if name not in self.paths:
+            raise KeyError(f"Path '{name}' not found in PathManager")
+
+        return self.paths[name]
+
+    def find_python_files(self, path_name: str, excluded_patterns: Optional[List[str]] = None) -> List[str]:
+        """
+        Find Python files in the specified path, excluding __init__.py files.
+
+        Args:
+            path_name: Name of the path to search in
+            excluded_patterns: Additional patterns to exclude
+
+        Returns:
+            List of absolute file paths to Python files
+        """
+        path_info = self.get_path(path_name)
+        path = path_info["path"]
+
+        # Combine default and additional exclusion patterns
+        all_excluded_patterns = self._excluded_patterns.copy()
+        if excluded_patterns:
+            all_excluded_patterns.extend(excluded_patterns)
+
+        # Add __init__.py to excluded patterns
+        all_excluded_patterns.append("__init__.py")
+
+        # Use DirectoryPath to find Python files
+        dir_path = DirectoryPath(path, all_excluded_patterns, [".py"])
+
+        # Convert Path objects to strings for compatibility with existing code
+        return [str(file_path) for file_path in dir_path.get_files_list()]
+
+    def iterate_python_files(self, path_name: str, excluded_patterns: Optional[List[str]] = None) -> Iterator[Path]:
+        """
+        Iterate over Python files in the specified path, excluding __init__.py files.
+
+        Args:
+            path_name: Name of the path to search in
+            excluded_patterns: Additional patterns to exclude
+
+        Returns:
+            Iterator of Path objects for Python files
+        """
+        path_info = self.get_path(path_name)
+        path = path_info["path"]
+
+        # Combine default and additional exclusion patterns
+        all_excluded_patterns = self._excluded_patterns.copy()
+        if excluded_patterns:
+            all_excluded_patterns.extend(excluded_patterns)
+
+        # Add __init__.py to excluded patterns
+        all_excluded_patterns.append("__init__.py")
+
+        # Use DirectoryPath to iterate over Python files
+        dir_path = DirectoryPath(path, all_excluded_patterns, [".py"])
+        return dir_path.iterate_files()
+
+    def get_drops(self, path_name: str) -> str:
+        """
+        Get the drops value for a path.
+
+        Args:
+            path_name: Name of the path
+
+        Returns:
+            The drops value for the path
+        """
+        path_info = self.get_path(path_name)
+        return path_info["drops"]
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> "PathManager":
+        """
+        Create a PathManager from configuration.
+
+        Args:
+            config: Configuration object containing path information
+
+        Returns:
+            Initialized PathManager instance
+        """
+        # Get use_pickle from config
+
+        use_pickle = config.get("parameters", {}).get("general", {}).get("use_pickle", False)
+
+        path_manager = cls(use_pickle=use_pickle)
+
+        config_paths = config.get("paths", {})
+        for path_name, path_info in config_paths.items():
+            path_manager.add_path(path_name, path_info["path"], path_info.get("drops", ""))
+
+        return path_manager
+
+    def find_parent_path(self, path: str) -> Optional[str]:
+        """
+        Find the parent path from self.paths that contains the given path.
+
+        Args:
+            path: The path to find the parent of
+
+        Returns:
+            The name of the parent path from self.paths that contains the given path,
+            or None if no parent is found
+        """
+        path_obj = Path(path)
+
+        # Convert to the absolute path if needed
+        if not path_obj.is_absolute():
+            path_obj = self.root_dir / path_obj
+
+        # Search through all paths in self.paths to find the one that contains path_obj
+        best_match = None
+        best_match_depth = -1
+
+        for path_name, path_info in self.paths.items():
+            registered_path = Path(path_info["path"])
+
+            try:
+                # Check if path_obj is under this registered path
+                path_obj.relative_to(registered_path)
+
+                # Calculate depth (number of parent directories)
+                depth = len(registered_path.parts)
+
+                # Keep the deepest match (the most specific parent)
+                if depth > best_match_depth:
+                    best_match = path_name
+                    best_match_depth = depth
+
+            except ValueError:
+                # path_obj is not under this registered path, continue
+                continue
+
+        return best_match
+
+path_manager = PathManager()

simstack/util/project_root_finder.py

@@ -0,0 +1,41 @@
+import os
+from pathlib import Path
+
+
+def find_project_root(current_file=None, marker_files=(".git", "simstack.toml", "setup.py"),
+                      skip_files=("simstack_marker.com",)) -> Path:
+    """
+    Find the project root directory by searching for common marker files
+
+    Args:
+        skip_files: directories with these names will be skipped
+        current_file: Path to the current file, defaults to __file__ if None
+        marker_files: Tuple of files/directories that indicate the project root
+
+    Returns:
+        Absolute path to the project root directory
+    """
+    if "SIMSTACK_PROJECT_ROOT" in os.environ:
+        return Path(os.environ["SIMSTACK_PROJECT_ROOT"]).resolve()
+
+    if current_file is None:
+        current_file = __file__
+
+    # Get the directory from where the process was started
+    current_dir = Path.cwd()
+
+    # Walk up the directory tree until we find a marker file
+    prev_dir = None
+    while current_dir != prev_dir:
+        # Check if any marker files/directories exist in the current directory
+        found_marker = any((current_dir / marker).exists() for marker in marker_files)
+        found_skip = any((current_dir / marker).exists() for marker in skip_files)
+        if found_marker and not found_skip:
+            return current_dir
+
+        # Move up one directory
+        prev_dir = current_dir
+        current_dir = current_dir.parent
+
+    # If we can't find any markers, return the directory of the current file
+    return Path(current_file).resolve().parent

simstack/util/queue_watcher.py

@@ -0,0 +1,262 @@
+#!/usr/bin/env python3
+import argparse
+import os
+import signal
+import subprocess
+import sys
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+from typing import Dict, Optional, Tuple
+
+
+def parse_spec(
+    cmd_file: Path
+) -> Tuple[Optional[str], Dict[str, str], Optional[str], Optional[int]]:
+    cwd: Optional[str] = None
+    env: Dict[str, str] = {}
+    command: Optional[str] = None
+    timeout: Optional[int] = None
+
+    with cmd_file.open("r", encoding="utf-8") as f:
+        for raw in f:
+            line = raw.strip()
+            if not line:
+                continue
+            if line.startswith("CWD="):
+                cwd = line[len("CWD=") :]
+            elif line.startswith("ENV[") and "]=" in line:
+                key = line[len("ENV[") : line.index("]=", len("ENV["))]
+                val = line.split("]=", 1)[1]
+                env[key] = val
+            elif line.startswith("CMD="):
+                command = line[len("CMD=") :]
+            elif line.startswith("TIMEOUT="):
+                try:
+                    timeout = int(line[len("TIMEOUT=") :])
+                except ValueError:
+                    timeout = None
+    return cwd, env, command, timeout
+
+
+def run_command(
+    command: str,
+    cwd: Optional[str],
+    env: Dict[str, str],
+    timeout: Optional[int],
+    out_file: Path,
+    err_file: Path,
+) -> int:
+    proc_env = os.environ.copy()
+    proc_env.update(env)
+    with out_file.open("w", encoding="utf-8") as out, err_file.open(
+        "w", encoding="utf-8"
+    ) as err:
+        try:
+            proc = subprocess.run(
+                command,
+                shell=True,
+                cwd=cwd if cwd else None,
+                env=proc_env,
+                stdout=out,
+                stderr=err,
+                text=True,
+                timeout=timeout if (timeout and timeout > 0) else None,
+                check=False,
+            )
+            return proc.returncode
+        except subprocess.TimeoutExpired:
+            err.write("\n[watcher] Command timed out\n")
+            return 124
+        except Exception as e:
+            err.write(f"\n[watcher] Execution error: {e}\n")
+            return 1
+
+
+def process_job_from_claim(base_stem: Path, claimed_by_lock: bool) -> None:
+    """
+    Process a job using base path (without suffix). If claimed_by_lock is True, we did not
+    rename .ready; we will remove it at the end. Otherwise, we operate on .claim and remove it.
+    """
+    cmd_file = base_stem.with_suffix(".cmd")
+    ready_file = base_stem.with_suffix(".ready")
+    claim_file = base_stem.with_suffix(".claim")
+    lock_file = base_stem.with_suffix(".lock")
+    out_file = base_stem.with_suffix(".stdout")
+    err_file = base_stem.with_suffix(".stderr")
+    exit_file = base_stem.with_suffix(".exit")
+    done_file = base_stem.with_suffix(".done")
+    fail_file = base_stem.with_suffix(".fail")
+
+    try:
+        if not cmd_file.exists():
+            fail_file.touch()
+            return
+
+        cwd, env, command, timeout = parse_spec(cmd_file)
+        if not command:
+            err_file.write_text("[watcher] CMD missing in spec\n", encoding="utf-8")
+            exit_file.write_text("2\n", encoding="utf-8")
+            fail_file.touch()
+            return
+
+        rc = run_command(command, cwd, env, timeout, out_file, err_file)
+        print(f"[watcher] Job {base_stem.name} finished with rc={rc}")
+        exit_file.write_text(f"{rc}\n", encoding="utf-8")
+        print(f"[watcher] Job {base_stem.name} finished")
+        (done_file if rc == 0 else fail_file).touch()
+    finally:
+        # Cleanup markers depending on the claim type
+        if claimed_by_lock:
+            try:
+                if ready_file.exists():
+                    ready_file.unlink()
+            except Exception as e:
+                print(
+                    f"[watcher] Warning: failed to remove ready {ready_file}: {e}",
+                    file=sys.stderr,
+                )
+            try:
+                if lock_file.exists():
+                    lock_file.unlink()
+            except Exception as e:
+                print(
+                    f"[watcher] Warning: failed to remove lock {lock_file}: {e}",
+                    file=sys.stderr,
+                )
+        else:
+            try:
+                if claim_file.exists():
+                    claim_file.unlink()
+            except Exception as e:
+                print(
+                    f"[watcher] Warning: failed to remove claim {claim_file}: {e}",
+                    file=sys.stderr,
+                )
+
+
+def try_claim(
+    ready_file: Path, retries: int = 5, backoff: float = 0.05
+) -> Optional[Path]:
+    """
+    Atomically rename *.ready -> *.claim. Retries on transient failures.
+    Returns claim path if successful, else None.
+    """
+    claim = ready_file.with_suffix(".claim")
+    for i in range(max(1, retries)):
+        try:
+            os.replace(ready_file, claim)
+            return claim
+        except FileNotFoundError:
+            return None
+        except PermissionError:
+            # Windows/NFS transient locks: backoff and retry
+            time.sleep(backoff * (i + 1))
+        except OSError:
+            # Other transient FS issues: backoff and retry
+            time.sleep(backoff * (i + 1))
+    return None
+
+
+def try_lock(base_stem: Path) -> bool:
+    """
+    Create an exclusive lock file <id>.lock using O_CREAT|O_EXCL.
+    Returns True if lock acquired, False otherwise.
+    """
+    lock_path = base_stem.with_suffix(".lock")
+    flags = os.O_CREAT | os.O_EXCL | os.O_WRONLY
+    try:
+        fd = os.open(lock_path, flags, 0o644)
+        os.close(fd)
+        return True
+    except FileExistsError:
+        return False
+    except OSError as e:
+        print(f"[watcher] Lock create failed for {lock_path}: {e}", file=sys.stderr)
+        return False
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description="Concurrent file-queue watcher with robust claiming"
+    )
+    parser.add_argument(
+        "--queue-dir", type=Path, required=False, default=Path.cwd() / "queue"
+    )
+    parser.add_argument("--poll-interval", type=float, default=1.0)
+    parser.add_argument("--workers", type=int, default=min(4, (os.cpu_count() or 2)))
+    parser.add_argument(
+        "--use-srun",
+        action="store_true",
+        help="(deprecated; wrapper removed for simplicity)",
+    )
+    parser.add_argument("--srun-args", type=str, default="")
+    parser.add_argument(
+        "--once", action="store_true", help="Process current claims and exit"
+    )
+    args = parser.parse_args()
+
+    qd = args.queue_dir
+    qd.mkdir(parents=True, exist_ok=True)
+    print(f"[watcher] Watching {qd}, workers={args.workers}")
+
+    stop = {"flag": False}
+
+    def handle_sig(_s, _f):
+        stop["flag"] = True
+
+    signal.signal(signal.SIGINT, handle_sig)
+    signal.signal(signal.SIGTERM, handle_sig)
+
+    with ThreadPoolExecutor(max_workers=max(1, args.workers)) as pool:
+        futures = set()
+
+        while not stop["flag"]:
+            claimed_any = False
+
+            # Scan for all .ready files
+            ready_list = list(qd.glob("*.ready"))
+            if ready_list:
+                print(f"[watcher] Found {len(ready_list)} ready file(s)")
+            for ready in ready_list:
+                base_stem = ready.with_suffix("")  # drop .ready
+                # 1) Preferred: atomic rename to .claim (with retries)
+                claim = try_claim(ready)
+                if claim is not None:
+                    claimed_any = True
+                    futures.add(pool.submit(process_job_from_claim, base_stem, False))
+                    continue
+
+                # 2) Fallback: exclusive lock file if rename failed
+                if try_lock(base_stem):
+                    print(f"[watcher] Claimed by lock: {base_stem.name}")
+                    claimed_any = True
+                    futures.add(pool.submit(process_job_from_claim, base_stem, True))
+
+            # Reap done tasks
+            done_now = {f for f in futures if f.done()}
+            for f in done_now:
+                try:
+                    f.result()
+                except Exception as e:
+                    print(f"[watcher] Job raised: {e}", file=sys.stderr)
+                futures.discard(f)
+
+            if args.once:
+                # Drain: wait for all running tasks to finish
+                for f in as_completed(list(futures)):
+                    try:
+                        f.result()
+                    except Exception as e:
+                        print(f"[watcher] Job raised: {e}", file=sys.stderr)
+                break
+
+            if not claimed_any:
+                time.sleep(max(0.05, args.poll_interval))
+
+    print("[watcher] Exiting")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())