simstack 0.2.0__py3-none-any.whl

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())

simstack/util/route_finder.py

@@ -0,0 +1,74 @@
+# TODO rmeove duplicate route finder this is a test implementation
+routes_table = [
+    {"target": "local", "source": "int-nano", "host": "local"},
+    {"target": "int-nano", "source": "local", "host": "local"},
+    {
+        "target": "horeka",
+        "source": "local",
+        "host": "horeka",
+    },
+    {
+        "target": "local",
+        "source": "horeka",
+        "host": "horeka",
+    },
+]
+
+
+def find_route(target, source):
+    """
+    Find the minimal route from source to target using available routes.
+
+    Args:
+        target (str): The target system
+        source (str): The source system
+
+    Returns:
+        list: A list of systems to visit in order, or None if no route exists
+    """
+    # If source and target are the same, no transfer needed
+    if source == target:
+        return [source]
+
+    # Build a graph of available routes
+    graph = {}
+    for route in routes_table:
+        if route["source"] not in graph:
+            graph[route["source"]] = []
+        graph[route["source"]].append(route["target"])
+
+    # If source isn't in our graph, we can't start a route
+    if source not in graph:
+        return None
+
+    # Use BFS to find the shortest path
+    visited = {source}
+    queue = [(source, [source])]  # (current_node, path_so_far)
+
+    while queue:
+        current, path = queue.pop(0)
+
+        # Check neighbors (available transfer targets)
+        if current in graph:
+            for neighbor in graph[current]:
+                if neighbor == target:
+                    # Found a path to target
+                    return path + [neighbor]
+
+                if neighbor not in visited:
+                    visited.add(neighbor)
+                    queue.append((neighbor, path + [neighbor]))
+
+    # If we've exhausted all possibilities without finding a path
+    return None
+
+
+if __name__ == "__main__":
+    # Example usage
+    target = "int-nano"
+    source = "horeka"
+    route = find_route(target, source)
+    if route:
+        print(f"Route from {source} to {target}: {' -> '.join(route)}")
+    else:
+        print(f"No route found from {source} to {target}")

simstack/util/runner_util.py

@@ -0,0 +1,162 @@
+import asyncio
+import logging
+import os
+import re
+import socket
+import subprocess
+from pathlib import Path
+
+from simstack.core.context import context
+from simstack.models.runner_model import RunnerEvent, RunnerType, RunnerEventEnum
+
+logger = logging.getLogger("RunnerUtil")
+
+
+def ensure_crontab_entry(command, schedule="*/10 * * * *"):
+    """
+    Checks if a crontab entry exists and adds it if it doesn't.
+
+    Args:
+        command (str): The command to be executed by cron
+        schedule (str): The cron schedule expression (default: every 10 minutes)
+
+    Returns:
+        bool: True if entry was added, False if it already existed
+    """
+    # The full crontab entry to look for
+    cron_entry = f"{schedule} {command}"
+
+    try:
+        # Get current crontab
+        result = subprocess.run(["crontab", "-l"], capture_output=True, text=True)
+
+        # Check if command succeeded
+        if result.returncode == 0:
+            current_crontab = result.stdout
+        else:
+            # If crontab is empty or doesn't exist
+            current_crontab = ""
+
+        # Check if our entry already exists (using regex to handle whitespace variations)
+        entry_pattern = (
+            re.escape(schedule.strip()) + r"\s+" + re.escape(command.strip())
+        )
+        if re.search(entry_pattern, current_crontab):
+            logger.info(f"Crontab entry for '{command}' already exists.")
+
+            return False
+
+        # Add our entry to crontab
+        new_crontab = current_crontab
+        if new_crontab and not new_crontab.endswith("\n"):
+            new_crontab += "\n"
+        new_crontab += cron_entry + "\n"
+
+        # Write back to crontab
+        process = subprocess.Popen(["crontab", "-"], stdin=subprocess.PIPE, text=True)
+        process.communicate(input=new_crontab)
+
+        if process.returncode == 0:
+            logger.info(f"Successfully added crontab entry: '{cron_entry}'")
+            return True
+        else:
+            logger.error(f"Failed to update crontab, return code: {process.returncode}")
+            return False
+
+    except Exception as e:
+        logger.error(f"Error managing crontab: {str(e)}")
+        return False
+
+
+async def graceful_shutdown(resource_name):
+    # Log shutdown
+
+    pid = os.getpid()
+    user = os.environ.get("USER", os.environ.get("USERNAME", "unknown"))
+    hostname = socket.gethostname()
+
+    # Capture current user info for the subprocess
+    current_user = os.environ.get("USER", os.environ.get("USERNAME"))
+
+    runner_event = RunnerEvent(
+        resource=resource_name,
+        runner_type=RunnerType.RESOURCE_RUNNER,
+        event=RunnerEventEnum.SHUTDOWN,
+        pid=pid,
+        user=user,
+        hostname=hostname,
+        message=f"Graceful shutdown of {pid} by {current_user}",
+    )
+    await context.db.save(runner_event)
+
+    # Cancel all remaining tasks
+    for task in asyncio.all_tasks():
+        if task is not asyncio.current_task():
+            task.cancel()
+
+    # Give tasks time to cancel
+    await asyncio.sleep(0.5)
+
+    # Signal to the main function to exit
+    raise SystemExit(0)
+
+
+async def restart(resource_name):
+    logger.info(
+        f"Restarting node runner for resource: {resource_name} {context.config.python_path}"
+    )
+
+    cmd = Path(context.config.python_path[0]) / "scripts" / "check_runner.sh"
+
+    # Get current user information
+    current_uid = os.getuid()
+    current_gid = os.getgid()
+    current_pid = os.getpid()
+    logger.info(f"Current user ID: {current_uid}, group ID: {current_gid}")
+
+    # Create environment with preserved user information
+    env = os.environ.copy()
+
+    # Start a new process with the same arguments
+
+    logger.info(f"Restarting node runner for resource: {resource_name} {str(cmd)}")
+
+    result = subprocess.Popen(str(cmd), start_new_session=True, env=env)
+    # Log the restart
+    message = f"user: {current_uid} group: {current_gid} pid: {current_pid} result: {result.pid}"
+    runner_event = RunnerEvent(
+        resource=resource_name,
+        runner_type=RunnerType.RESOURCE_RUNNER,
+        event=RunnerEventEnum.RESTART,
+        message=message,
+    )
+    await context.db.save(runner_event)
+    # Exit current process
+    logger.info(f"Restarted node runner for resource: {resource_name}")
+    # await graceful_shutdown(resource_name)
+
+
+async def schedule_restart(resource_name, restart_minutes):
+    """Schedule a restart after the specified minutes"""
+    try:
+        logger.info(
+            f"Runner will restart in {restart_minutes} minutes on resource {resource_name} "
+        )
+        await asyncio.sleep(restart_minutes * 60)  # Convert minutes to seconds
+        # try:
+        #     if context.config.python_path:
+        #         path = Path(context.config.python_path[0])
+        #         command = str(path / "scripts" / "check_runner.sh") + " >> " + str(path / "cron_log.log") + " 2>&1"
+        #         logger.info(f"Path: {command} ")
+        #         if not ensure_crontab_entry(command):
+        #             runner_entry = RunnerEvent(resource=resource_name, runner_type=RunnerType.RESOURCE_RUNNER,
+        #                                        event=RunnerEventEnum.CRONTAB_GONE,
+        #                                        message="Crontab entry for resource runner was not found")
+        #             await context.db.save(runner_entry)
+        # except Exception as e:
+        #     logger.error("Something went wrong for crontab")
+        await restart(resource_name)
+        await graceful_shutdown(resource_name)
+    except asyncio.CancelledError:
+        logger.info("Restart task was cancelled")
+        raise

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