researchloop 0.1.0__py3-none-any.whl

researchloop/clusters/__init__.py

@@ -0,0 +1,4 @@
+from researchloop.clusters.monitor import JobMonitor
+from researchloop.clusters.ssh import SSHConnection, SSHManager
+
+__all__ = ["SSHConnection", "SSHManager", "JobMonitor"]

researchloop/clusters/monitor.py

@@ -0,0 +1,199 @@
+"""Job monitoring - polls active jobs via SSH and updates the database."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from datetime import datetime, timezone
+from typing import Any
+
+from researchloop.clusters.ssh import SSHManager
+from researchloop.db import queries
+from researchloop.schedulers.base import BaseScheduler
+
+logger = logging.getLogger(__name__)
+
+# If a job's heartbeat is older than this many seconds AND the job is not
+# visible in the scheduler queue, consider it abandoned.
+_HEARTBEAT_STALE_SECONDS = 5 * 60  # 5 minutes
+
+
+class JobMonitor:
+    """Monitors submitted jobs by periodically polling their status."""
+
+    def __init__(
+        self,
+        ssh_manager: SSHManager,
+        db: Any,
+        schedulers: dict[str, BaseScheduler],
+        config: Any = None,
+    ) -> None:
+        self.ssh_manager = ssh_manager
+        self.db = db
+        self.schedulers = schedulers
+        self.config = config
+        self._polling_task: asyncio.Task[None] | None = None
+        self._stop_event = asyncio.Event()
+
+    # ------------------------------------------------------------------
+    # Single-job check
+    # ------------------------------------------------------------------
+
+    async def check_job(self, sprint: dict[str, Any]) -> str:
+        """Check the status of the job for *sprint*.
+
+        Returns the scheduler-normalised status string:
+        ``pending``, ``running``, ``completed``, ``failed``, or ``unknown``.
+        """
+        job_id: str | None = sprint.get("job_id")
+        if not job_id:
+            logger.warning("Sprint %s has no job_id", sprint["id"])
+            return "unknown"
+
+        # Resolve scheduler for this sprint's study/cluster.
+        study = await queries.get_study(self.db, sprint["study_name"])
+        if study is None:
+            return "unknown"
+
+        cluster_name = study["cluster"]
+        scheduler = self.schedulers.get(cluster_name)
+        if scheduler is None:
+            logger.debug("No scheduler for cluster %r", cluster_name)
+            return "unknown"
+
+        # Resolve cluster config for SSH connection.
+        cluster_cfg = None
+        if self.config:
+            for c in self.config.clusters:
+                if c.name == cluster_name:
+                    cluster_cfg = c
+                    break
+
+        if cluster_cfg is None:
+            logger.debug("No cluster config for %r", cluster_name)
+            return "unknown"
+
+        try:
+            ssh = await self.ssh_manager.get_connection(
+                {
+                    "host": cluster_cfg.host,
+                    "port": cluster_cfg.port,
+                    "user": cluster_cfg.user,
+                    "key_path": cluster_cfg.key_path,
+                }
+            )
+            status = await scheduler.status(ssh, job_id)
+        except Exception:
+            logger.debug("SSH check failed for job %s", job_id, exc_info=True)
+            return "unknown"
+
+        logger.info("Sprint %s (job %s) status: %s", sprint["id"], job_id, status)
+        return status
+
+    # ------------------------------------------------------------------
+    # Poll all active jobs
+    # ------------------------------------------------------------------
+
+    async def poll_active_jobs(self) -> None:
+        """Check every active sprint, update the DB, and detect abandoned jobs."""
+        sprints = await queries.get_active_sprints(self.db)
+        if not sprints:
+            logger.debug("No active sprints to poll.")
+            return
+
+        logger.info("Polling %d active sprint(s)...", len(sprints))
+        now = datetime.now(timezone.utc)
+
+        for sprint in sprints:
+            sprint_id: str = sprint["id"]
+            try:
+                status = await self.check_job(sprint)
+            except Exception:
+                logger.exception("Error checking status for sprint %s", sprint_id)
+                status = "unknown"
+
+            # --- Abandoned-job detection ---
+            if status == "unknown":
+                # Check heartbeat from metadata_json
+                metadata_str = sprint.get("metadata_json")
+                heartbeat_str: str | None = None
+                if metadata_str:
+                    try:
+                        metadata = json.loads(metadata_str)
+                        heartbeat_str = metadata.get("last_heartbeat")
+                    except (json.JSONDecodeError, TypeError):
+                        pass
+
+                if heartbeat_str is not None:
+                    heartbeat = datetime.fromisoformat(heartbeat_str)
+                    if heartbeat.tzinfo is None:
+                        heartbeat = heartbeat.replace(tzinfo=timezone.utc)
+
+                    stale_seconds = (now - heartbeat).total_seconds()
+                    if stale_seconds > _HEARTBEAT_STALE_SECONDS:
+                        logger.warning(
+                            "Sprint %s appears abandoned: heartbeat %.0fs ago "
+                            "and job not in scheduler queue. Marking as failed.",
+                            sprint_id,
+                            stale_seconds,
+                        )
+                        status = "failed"
+
+            # Persist the updated status if it changed.
+            if status in ("completed", "failed"):
+                try:
+                    await queries.update_sprint(
+                        self.db,
+                        sprint_id,
+                        status=status,
+                        completed_at=datetime.now(timezone.utc).isoformat(),
+                    )
+                except Exception:
+                    logger.exception(
+                        "Failed to update DB status for sprint %s", sprint_id
+                    )
+
+    # ------------------------------------------------------------------
+    # Background polling loop
+    # ------------------------------------------------------------------
+
+    async def start_polling(self, interval: int = 120) -> None:
+        """Start a background task that polls active jobs every *interval* seconds."""
+        if self._polling_task is not None and not self._polling_task.done():
+            logger.warning("Polling is already running.")
+            return
+
+        self._stop_event.clear()
+        self._polling_task = asyncio.create_task(
+            self._poll_loop(interval), name="job-monitor-poll"
+        )
+        logger.info("Job monitor polling started (interval=%ds).", interval)
+
+    async def _poll_loop(self, interval: int) -> None:
+        """Internal loop executed by the background task."""
+        while not self._stop_event.is_set():
+            try:
+                await self.poll_active_jobs()
+            except Exception:
+                logger.exception("Unhandled error during job polling")
+
+            # Wait for the interval, but break early if stop is requested.
+            try:
+                await asyncio.wait_for(self._stop_event.wait(), timeout=interval)
+                # If we reach here, stop was requested.
+                break
+            except asyncio.TimeoutError:
+                # Normal timeout - continue polling.
+                pass
+
+    async def stop_polling(self) -> None:
+        """Stop the background polling task."""
+        self._stop_event.set()
+        if self._polling_task is not None:
+            try:
+                await self._polling_task
+            except asyncio.CancelledError:
+                pass
+            self._polling_task = None
+            logger.info("Job monitor polling stopped.")

researchloop/clusters/ssh.py

@@ -0,0 +1,183 @@
+"""SSH connection manager using asyncssh."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from pathlib import Path
+from typing import Any
+
+import asyncssh
+
+logger = logging.getLogger(__name__)
+
+
+class SSHConnection:
+    """Manages a single SSH connection to a remote host."""
+
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        user: str,
+        key_path: str,
+        known_hosts: str | None = None,
+    ) -> None:
+        self.host = host
+        self.port = port
+        self.user = user
+        self.key_path = key_path
+        self.known_hosts = known_hosts
+        self._conn: asyncssh.SSHClientConnection | None = None
+
+    async def connect(self) -> SSHConnection:
+        """Establish the SSH connection and return self."""
+        connect_kwargs: dict[str, Any] = {
+            "host": self.host,
+            "port": self.port,
+            "username": self.user,
+            "client_keys": [self.key_path],
+            "agent_path": None,  # Don't use SSH agent; we have explicit keys.
+        }
+        if self.known_hosts is not None:
+            connect_kwargs["known_hosts"] = self.known_hosts
+        else:
+            # Disable host key checking when no known_hosts file is provided.
+            connect_kwargs["known_hosts"] = None
+
+        logger.info("Connecting to %s@%s:%d", self.user, self.host, self.port)
+        self._conn = await asyncssh.connect(**connect_kwargs)
+        logger.info("Connected to %s@%s:%d", self.user, self.host, self.port)
+        return self
+
+    @property
+    def connection(self) -> asyncssh.SSHClientConnection:
+        if self._conn is None:
+            raise RuntimeError(
+                "SSH connection is not established. Call connect() first."
+            )
+        return self._conn
+
+    async def run(self, command: str, timeout: float = 30) -> tuple[str, str, int]:
+        """Run a command over SSH.
+
+        Returns:
+            A tuple of (stdout, stderr, exit_code).
+        """
+        logger.debug("Running command on %s: %s", self.host, command)
+        try:
+            result = await asyncio.wait_for(
+                self.connection.run(command, check=False),
+                timeout=timeout,
+            )
+        except asyncio.TimeoutError:
+            logger.error(
+                "Command timed out after %.1fs on %s: %s",
+                timeout,
+                self.host,
+                command,
+            )
+            raise
+
+        stdout = str(result.stdout or "")
+        stderr = str(result.stderr or "")
+        exit_code = result.exit_status if result.exit_status is not None else -1
+
+        logger.debug("Command on %s finished with exit_code=%d", self.host, exit_code)
+        return stdout, stderr, exit_code
+
+    async def upload_file(self, local_path: str, remote_path: str) -> None:
+        """Upload a local file to the remote host via SFTP."""
+        logger.info("Uploading %s -> %s:%s", local_path, self.host, remote_path)
+        async with self.connection.start_sftp_client() as sftp:
+            await sftp.put(local_path, remote_path)
+        logger.info("Upload complete: %s -> %s:%s", local_path, self.host, remote_path)
+
+    async def download_file(self, remote_path: str, local_path: str) -> None:
+        """Download a file from the remote host via SFTP."""
+        logger.info("Downloading %s:%s -> %s", self.host, remote_path, local_path)
+        # Ensure local parent directory exists.
+        Path(local_path).parent.mkdir(parents=True, exist_ok=True)
+        async with self.connection.start_sftp_client() as sftp:
+            await sftp.get(remote_path, local_path)
+        logger.info(
+            "Download complete: %s:%s -> %s", self.host, remote_path, local_path
+        )
+
+    async def close(self) -> None:
+        """Close the SSH connection."""
+        if self._conn is not None:
+            self._conn.close()
+            await self._conn.wait_closed()
+            logger.info(
+                "Closed connection to %s@%s:%d", self.user, self.host, self.port
+            )
+            self._conn = None
+
+    # --- Context manager support ---
+
+    async def __aenter__(self) -> SSHConnection:
+        return await self.connect()
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        await self.close()
+
+
+class SSHManager:
+    """Manages a pool of SSH connections keyed by cluster configuration."""
+
+    def __init__(self) -> None:
+        self._connections: dict[str, SSHConnection] = {}
+        self._lock = asyncio.Lock()
+
+    @staticmethod
+    def _config_key(cluster_config: dict[str, Any]) -> str:
+        """Derive a unique key from cluster configuration."""
+        host = cluster_config["host"]
+        port = cluster_config["port"]
+        user = cluster_config["user"]
+        return f"{user}@{host}:{port}"
+
+    async def get_connection(self, cluster_config: dict[str, Any]) -> SSHConnection:
+        """Return an existing connection or create a new one.
+
+        ``cluster_config`` must contain at minimum::
+
+            {
+                "host": str,
+                "port": int,
+                "user": str,
+                "key_path": str,
+                "known_hosts": str | None,  # optional
+            }
+        """
+        key = self._config_key(cluster_config)
+
+        async with self._lock:
+            existing = self._connections.get(key)
+            if existing is not None and existing._conn is not None:
+                logger.debug("Reusing existing SSH connection for %s", key)
+                return existing
+
+            # Create a fresh connection.
+            conn = SSHConnection(
+                host=cluster_config["host"],
+                port=cluster_config["port"],
+                user=cluster_config["user"],
+                key_path=cluster_config["key_path"],
+                known_hosts=cluster_config.get("known_hosts"),
+            )
+            await conn.connect()
+            self._connections[key] = conn
+            return conn
+
+    async def close_all(self) -> None:
+        """Close every managed SSH connection."""
+        async with self._lock:
+            for key, conn in self._connections.items():
+                try:
+                    await conn.close()
+                except Exception:
+                    logger.exception("Error closing SSH connection %s", key)
+            self._connections.clear()
+            logger.info("All SSH connections closed.")

researchloop/comms/base.py

@@ -0,0 +1,34 @@
+"""Abstract notification interface for researchloop."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+
+class BaseNotifier(ABC):
+    """Every notification backend must implement this interface."""
+
+    @abstractmethod
+    async def notify_sprint_started(
+        self, sprint_id: str, study_name: str, idea: str
+    ) -> None:
+        """Called when a sprint has been submitted to a cluster."""
+        ...
+
+    @abstractmethod
+    async def notify_sprint_completed(
+        self,
+        sprint_id: str,
+        study_name: str,
+        summary: str,
+        pdf_path: str | None = None,
+    ) -> None:
+        """Called when a sprint finishes successfully."""
+        ...
+
+    @abstractmethod
+    async def notify_sprint_failed(
+        self, sprint_id: str, study_name: str, error: str
+    ) -> None:
+        """Called when a sprint fails."""
+        ...