simstack 2.0.1__py3-none-any.whl

simstack/core/services/base_service.py

@@ -0,0 +1,191 @@
+import asyncio
+import logging
+import os
+import socket
+import sys
+import platform
+import subprocess
+from abc import ABC, abstractmethod
+from datetime import datetime
+from pathlib import Path
+from odmantic import ObjectId
+
+from simstack.core.context import context
+from simstack.models.parameters import Resource
+from simstack.models.runner_model import RunnerEvent, RunnerType, RunnerEventEnum
+from simstack.util.runner_utils import make_git_status_list
+
+logger = logging.getLogger("NodeRunner")
+
+class BaseService(ABC):
+    """A managed periodic service that can be stopped gracefully"""
+
+    def __init__(self, name: str, resource: Resource, interval: int, shutdown_event: asyncio.Event = None):
+        self._name = name
+        self._resource = resource
+        self._interval: int = interval
+        self._stop_event = asyncio.Event()
+        self._shutdown_event = shutdown_event
+        self._task = None
+        self._consecutive_failures = 0
+        self._log_traceback_on_failure = True
+
+        # Common identity attributes
+        self._pid = os.getpid()
+        self._hostname = socket.gethostname()
+        self._username = os.environ.get("USER", os.environ.get("USERNAME", "unknown"))
+        self._time_started = datetime.now()
+
+    def _get_uptime_string(self) -> str:
+        time_diff = datetime.now() - self._time_started
+        total_seconds = int(time_diff.total_seconds())
+        days, remainder = divmod(total_seconds, 86400)
+        hours, remainder = divmod(remainder, 3600)
+        minutes, seconds = divmod(remainder, 60)
+        return f"{days}d {hours}h {minutes}m {seconds}s"
+
+    async def write_node_event(self, event: RunnerEventEnum, node_id: ObjectId, message: str = None):
+        runner_event = RunnerEvent(
+            runner_type=RunnerType.NODE_RUNNER,
+            event=event,
+            pid=self._pid,
+            hostname=self._hostname,
+            user=self._username,
+            resource=self._resource,
+            node_id=node_id,
+            message=message,
+        )
+        await context.db.save(runner_event)
+
+    async def write_resource_event(self, event: RunnerEventEnum, message: str = None):
+        git_list = make_git_status_list()
+        if event == RunnerEventEnum.ALIVE:
+            uptime = self._get_uptime_string()
+            message = f"Uptime: {uptime}"
+
+            runner_event = await context.db.find_one(
+                RunnerEvent,
+                (RunnerEvent.runner_type == RunnerType.RESOURCE_RUNNER)
+                & (RunnerEvent.resource == self._resource)
+                & (RunnerEvent.event == RunnerEventEnum.ALIVE)
+                & (RunnerEvent.pid == self._pid),
+            )
+            if runner_event:
+                runner_event.message = message
+                runner_event.timestamp = datetime.now()
+                runner_event.git_status = git_list
+                await context.db.save(runner_event)
+                return
+
+        runner_event = RunnerEvent(
+            runner_type=RunnerType.RESOURCE_RUNNER,
+            pid=self._pid,
+            hostname=self._hostname,
+            user=self._username,
+            timestamp=datetime.now(),
+            git_status=git_list,
+            event=event,
+            resource=self._resource,
+            message=message,
+        )
+        await context.db.save(runner_event)
+
+    async def _run_loop(self):
+        """Internal loop that respects the stop event"""
+        logger.info(f"Service {self._name} started.")
+        while not self._stop_event.is_set():
+            start_time = asyncio.get_event_loop().time()
+            try:
+                await self.execute()
+                self._consecutive_failures = 0
+            except Exception as e:
+                self._consecutive_failures += 1
+                msg = f"Error in service {self._name} (failure {self._consecutive_failures}/3): {e}"
+                if self._log_traceback_on_failure:
+                    logger.exception(msg)
+                else:
+                    logger.error(msg)
+
+                if self._consecutive_failures >= 3:
+                    logger.error(f"Service {self._name} failed 3 times in a row. Shutting down.")
+                    self._stop_event.set()
+                    if self._shutdown_event:
+                        self._shutdown_event.set()
+
+            # Calculate wait time to maintain interval regardless of execution duration
+            elapsed = asyncio.get_event_loop().time() - start_time
+            wait_time = max(0, int(self._interval - elapsed))
+
+            try:
+                # Wait for interval OR until stop event is set
+                await asyncio.wait_for(self._stop_event.wait(), timeout=wait_time)
+            except asyncio.TimeoutError:
+                # Normal path: the interval passed, loop again
+                pass
+
+        logger.info(f"Service {self._name} stopped.")
+
+    @abstractmethod
+    async def execute(self):
+        raise NotImplementedError("Services must implement execute()")
+
+    def start(self):
+        if self._task is None or self._task.done():
+            self._stop_event.clear()
+            self._task = asyncio.create_task(self._run_loop())
+        return self._task
+
+    async def stop(self):
+        self._stop_event.set()
+        if self._task:
+            await self._task
+
+
+class RestartService(BaseService, ABC):
+    """
+    Intermediate abstract class for services that can trigger a runner restart.
+    """
+
+    def __init__(self, name: str, resource: Resource, interval: int):
+        super().__init__(name, resource, interval)
+        self._pid_file = context.config.workdir / "runner.pid"
+
+    async def trigger_restart(self):
+        """Spawns an independent process to kill current PID and start new one"""
+        current_pid = os.getpid()
+
+        # Write current PID to file so the restarter knows who to kill
+        self._pid_file.write_text(str(current_pid))
+
+        # Command to restart depends on OS
+        # We use sys.executable to ensure we use the same Python interpreter
+        # We'll call the runner module again
+        # IMPORTANT: We need the path to the original runner.py script.
+        # Since we are now in services/base_service.py, we need to go up two levels.
+        # But actually Path(__file__).resolve() in the original runner.py was used.
+        # We should probably pass the runner script path or rely on how it was invoked.
+        # In runner.py it was: script_path = Path(__file__).resolve()
+        
+        # Let's try to find runner.py relative to this file.
+        script_path = (Path(__file__).parent.parent / "runner.py").resolve()
+
+        log_file = context.config.workdir / "runner.out"
+
+        args = [sys.executable, str(script_path)] + sys.argv[1:]
+
+        # We need a small helper script or a one-liner that:
+        # 1. Waits for this process to die (to avoid port/resource conflicts)
+        # 2. Starts the new one
+        if platform.system() == "Windows":
+            # Windows 'start' command handles detachment well
+            # Use ping for a 2-second delay as 'timeout' fails in non-interactive shells
+            cmd = f"taskkill /F /PID {current_pid} && ping 127.0.0.1 -n 3 > nul && {' '.join(args)} >> \"{log_file}\" 2>&1"
+            subprocess.Popen(cmd, shell=True, creationflags=subprocess.CREATE_NEW_PROCESS_GROUP)
+        else:
+            # Linux: use a subshell that nohup/disowns
+            # Kill, wait, and start
+            cmd = f"kill -9 {current_pid} && sleep 2 && {' '.join(args)} >> \"{log_file}\" 2>&1"
+            subprocess.Popen(["/bin/bash", "-c", cmd], start_new_session=True)
+
+        # The above commands kill us, so this line might not even log
+        logger.info("Restart signal sent. Goodbye!")

simstack/core/services/file_transfer_service.py

@@ -0,0 +1,117 @@
+import asyncio
+import logging
+from pathlib import Path
+from typing import Any
+
+from simstack.core.context import context
+from simstack.core.services.base_service import BaseService
+from simstack.models.parameters import Resource
+from simstack.util.file_transfer_client import (
+    FileTransferClient,
+    FileTransferError,
+    resolve_instance_path,
+)
+
+logger = logging.getLogger("NodeRunner")
+
+
+class FileTransferService(BaseService):
+    """Poll the SimStack server for hidden file transfer uploads assigned to this runner."""
+
+    def __init__(
+        self,
+        resource: Resource,
+        interval: int = 10,
+        max_concurrent: int = 2,
+        shutdown_event: asyncio.Event | None = None,
+    ) -> None:
+        super().__init__(
+            "FileTransfer", resource, interval, shutdown_event=shutdown_event
+        )
+        self._log_traceback_on_failure = False
+        self._semaphore = asyncio.Semaphore(max_concurrent)
+        self._running_tasks: set[asyncio.Task[bool]] = set()
+        self._client: FileTransferClient | None = None
+        self._configuration_checked = False
+
+    def _get_client(self) -> FileTransferClient | None:
+        if self._client is not None:
+            return self._client
+        if self._configuration_checked:
+            return None
+        self._configuration_checked = True
+        self._client = FileTransferClient.from_context(required=False)
+        if self._client is None:
+            logger.info(
+                "File transfer service disabled; SIMSTACK_SERVER_URL or SIMSTACK_RUNNER_TOKEN is not configured."
+            )
+        return self._client
+
+    async def execute(self) -> None:
+        completed_tasks = {task for task in self._running_tasks if task.done()}
+        for task in completed_tasks:
+            try:
+                await task
+            except Exception as exc:
+                logger.exception("File transfer task completed with error: %s", exc)
+            self._running_tasks.remove(task)
+
+        client = self._get_client()
+        if client is None:
+            return
+
+        try:
+            transfers = client.list_transfers(role="source", status="created", limit=10)
+        except Exception as exc:
+            raise exc
+
+        for transfer in transfers:
+            transfer_id = str(transfer.get("transfer_id") or "")
+            if not transfer_id:
+                continue
+            if any(
+                getattr(task, "transfer_id", None) == transfer_id
+                for task in self._running_tasks
+            ):
+                continue
+            task = asyncio.create_task(self._run_with_semaphore(client, transfer))
+            setattr(task, "transfer_id", transfer_id)
+            self._running_tasks.add(task)
+
+    async def _run_with_semaphore(
+        self, client: FileTransferClient, transfer: dict[str, Any]
+    ) -> bool:
+        async with self._semaphore:
+            return await asyncio.to_thread(self._upload_transfer, client, transfer)
+
+    def _upload_transfer(
+        self, client: FileTransferClient, transfer: dict[str, Any]
+    ) -> bool:
+        transfer_id = str(transfer.get("transfer_id"))
+        source_path = transfer.get("source_path")
+        if not source_path:
+            client.fail_transfer(
+                transfer_id,
+                error_message="Transfer is missing source_path metadata.",
+                error_code="SOURCE_PATH_MISSING",
+            )
+            return False
+
+        path = resolve_instance_path(str(source_path), Path(context.config.workdir))
+        try:
+            client.upload_file(transfer_id, path)
+            logger.info("Uploaded FileStack transfer %s from %s", transfer_id, path)
+            return True
+        except FileTransferError as exc:
+            logger.warning("File transfer upload failed for %s: %s", transfer_id, exc)
+            try:
+                client.fail_transfer(
+                    transfer_id,
+                    error_message=str(exc),
+                    error_code="SOURCE_UPLOAD_FAILED",
+                )
+            except Exception:
+                logger.exception(
+                    "Failed to report failed file transfer %s", transfer_id
+                )
+            return False

simstack/core/services/git_restart_service.py

@@ -0,0 +1,12 @@
+from simstack.models.parameters import Resource
+from simstack.core.services.base_service import BaseService
+
+class GitRestartService(BaseService):
+    """Service that runs an external script to check git/restart the runner"""
+
+    def __init__(self, resource: Resource, interval):
+        super().__init__("GitRestart", resource, interval)
+        self._resource_name = str(resource)
+
+    def execute(self):
+        raise NotImplementedError("GitRestartService must implement execute()")

simstack/core/services/git_uv_update_service.py

@@ -0,0 +1,105 @@
+import asyncio
+import logging
+from pathlib import Path
+import hashlib
+
+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")
+
+def get_file_checksum(filepath: Path) -> str:
+    """Calculate SHA256 checksum of a file"""
+    if not filepath.exists():
+        return ""
+    hash_sha256 = hashlib.sha256()
+    with open(filepath, "rb") as f:
+        for chunk in iter(lambda: f.read(4096), b""):
+            hash_sha256.update(chunk)
+    return hash_sha256.hexdigest()
+
+class GitUvUpdateService(RestartService):
+    """
+    Advanced service that performs git pull and uv lock upgrades.
+    If changes are detected, it restarts the runner.
+    """
+
+    def __init__(self, resource: Resource, interval):
+        super().__init__("GitUvUpdate", resource, interval)
+        # Resolve project root (assuming we are in src/simstack/core/runner.py)
+        self._project_dir = context.config.project_root.resolve(strict=True)
+        self._uv_lock_path = context.config.project_root / "uv.lock"
+        self._uv_extra_depencency_path = context.config.project_root / "user_extra_config.toml"
+        self.extras=False
+        self.desired_extras=[]
+        #check if the file _uv_extra_depencency_path exists if yes - read toml and then extras=True
+        try :
+            import tomllib
+            logger.info(f" trying to read user-specified (extras) optional dependencies from {self._uv_extra_depencency_path}")
+            with open (self._uv_extra_depencency_path,"rb") as f:
+                desired_extras_all=tomllib.load(f)
+                # Access the nested section - this was the issue with the keys
+                extras_section = desired_extras_all.get("optional_dependencies_desired_by_ressource", {})
+
+            
+            # Access raw value without triggering validation against allowed_resources
+            resource_name = object.__getattribute__(resource, "__dict__").get("value") or str(resource)
+            #logger.info(f" name {resource_name} , type {type(resource_name)}, compared against {extras_section.keys()}") 
+            if resource_name in extras_section.keys():
+                self.desired_extras = extras_section[resource_name]
+                self.extras = True
+                
+        except Exception as e:
+            logger.warning(f"resource {resource} spec  with user file {self._uv_extra_depencency_path} failed with the exception {e}")
+
+
+        
+
+    async def _run_command(self, cmd: list, ignore_error=False) -> str:
+        """Run a shell command and return output"""
+        process = await asyncio.create_subprocess_exec(
+            *cmd,
+            cwd=str(self._project_dir),
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE
+        )
+        stdout, stderr = await process.communicate()
+        if process.returncode != 0:
+            if ignore_error:
+                logger.warning(f"Command {' '.join(cmd)} failed: {stderr.decode()}")
+            else:
+                raise RuntimeError(f"Command {' '.join(cmd)} failed: {stderr.decode()}")
+        return stdout.decode().strip()
+    
+    async def execute(self):
+        # 1. Git Pull
+        # Check checksum before and after pull to see if Git brought a new lockfile
+        # old_uv_checksum = get_file_checksum(self._uv_lock_path)
+        
+        # Ensure we don't have local lockfile changes that block the pull
+
+        await self._run_command(["git", "stash"], ignore_error=True)
+
+        git_output = await self._run_command(["git", "pull", "--recurse-submodules" ])
+        git_changed = "Already up to date." not in git_output
+
+        # Clear the stash now that we've pulled
+        # await self._run_command(["git", "stash", "drop"], ignore_error=True)
+
+        if git_changed: # or uv_locally_upgraded:
+            command_list=["uv", "sync", "--locked"]
+            if self.extras and len(self.desired_extras) > 0:
+                
+                for desired_extra in self.desired_extras:
+                    command_list.extend(["--extra", str(desired_extra)])
+                    logger.info(f"user specified extra dependency {desired_extra}")
+                    
+                await self._run_command(command_list)  # Update local .venv but with extras as specified by the user for the ressource
+            else:
+                await self._run_command(command_list)  # Update local .venv
+            reason = "Git pull" if git_changed else "Local UV upgrade"
+            await self.write_resource_event(RunnerEventEnum.SHUTDOWN, message=reason)
+            logger.info(f"Update detected ({reason}). Triggering restart...")
+            await self.trigger_restart()

simstack/core/services/node_execution_service.py

@@ -0,0 +1,182 @@
+import asyncio
+import logging
+import platform
+import subprocess
+
+from simstack.core.context import context
+from simstack.core.definitions import TaskStatus
+from simstack.core.node import node_from_database
+from simstack.core.node_claim import claim_submitted_node
+from simstack.core.run_docker import run_docker
+from simstack.core.services.base_service import BaseService
+from simstack.core.submit_node import submit_node
+from simstack.models import NodeRegistry
+from simstack.models.parameters import Resource
+from simstack.models.runner_model import RunnerEventEnum
+
+logger = logging.getLogger("NodeRunner")
+
+
+async def run_node_from_registry(registry_entry: NodeRegistry) -> bool:
+    # Create the node from the registry entry
+    node = await node_from_database(registry_entry)
+    if not node:
+        logger.error(
+            f"Failed to create node from registry entry task_id: {registry_entry.id} on resource {context.config.resource}"
+        )
+        registry_entry.status = TaskStatus.FAILED
+        await context.db.save(registry_entry)
+        return False
+    registry_entry = node.registry_entry  # it may have changed
+    assert registry_entry is not None
+    if (
+        node.status == TaskStatus.RETRIEVED
+        or node.status == TaskStatus.SUBMITTED
+        or node.status == TaskStatus.SLURM_QUEUED
+    ) or (
+        node.status == TaskStatus.COMPLETED and registry_entry.parameters.force_rerun
+    ):
+        await node.execute_node_locally()
+    else:
+        logger.info(
+            f"task_id: {registry_entry.id} skipping task: {registry_entry.name} with status {registry_entry.status}"
+        )
+    return bool(node.status == TaskStatus.COMPLETED)
+
+
+class NodeExecutionService(BaseService):
+    def __init__(
+        self,
+        resource: Resource,
+        interval: int,
+        max_concurrent: int,
+        shutdown_event: asyncio.Event | None,
+        detach: bool = True,
+        is_default: bool = False,
+    ) -> None:
+        super().__init__(
+            "JobPolling", resource, interval, shutdown_event=shutdown_event
+        )
+        self._resource_name = str(resource)
+        self._semaphore = asyncio.Semaphore(max_concurrent)
+        self._running_tasks: set[asyncio.Task[bool]] = set()
+        self._started = False
+        self._detach = detach
+        self._is_default = is_default
+
+    async def run_node(self, registry_entry: NodeRegistry) -> bool:
+        """Run a single node by its ID from the database"""
+
+        await self.write_node_event(RunnerEventEnum.NODE_STARTED, registry_entry.id)
+        try:
+            logger.info(
+                f"Running node task_id: {registry_entry.id} on resource {context.config.resource} with status {registry_entry.status}"
+            )
+            queue = (
+                registry_entry.parameters.queue
+                if hasattr(registry_entry.parameters, "queue")
+                else "default"
+            )
+            if queue is None:
+                logger.error(
+                    f"Queue parameter not found for task_id: {registry_entry.id}"
+                )
+                return False
+
+            if queue == "slurm-queue":
+                return await submit_node(registry_entry)
+            elif queue == "docker":
+                return await run_docker(registry_entry)
+
+            elif queue == "default":
+                if self._detach:
+                    # Spawn independent process that survives when the runner dies
+                    cmd = [
+                        "uv",
+                        "run",
+                        "--directory",
+                        str(context.config.project_root),
+                        "run_node",
+                        "--node-id",
+                        str(registry_entry.id),
+                        "--resource",
+                        str(self._resource_name),
+                    ]
+
+                    # Use platform specific flags to ensure the process survives if runner is killed
+                    creationflags = 0
+                    if platform.system() == "Windows":
+                        create_new_process_group = getattr(
+                            subprocess, "CREATE_NEW_PROCESS_GROUP", 0
+                        )
+                        detached_process = getattr(subprocess, "DETACHED_PROCESS", 0)
+                        creationflags = create_new_process_group | detached_process
+
+                    try:
+                        process = await asyncio.create_subprocess_exec(
+                            *cmd,
+                            stdout=asyncio.subprocess.DEVNULL,
+                            stderr=asyncio.subprocess.DEVNULL,
+                            creationflags=creationflags,
+                            start_new_session=True
+                            if platform.system() != "Windows"
+                            else False,
+                        )
+                    except Exception as e:
+                        logger.error(
+                            f"Failed to spawn detached process for task id: {registry_entry.id}. task_id: {registry_entry.id}. "
+                            f"Command: {' '.join(cmd)}. Error: {str(e)}"
+                        )
+                        raise
+
+                    logger.info(
+                        f"Spawned detached process for task_id: {registry_entry.id} with PID: {process.pid}"
+                    )
+                    return True
+                else:
+                    return await run_node_from_registry(registry_entry)
+            else:
+                logger.error(
+                    f"Queue {queue} not supported for task_id: {registry_entry.id}"
+                )
+                return False
+
+        except Exception as e:
+            logger.exception(
+                f"Error running node task_id: {registry_entry.id} on resource {context.config.resource} : {str(e)}"
+            )
+            if registry_entry:
+                registry_entry.status = TaskStatus.FAILED
+                await context.db.save(registry_entry)
+            return False
+
+    async def execute(self) -> None:
+        if not self._started:
+            await self.write_resource_event(RunnerEventEnum.RUNNER_STARTED)
+            self._started = True
+
+        # Clean up the completed tasks
+        completed_tasks = {task for task in self._running_tasks if task.done()}
+        for task in completed_tasks:
+            try:
+                await task
+            except Exception as e:
+                logger.exception(f"Task completed with error: {e}")
+            self._running_tasks.remove(task)
+
+        # Load tasks
+        registry_entry_list = await context.db.load_waiting_tasks_for_resource(self._resource_name)
+
+        if registry_entry_list:
+            logger.info(
+                f"Retrieved {len(registry_entry_list)} tasks for {self._resource_name}"
+            )
+            for entry in registry_entry_list:
+                if not await claim_submitted_node(entry):
+                    continue
+                task = asyncio.create_task(self._run_with_semaphore(entry))
+                self._running_tasks.add(task)
+
+    async def _run_with_semaphore(self, entry: NodeRegistry) -> bool:
+        async with self._semaphore:
+            return await self.run_node(entry)

simstack/core/services/resource_branch_monitor_service.py

@@ -0,0 +1,71 @@
+import asyncio
+import logging
+
+from simstack.core.context import context
+from simstack.models.parameters import Resource
+from simstack.models.runner_model import RunnerEventEnum
+from simstack.models.resource_definition import ResourceDefinition
+from simstack.core.services.base_service import RestartService
+
+logger = logging.getLogger("NodeRunner")
+
+class ResourceBranchMonitorService(RestartService):
+    """
+    Monitors the ResourceDefinition for the current resource.
+    If the git_branch field changes, it stashes, switches branch, syncs, and restarts.
+    """
+
+    def __init__(self, resource: Resource, interval: int):
+        super().__init__("ResourceBranchMonitor", resource, interval)
+        self._project_dir = context.config.project_root.resolve(strict=True)
+
+    async def _run_command(self, cmd: list) -> str:
+        process = await asyncio.create_subprocess_exec(
+            *cmd,
+            cwd=str(self._project_dir),
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE
+        )
+        stdout, stderr = await process.communicate()
+        if process.returncode != 0:
+            logger.error(f"Command {' '.join(cmd)} failed: {stderr.decode()}")
+            return ""
+        return stdout.decode().strip()
+
+    async def _get_current_branch(self) -> str:
+        return await self._run_command(["git", "rev-parse", "--abbrev-ref", "HEAD"])
+
+    async def execute(self):
+        resource_def = await context.db.find_one(
+            ResourceDefinition, 
+            ResourceDefinition.resource_str == str(self._resource)
+        )
+        
+        if not resource_def:
+            return
+
+        target_branch = resource_def.git_branch
+        current_branch = await self._get_current_branch()
+
+        if target_branch and current_branch and target_branch != current_branch:
+            logger.info(f"Branch change detected: {current_branch} -> {target_branch}. Updating...")
+            
+            # 1. Stash existing changes
+            await self._run_command(["git", "stash"])
+            
+            # 2. Checkout new branch
+            checkout_res = await self._run_command(["git", "checkout", target_branch])
+            if not checkout_res:
+                logger.error(f"Failed to checkout branch {target_branch}")
+                return
+
+            # 3. UV Sync
+            logger.info("Running uv sync --locked...")
+            await self._run_command(["uv", "sync", "--locked"])
+
+            # 4. Trigger Restart
+            await self.write_resource_event(
+                RunnerEventEnum.SHUTDOWN, 
+                message=f"Branch switched to {target_branch}"
+            )
+            await self.trigger_restart()