simstack 2.0.1__py3-none-any.whl

simstack/util/resource_config.py

@@ -0,0 +1,198 @@
+import tomllib
+import os
+import shutil
+import subprocess
+import tempfile
+import uuid
+from pathlib import Path
+from typing import Dict, Any, Optional, List, Union
+
+from odmantic import ObjectId
+
+import logging
+logger = logging.getLogger("ResourceConfig")
+
+class ResourceConfig:
+    """
+    ResourceConfig is responsible for managing configuration settings, setup,
+    execution, and post-processing parameters for specified resources.
+
+    This class is designed to read configuration data from a TOML file, provide
+    an interface to access resource-specific parameters, and execute resource-related
+    operations such as setup and running commands. It encapsulates functionality for
+    handling temporary directories, file manipulation, and subprocess execution.
+
+    Attributes:
+        _config (Dict[str, Any]): The loaded configuration dictionary.
+        _resource (str): The name of the current resource.
+    """
+    def __init__(self, config_path: Path, resource: str):
+        self._config: Dict[str, Any] = {}
+        self._resource = resource
+
+        actual_path = Path(config_path)
+        if actual_path.is_dir():
+            actual_path = actual_path / "config.toml"
+
+        if actual_path.exists():
+            with open(actual_path, "rb") as f:
+                self._config = tomllib.load(f)
+        else:
+            # If it doesn't exist, we just have an empty config
+            pass
+
+    @property
+    def os(self) -> str:
+        """
+        Returns the OS of the current resource, defaults to 'linux'.
+        """
+        try:
+            return self._config[self._resource].get("os", "linux")
+        except KeyError:
+            return "linux"
+
+    def setup(self, node_runner: Optional[Any] = None):
+        """
+        Executes the setup scripts for the current resource.
+        """
+        setup_params = self.get_setup_params()
+        scripts = setup_params.get("scripts", [])
+        for i, script in enumerate(scripts):
+            # Note: shell=True is needed to execute shell commands/scripts
+            # and might be OS dependent for the actual commands.
+            if node_runner and hasattr(node_runner, "subprocess"):
+                node_runner.subprocess(f"setup_{i}", script)
+            else:
+                subprocess.run(script, shell=True, check=True)
+
+    def tmp_dir(self, task_id: ObjectId | str) -> Path:
+        if isinstance(task_id, ObjectId):
+            task_id = str(task_id)
+        tmp_dir = self.tmp_base_dir / str(task_id)
+        Path(tmp_dir).mkdir(parents=True, exist_ok=True)
+        return tmp_dir
+
+    @property
+    def tmp_base_dir(self) -> Path:
+        tmp_base_dir_str = self.get_setup_params().get("tmp_base_dir", "")
+
+        if tmp_base_dir_str:
+            expanded_path_str = os.path.expandvars(os.path.expanduser(tmp_base_dir_str))
+            path = Path(expanded_path_str)
+            path.mkdir(parents=True, exist_ok=True)
+            return path
+
+        # Default to system temp directory if no command/path specified
+        return Path(tempfile.gettempdir())
+
+    def run(self,
+        program_name: str,
+        input_files: Optional[List[Union[str, "FileStack"]]] = None,
+        output_files: Optional[List[Union[str, "FileStack"]]] = None,
+        node_runner: Optional[Any] = None,
+    ):
+        """
+        Executes the run command with optional temporary directory usage and file handling.
+        Retrieves parameters from the configuration for the specified program.
+
+        Args:
+            program_name: Name of the program to run.
+            input_files: List of input files (str or FileStack). Overrides TOML input_files if provided.
+            output_files: List of output files (str or FileStack). Overrides TOML output_files if provided.
+            node_runner: Optional NodeRunner instance for execution.
+        """
+
+        params = self.get_program(program_name)
+        run_command = params.get("run_command", "")
+        
+        if input_files is None:
+            input_files = params.get("input_files", [])
+        
+        if output_files is None:
+            output_files = params.get("output_files", [])
+            
+        use_temp = params.get("use_temp", False)
+        
+        # tmp_base_dir can come from setup or the program itself, but usually it's in setup for the resource
+    
+        # scratch_cleanup from postprocessing
+        post_params = self.get_postprocessing_params()
+        scratch_cleanup = params.get("scratch_cleanup", post_params.get("scratch_cleanup", False))
+
+        tmp_dir = None
+        try:
+            exec_dir = Path.cwd()
+            if use_temp:
+                tmp_id = node_runner.task_id if node_runner else uuid.uuid4()
+                tmp_dir = self.tmp_dir(tmp_id)
+                exec_dir = tmp_dir
+                
+                # Copy input files to exec_dir
+                for f in input_files:
+                    if hasattr(f, "get"):  # It's a FileStack
+                        f.get(local_dir=exec_dir)
+                    else:  # It's a string (filename)
+                        src = Path.cwd() / f
+                        if src.exists() and src != exec_dir / f:
+                            shutil.copy(src, exec_dir / f)
+
+            # Execute run_command
+            if node_runner and hasattr(node_runner, "subprocess"):
+                node_runner.subprocess("run", run_command, cwd=str(exec_dir))
+            else:
+                subprocess.run(run_command, shell=True, check=True, cwd=exec_dir)
+
+
+            # Copy output files back to cwd
+            if use_temp and tmp_dir:
+                for f in output_files:
+                    # output_files can also be FileStack in the new pattern?
+                    # "output_files: List[str | FileStack] = [...]"
+                    # If it's a FileStack, we might need to update its content from the local file
+                    # but the standard behavior for output_files here is copying back to cwd.
+                    filename = f.name if hasattr(f, "name") else f
+                    src = tmp_dir / filename
+                    if src.exists():
+                        shutil.copy(src, Path.cwd() / filename)
+
+        finally:
+            if scratch_cleanup and tmp_dir and tmp_dir.exists():
+                shutil.rmtree(tmp_dir)
+
+    def get_program(self, program_name: str) -> Dict[str, Any]:
+        """
+        Returns the dict from resource.program.name for program with name and the current resource.
+        Expected structure in TOML: [resource_name.program.program_name]
+        """
+        try:
+            return self._config[self._resource]["program"][program_name]
+        except (KeyError, TypeError):
+            return {}
+
+    def get_setup_params(self) -> Dict[str, Any]:
+        """
+        Returns the setup dict for the specified resource.
+        Expected structure in TOML: [resource_name.setup]
+        """
+        try:
+            return self._config[self._resource]["setup"]
+        except KeyError:
+            return {}
+
+    def get_postprocessing_params(self) -> Dict[str, Any]:
+        """
+        Returns the post-processing dict for the specified resource.
+        Expected structure in TOML: [resource_name.post-processing] or [resource_name.postprocessing]
+        """
+        try:
+            resource_cfg = self._config[self._resource]
+            if "post-processing" in resource_cfg:
+                return resource_cfg["post-processing"]
+            if "postprocessing" in resource_cfg:
+                return resource_cfg["postprocessing"]
+        except KeyError:
+            pass
+        return {}
+
+    def __str__(self):
+        return f"ResourceConfig(resource={self._resource}, config={self._config})"

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,166 @@
+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_paths}"
+    )
+
+    cmd = Path(context.config.project_root) / "scripts" / "old" / "check_runner.sh"
+    if not cmd.is_file():
+        raise FileNotFoundError(f"Runner restart script not found: {cmd}")
+
+    # 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(
+        ["bash", 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" / "old" / "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/runner_utils.py

@@ -0,0 +1,143 @@
+import re
+import subprocess
+import uuid
+from datetime import datetime
+from typing import List
+
+from odmantic import ObjectId
+from simstack.core.context import context
+from simstack.models.parameters import Resource
+from simstack.models.slurm_info import SlurmInfo
+from simstack.util.git_repository_status import get_git_status
+from simstack.util.submit_to_watchdog import submit_to_watchdog
+
+import logging
+
+logger = logging.getLogger("runner_utils")
+
+
+def make_git_status_list() -> List[str]:
+    git_status_list = []
+    git_path_list = [context.config.project_root]
+    for path in git_path_list:
+        result = get_git_status(path)
+        if result["branch"]:
+            value = result["branch"] + "[" + result["short_hash"] + "]"
+            if result["up_to_date"]:
+                value += " (up-to-date)"
+            else:
+                value += " (behind " + str(result["behind"]) + " commits)"
+            git_status_list.append(value)
+        else:
+            git_status_list.append("No branch found")
+    return git_status_list
+
+
+def run_squeue_for_job(job_id: str) -> str:
+    result = subprocess.run(
+        f"squeue -j {job_id}",
+        shell=True,
+        capture_output=True,
+        text=True,
+        timeout=30,
+    )
+    return result.stdout
+
+
+def get_job_info(
+    job_id: str, task_id: ObjectId, resource: Resource
+) -> SlurmInfo | None:
+    """Get job information from SLURM queue using squeue"""
+    try:
+        stdout = run_squeue_for_job(job_id)
+        # logger.info(f"task_id: {task_id} running squeue for job {job_id}: result: {stdout}")
+
+        if not stdout or stdout.strip() == "":
+            # after a while slurm will stop returning info for jobs that are no longer running
+            return None
+
+        lines = stdout.splitlines()
+        # logger.info(f"task_id: {task_id} slurm info for job {job_id}: {lines}")
+        if len(lines) < 2:
+            return None
+        # The first line is the header; the second line is the single info line
+
+        info_line = lines[1].strip()
+        if not info_line:
+            return None
+        # Split the single line into parts separated by whitespace
+        parts = re.split(r"\s+", info_line)
+        # logger.info(f"task_id: {task_id} slurm info for job {job_id}: {parts}")
+        # Expected default squeue columns:
+        # JOBID PARTITION NAME USER ST TIME NODES NODELIST(REASON)
+        name = parts[2] if len(parts) > 2 else ""
+        user = parts[3] if len(parts) > 3 else ""
+        code = parts[4] if len(parts) > 4 else ""
+        time_str = parts[5] if len(parts) > 5 else ""
+        nodelist_raw = parts[7] if len(parts) > 7 else ""
+        # Split nodelist on commas or whitespace, filter empties
+        nodes = [n for n in re.split(r"[,\s]+", nodelist_raw) if n]
+
+        slurm_info = SlurmInfo(
+            node_registry=task_id,
+            resource=resource,
+            job_id=job_id,
+            updated=datetime.now(),
+            name=name,
+            user=user,
+            code=code,
+            time=time_str,
+            nodes=nodes,
+        )
+
+        return slurm_info
+    except Exception as e:
+        logger.exception(f"Error getting job info for {job_id}: {str(e)}")
+        return None
+
+
+async def clean_slurm_info(resource: Resource, user: str | None = None) -> None:
+    """Clean up old slurm info entries"""
+    try:
+        squeue_cmd = "squeue"
+        if user:
+            squeue_cmd += f" -u {user}"
+
+        if context.config.docker:
+            watchdog_id = f"slurm_{uuid.uuid4()}"
+            queue_dir = context.config.workdir / "queue"
+            result = submit_to_watchdog(squeue_cmd, watchdog_id, queue_dir=queue_dir)
+        else:
+            result = subprocess.run(
+                squeue_cmd,
+                shell=True,
+                capture_output=True,
+                text=True,
+                timeout=30,
+            )
+
+        if result.returncode == 0:
+            active_job_ids = set()
+            for line in result.stdout.splitlines():
+                parts = line.split()
+                if not parts or parts[0] == "JOBID":
+                    continue
+                active_job_ids.add(parts[0])
+
+            # Find all SLURM info entries for this resource
+            # the user id is truncated on saving
+            running_jobs = await context.db.find(
+                SlurmInfo, SlurmInfo.resource.value == resource.value
+            )
+            # logger.info(f"Found {running_jobs} slurm info entries for {resource}")
+            # logger.info(f"Active job IDs: {active_job_ids} Slurm info IDs: {[job.job_id for job in running_jobs]}")
+            # logger.info(f"User: {user} resource: {resource} ")
+
+            # Delete entries for jobs that are no longer running
+            for job in running_jobs:
+                if job.job_id not in active_job_ids:
+                    await context.db.delete(job)
+                    logger.info(f"Deleted SLURM info for completed job {job.job_id}")
+
+    except Exception as e:
+        logger.exception(f"Error cleaning slurm info for {resource}: {str(e)}")