simstack 2.0.1__py3-none-any.whl

simstack/core/services/runner_cleanup_service.py

@@ -0,0 +1,149 @@
+import logging
+import os
+from datetime import datetime, timedelta
+from pathlib import Path
+from typing import Any
+
+from simstack.core.context import context
+from simstack.models.files import FileStack
+from simstack.models.parameters import Resource
+from simstack.models.runner_model import RunnerEvent, RunnerType
+from simstack.core.services.base_service import BaseService
+
+logger = logging.getLogger("NodeRunner")
+
+
+def _resource_to_str(resource: Any) -> str:
+    if hasattr(resource, "resource_str"):
+        return str(getattr(resource, "resource_str"))
+    raw = getattr(resource, "__dict__", {}).get("value")
+    if raw is not None:
+        return str(raw)
+    return str(resource)
+
+
+def _cache_ttl_days() -> int:
+    raw = os.environ.get("SIMSTACK_FILE_CACHE_TTL_DAYS")
+    if raw is None:
+        return 14
+    try:
+        value = int(raw)
+    except ValueError:
+        return 14
+    return max(value, 0)
+
+
+def _cache_cleanup_enabled() -> bool:
+    return os.environ.get("SIMSTACK_FILE_CACHE_CLEANUP_ENABLED", "1").lower() not in {
+        "0",
+        "false",
+        "no",
+    }
+
+
+def _as_naive_datetime(value: datetime | None) -> datetime | None:
+    if value is None:
+        return None
+    if value.tzinfo is None:
+        return value
+    return value.astimezone().replace(tzinfo=None)
+
+
+class RunnerCleanupService(BaseService):
+    """
+    Service that cleans up old RunnerEvent logs and expired local file cache
+    copies for the current resource.
+    """
+
+    def __init__(self, resource: Resource, interval: int = 300) -> None:
+        # Default interval 5 minutes
+        super().__init__("RunnerCleanup", resource, interval)
+
+    async def execute(self) -> None:
+        cutoff_time = datetime.now() - timedelta(minutes=30)
+
+        # Find and delete events matching the criteria
+        old_events = await context.db.find(
+            RunnerEvent,
+            (RunnerEvent.runner_type == RunnerType.RESOURCE_RUNNER)
+            & (RunnerEvent.resource == self._resource)
+            & (RunnerEvent.timestamp < cutoff_time),
+        )
+
+        if old_events:
+            logger.info(
+                f"Cleaning up {len(old_events)} old RunnerEvent logs for resource {self._resource}"
+            )
+            for event in old_events:
+                await context.db.delete(event)
+
+        if _cache_cleanup_enabled():
+            await self._cleanup_file_cache()
+
+    async def _cleanup_file_cache(self) -> None:
+        ttl_days = _cache_ttl_days()
+        cutoff_time = datetime.now() - timedelta(days=ttl_days)
+        resource_name = _resource_to_str(self._resource)
+        workdir = Path(context.config.workdir).resolve()
+
+        file_stacks = await context.db.find_all(FileStack)
+        updated_count = 0
+        deleted_count = 0
+
+        for file_stack in file_stacks:
+            changed = False
+            for location in file_stack.locations or []:
+                if _resource_to_str(location.resource) != resource_name:
+                    continue
+                if getattr(location, "location_type", "local_path") != "local_path":
+                    continue
+                if getattr(location, "status", "available") != "available":
+                    continue
+                if not getattr(location, "is_cached", False):
+                    continue
+
+                accessed_at = _as_naive_datetime(
+                    getattr(location, "last_accessed_at", None) or location.created_at
+                )
+                if accessed_at and accessed_at > cutoff_time:
+                    continue
+
+                path = Path(location.path)
+                resolved_path = path if path.is_absolute() else workdir / path
+                try:
+                    resolved_path = resolved_path.resolve()
+                    if (
+                        workdir != resolved_path
+                        and workdir not in resolved_path.parents
+                    ):
+                        logger.warning(
+                            "Skipping cached FileInstance outside workdir: file_stack=%s path=%s",
+                            file_stack.id,
+                            resolved_path,
+                        )
+                        continue
+                    if resolved_path.exists():
+                        resolved_path.unlink()
+                        deleted_count += 1
+                    location.status = "deleted"
+                    location.expires_at = datetime.now()
+                    changed = True
+                except Exception as exc:
+                    logger.warning(
+                        "Failed to clean cached FileInstance file_stack=%s path=%s: %s",
+                        file_stack.id,
+                        resolved_path,
+                        exc,
+                    )
+
+            if changed:
+                await context.db.save(file_stack)
+                updated_count += 1
+
+        if updated_count:
+            logger.info(
+                "Cleaned file cache for resource %s: marked %s FileStack(s), deleted %s file(s)",
+                resource_name,
+                updated_count,
+                deleted_count,
+            )

simstack/core/services/runner_manager.py

@@ -0,0 +1,163 @@
+import asyncio
+import logging
+import os
+import platform
+from pathlib import Path
+
+from simstack.core.services.base_service import BaseService
+from simstack.core.context import context
+from simstack.models.parameters import Resource
+from simstack.core.services.node_execution_service import NodeExecutionService
+from simstack.core.services.runner_status_service import RunnerStatusService
+from simstack.core.services.runner_cleanup_service import RunnerCleanupService
+from simstack.core.services.file_transfer_service import FileTransferService
+from simstack.core.services.slurm_status_service import SlurmStatusService
+from simstack.core.services.resource_branch_monitor_service import (
+    ResourceBranchMonitorService,
+)
+from simstack.core.services.stop_check_service import StopCheckService
+from simstack.core.services.git_uv_update_service import GitUvUpdateService
+from simstack.core.services.timeout_restart_service import TimeoutRestartService
+
+logger = logging.getLogger("NodeRunner")
+
+
+class RunnerManager:
+    def __init__(
+        self,
+        resource: Resource,
+        detach: bool = True,
+        no_pull: bool = False,
+        is_default: bool = False,
+        with_file_transfer: bool = True,
+    ) -> None:
+        self._resource = resource
+        self._detach = detach
+        self._no_pull = no_pull
+        self._pid = os.getpid()
+        self._services: list[BaseService] = []
+        self._shutdown_event: asyncio.Event = asyncio.Event()
+        self._pid_file: Path = Path(context.config.workdir) / f"runner_{resource}.pid"
+        self._is_default = is_default
+        self._with_file_transfer = with_file_transfer
+
+    def _is_process_running(self, pid: int) -> bool:
+        """Check if a process with given PID is running and is a simstack_runner process"""
+        try:
+            if platform.system() == "Windows":
+                # On Windows, os.kill with signal 0 doesn't work, use tasklist
+                import subprocess
+
+                result = subprocess.run(
+                    ["tasklist", "/FI", f"PID eq {pid}", "/V", "/FO", "CSV"],
+                    capture_output=True,
+                    text=False,  # keep bytes; avoid UnicodeDecodeError from console code pages
+                )
+                # Check if process exists and command line contains simstack_runner
+                stdout = result.stdout.decode("utf-8", errors="ignore")
+                return str(pid) in stdout and "simstack_runner" in stdout.lower()
+            else:
+                # Check if process exists
+                os.kill(pid, 0)
+                # Verify it's a simstack_runner process by checking command line
+                try:
+                    with open(f"/proc/{pid}/cmdline", "r") as f:
+                        cmdline = f.read()
+                        return "simstack_runner" in cmdline
+                except (FileNotFoundError, PermissionError):
+                    # If we can't read cmdline, fall back to just checking if process exists
+                    return True
+        except (OSError, ProcessLookupError):
+            return False
+
+    def _check_existing_runner(self) -> None:
+        """Check if another runner for this resource is already running on this host"""
+        if self._pid_file.exists():
+            try:
+                existing_pid = int(self._pid_file.read_text().strip())
+                if existing_pid != self._pid and self._is_process_running(existing_pid):
+                    error_msg = (
+                        f"Another runner for resource '{self._resource}' is already running "
+                        f"on this host with PID {existing_pid}. Exiting."
+                    )
+                    logger.error(error_msg)
+                    raise RuntimeError(error_msg)
+                else:
+                    logger.info(
+                        f"Stale PID file found for resource '{self._resource}'. "
+                        f"Overwriting with current PID {self._pid}."
+                    )
+            except (ValueError, OSError) as e:
+                logger.warning(
+                    f"Could not read PID file: {e}. Proceeding with startup."
+                )
+
+    async def stop_all_services(self) -> None:
+        """Gracefully stop all registered services"""
+        logger.info("Stopping all services...")
+        await asyncio.gather(
+            *(s.stop() for s in self._services), return_exceptions=True
+        )
+
+    async def run_nodes_for_resource(
+        self,
+        polling_interval: int = 5,
+        max_concurrent: int = 10,
+        timeout: int | None = None,
+    ) -> None:
+        """Orchestrates multiple independent services"""
+        # Check if another runner is already running for this resource
+        self._check_existing_runner()
+
+        # Save PID on startup
+        self._pid_file.write_text(str(self._pid))
+
+        self._services = [
+            NodeExecutionService(
+                self._resource,
+                polling_interval,
+                max_concurrent,
+                self._shutdown_event,
+                detach=self._detach,
+                is_default=self._is_default,
+            ),
+            RunnerStatusService(self._resource, interval=60),
+            RunnerCleanupService(self._resource, interval=300),
+            SlurmStatusService(self._resource, interval=60),
+            StopCheckService(
+                self._resource, interval=10, shutdown_event=self._shutdown_event
+            ),
+        ]
+
+        if self._with_file_transfer:
+            self._services.append(
+                FileTransferService(
+                    self._resource,
+                    interval=10,
+                    max_concurrent=2,
+                    shutdown_event=None,
+                )
+            )
+
+        if not self._no_pull:
+            self._services.append(GitUvUpdateService(self._resource, interval=60))
+            self._services.append(
+                ResourceBranchMonitorService(self._resource, interval=60)
+            )
+
+        # Add timeout restart service if timeout is specified
+        if timeout is not None:
+            self._services.append(TimeoutRestartService(self._resource, timeout))
+
+        # Start all services
+        service_tasks = [service.start() for service in self._services]
+
+        try:
+            # Wait for either a service task to finish or the shutdown event to be set
+            shutdown_task = asyncio.create_task(self._shutdown_event.wait())
+            tasks_to_wait = [*service_tasks, shutdown_task]
+            done, pending = await asyncio.wait(
+                tasks_to_wait, return_when=asyncio.FIRST_COMPLETED
+            )
+        finally:
+            await self.stop_all_services()

simstack/core/services/runner_status_service.py

@@ -0,0 +1,10 @@
+from simstack.models.parameters import Resource
+from simstack.models.runner_model import RunnerEventEnum
+from simstack.core.services.base_service import BaseService
+
+class RunnerStatusService(BaseService):
+    def __init__(self, resource: Resource, interval):
+        super().__init__("RunnerStatus", resource, interval)
+
+    async def execute(self):
+        await self.write_resource_event(RunnerEventEnum.ALIVE)

simstack/core/services/slurm_status_service.py

@@ -0,0 +1,67 @@
+import logging
+from datetime import datetime
+
+from simstack.core.context import context
+from simstack.core.definitions import TaskStatus
+from simstack.models import NodeRegistry
+from simstack.models.parameters import Resource
+from simstack.models.slurm_info import SlurmInfo
+from simstack.util.runner_utils import get_job_info, clean_slurm_info
+from simstack.core.services.base_service import BaseService
+
+logger = logging.getLogger("NodeRunner")
+
+
+class SlurmStatusService(BaseService):
+    def __init__(self, resource: Resource, interval: int) -> None:
+        super().__init__("SlurmStatus", resource, interval)
+        self._resource_name = str(resource)
+
+    async def execute(self) -> None:
+        try:
+            # logger.info(f"Running clean_slurm_info for {self._resource} and user {self._username}")
+            await clean_slurm_info(self._resource, user=self._username)
+
+            running_tasks = await context.db.find(
+                NodeRegistry,
+                (NodeRegistry.status == TaskStatus.RUNNING)
+                & (NodeRegistry.parameters.resource == self._resource),
+            )
+            queued_tasks = await context.db.find(
+                NodeRegistry,
+                (NodeRegistry.status == TaskStatus.SLURM_QUEUED)
+                & (NodeRegistry.parameters.resource == self._resource),
+            )
+            # logger.info(f"Checking Slurm status for {len(running_tasks)} running jobs on resource {self._resource}")
+
+            for task in list(running_tasks) + list(queued_tasks):
+                if task.job_id is not None:
+                    slurm_info = get_job_info(
+                        task.job_id, task.id, Resource(value=self._resource_name)
+                    )
+                    slurm_entry = await context.db.find_one(
+                        SlurmInfo, SlurmInfo.job_id == task.job_id
+                    )
+                    if slurm_info:
+                        if slurm_entry:
+                            slurm_entry.code = slurm_info.code
+                            slurm_entry.time = slurm_info.time
+                            slurm_entry.updated = datetime.now()
+                            await context.db.save(slurm_entry)
+                        else:
+                            await context.db.save(slurm_info)
+                    else:
+                        # Logic for finished/timed out jobs
+                        check_job = await context.db.find_one(
+                            NodeRegistry, NodeRegistry.id == task.id
+                        )
+                        if slurm_entry:
+                            await context.db.delete(slurm_entry)
+                        if check_job.status == TaskStatus.RUNNING:
+                            task.job_id = None
+                            task.status = TaskStatus.TIME_OUT
+                            await context.db.save(task)
+
+        except Exception as e:
+            logger.exception(f"Error checking Slurm status: {e}")
+            raise e

simstack/core/services/stop_check_service.py

@@ -0,0 +1,27 @@
+import asyncio
+import logging
+from pathlib import Path
+
+from simstack.core.context import context
+from simstack.models.parameters import Resource
+from simstack.models.runner_model import RunnerEventEnum
+from simstack.core.services.base_service import BaseService
+
+logger = logging.getLogger("NodeRunner")
+
+class StopCheckService(BaseService):
+    """Service that checks for STOP file and triggers shutdown"""
+
+    def __init__(self, resource: Resource, interval: int, shutdown_event: asyncio.Event):
+        super().__init__("StopCheck", resource, interval, shutdown_event=shutdown_event)
+
+    async def execute(self):
+        # Check for STOP file
+        for path in [context.config.workdir, context.config.project_root]:
+            stop_file = Path(path) / "STOP"
+            if stop_file.exists():
+                await self.write_resource_event(RunnerEventEnum.SHUTDOWN, message="STOP file found")
+                logger.info(f"STOP file found at {stop_file}, signaling shutdown...")
+                if self._shutdown_event:
+                    self._shutdown_event.set()
+                return

simstack/core/services/timeout_restart_service.py

@@ -0,0 +1,26 @@
+import logging
+
+from simstack.core.context import context
+from simstack.models.parameters import Resource
+from simstack.models.runner_model import RunnerEventEnum
+from simstack.core.services.base_service import RestartService
+
+logger = logging.getLogger("NodeRunner")
+
+class TimeoutRestartService(RestartService):
+    """Service that restarts the runner after a specified timeout"""
+
+    def __init__(self, resource: Resource, timeout_minutes: int):
+        super().__init__("TimeoutRestart", resource, interval=timeout_minutes * 60)
+        self._project_dir = context.config.project_root.resolve(strict=True)
+        self._timeout_minutes = timeout_minutes
+        self._executed = False
+
+    async def execute(self):
+        # Only execute once after the timeout interval
+        if not self._executed:
+            self._executed = True
+            logger.info(f"Timeout of {self._timeout_minutes} minutes reached. Triggering restart...")
+            await self.write_resource_event(RunnerEventEnum.SHUTDOWN,
+                                            message=f"Timeout restart after {self._timeout_minutes} minutes")
+            await self.trigger_restart()

simstack/core/simstack_result.py

@@ -0,0 +1,39 @@
+from typing import Optional, List
+
+from pydantic import BaseModel, Field, ConfigDict
+
+from simstack.core.definitions import TaskStatus
+from simstack.models.files import FileStack
+
+
+class SimstackResult(BaseModel):
+    model_config = ConfigDict(extra="allow")
+    """
+    Represents the results from a Simstack operation.
+
+    This class serves to encapsulate the status, messages, and files resulting
+    from a Simstack operation. It provides structured attributes for the error
+    message, operation message, and categorized files, ensuring clarity in the
+    representation of these aspects. It also supports additional fields through
+    flexible configuration.
+
+    :ivar status: Indicates the status of the task, defaulting to `TaskStatus.COMPLETED`.
+    :type status: TaskStatus
+    :ivar error_message: Optionally stores an error message if the task encountered issues.
+    :type error_message: Optional[str]
+    :ivar message: An optional general message providing additional information about the task.
+    :type message: Optional[str]
+    :ivar files: A list of `FileStack` objects that represent the primary files involved in
+        the operation. Defaults to an empty list.
+    :type files: List[FileStack]
+    :ivar info_files: A list of `FileStack` objects that provide additional or informational
+        files pertaining to the operation. Defaults to an empty list. These files are shown in the results
+        but not passed to the calling function
+    :type info_files: List[FileStack]
+    """
+
+    status: TaskStatus = TaskStatus.RETRIEVED
+    error_message: Optional[str] = None
+    message: Optional[str] = None
+    files: List[FileStack] = Field(default_factory=list)
+    info_files: List[FileStack] = Field(default_factory=list)