researchloop 0.1.0__py3-none-any.whl

researchloop/runner/upload.py

@@ -0,0 +1,153 @@
+"""Artifact upload and orchestrator webhooks.
+
+Uses ``httpx`` for HTTP calls since this runs on HPC nodes where
+``aiohttp`` may not be available.
+"""
+
+from __future__ import annotations
+
+import logging
+import mimetypes
+from pathlib import Path
+
+import httpx
+
+logger = logging.getLogger(__name__)
+
+# Files we always look for in the sprint directory.
+_KEY_FILENAMES = {"report.md", "summary.txt"}
+
+# Extensions we scan for as additional artifacts.
+_ARTIFACT_EXTENSIONS = {".pdf", ".png", ".jpg", ".jpeg"}
+
+# Generous timeout for large uploads from HPC → orchestrator.
+_UPLOAD_TIMEOUT = httpx.Timeout(connect=10.0, read=60.0, write=120.0, pool=10.0)
+_WEBHOOK_TIMEOUT = httpx.Timeout(connect=10.0, read=30.0, write=30.0, pool=10.0)
+
+
+def _auth_headers(shared_secret: str) -> dict[str, str]:
+    return {"Authorization": f"Bearer {shared_secret}"}
+
+
+async def upload_artifacts(
+    sprint_dir: str,
+    orchestrator_url: str,
+    shared_secret: str,
+    sprint_id: str,
+) -> list[str]:
+    """Scan *sprint_dir* for key artifacts and upload them to the orchestrator.
+
+    Returns a list of successfully uploaded filenames.
+    """
+    sprint_path = Path(sprint_dir)
+    files_to_upload: list[Path] = []
+
+    # Collect known filenames.
+    for name in _KEY_FILENAMES:
+        candidate = sprint_path / name
+        if candidate.is_file():
+            files_to_upload.append(candidate)
+
+    # Scan for image / PDF artifacts anywhere in the sprint directory (1 level deep).
+    for child in sprint_path.iterdir():
+        if child.is_file() and child.suffix.lower() in _ARTIFACT_EXTENSIONS:
+            if child not in files_to_upload:
+                files_to_upload.append(child)
+
+    # Also check results/ subdirectory.
+    results_dir = sprint_path / "results"
+    if results_dir.is_dir():
+        for child in results_dir.iterdir():
+            if child.is_file() and child.suffix.lower() in _ARTIFACT_EXTENSIONS:
+                files_to_upload.append(child)
+
+    if not files_to_upload:
+        logger.info("No artifacts found to upload in %s", sprint_dir)
+        return []
+
+    url = f"{orchestrator_url.rstrip('/')}/api/artifacts/{sprint_id}"
+    uploaded: list[str] = []
+
+    async with httpx.AsyncClient(timeout=_UPLOAD_TIMEOUT) as client:
+        for file_path in files_to_upload:
+            mime_type = (
+                mimetypes.guess_type(str(file_path))[0] or "application/octet-stream"
+            )
+            try:
+                with open(file_path, "rb") as f:
+                    resp = await client.post(
+                        url,
+                        headers=_auth_headers(shared_secret),
+                        files={"file": (file_path.name, f, mime_type)},
+                    )
+                resp.raise_for_status()
+                uploaded.append(file_path.name)
+                logger.info("Uploaded artifact: %s", file_path.name)
+            except httpx.HTTPError:
+                logger.exception("Failed to upload artifact: %s", file_path.name)
+
+    return uploaded
+
+
+async def send_webhook(
+    orchestrator_url: str,
+    shared_secret: str,
+    sprint_id: str,
+    status: str,
+    summary: str | None = None,
+    error: str | None = None,
+) -> None:
+    """Notify the orchestrator that a sprint has completed (or failed).
+
+    POSTs to ``/api/webhook/sprint-complete``.
+    """
+    url = f"{orchestrator_url.rstrip('/')}/api/webhook/sprint-complete"
+    payload = {
+        "sprint_id": sprint_id,
+        "status": status,
+        "summary": summary,
+        "error": error,
+    }
+
+    async with httpx.AsyncClient(timeout=_WEBHOOK_TIMEOUT) as client:
+        resp = await client.post(
+            url,
+            headers={
+                **_auth_headers(shared_secret),
+                "Content-Type": "application/json",
+            },
+            json=payload,
+        )
+        resp.raise_for_status()
+
+    logger.info("Sent sprint-complete webhook: sprint=%s status=%s", sprint_id, status)
+
+
+async def send_heartbeat(
+    orchestrator_url: str,
+    shared_secret: str,
+    sprint_id: str,
+    status: str,
+    step: int,
+) -> None:
+    """Send a heartbeat ping to the orchestrator.
+
+    POSTs to ``/api/webhook/heartbeat``.
+    """
+    url = f"{orchestrator_url.rstrip('/')}/api/webhook/heartbeat"
+    payload = {
+        "sprint_id": sprint_id,
+        "status": status,
+        "step": step,
+    }
+
+    async with httpx.AsyncClient(timeout=_WEBHOOK_TIMEOUT) as client:
+        resp = await client.post(
+            url,
+            headers={
+                **_auth_headers(shared_secret),
+                "Content-Type": "application/json",
+            },
+            json=payload,
+        )
+        resp.raise_for_status()

researchloop/schedulers/__init__.py

@@ -0,0 +1,11 @@
+from researchloop.schedulers.base import BaseScheduler
+from researchloop.schedulers.local import LocalScheduler
+from researchloop.schedulers.sge import SGEScheduler
+from researchloop.schedulers.slurm import SlurmScheduler
+
+__all__ = [
+    "BaseScheduler",
+    "SGEScheduler",
+    "SlurmScheduler",
+    "LocalScheduler",
+]

researchloop/schedulers/base.py

@@ -0,0 +1,43 @@
+"""Abstract base class for job schedulers."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+
+class BaseScheduler(ABC):
+    """Interface that every scheduler backend must implement."""
+
+    @abstractmethod
+    async def submit(
+        self,
+        ssh: object,
+        script: str,
+        job_name: str,
+        working_dir: str,
+        env: dict[str, str] | None = None,
+    ) -> str:
+        """Submit a job and return the scheduler-assigned job ID."""
+
+    @abstractmethod
+    async def status(self, ssh: object, job_id: str) -> str:
+        """Return the current status of a job.
+
+        Must return one of:
+        ``"pending"``, ``"running"``, ``"completed"``, ``"failed"``, ``"unknown"``.
+        """
+
+    @abstractmethod
+    async def cancel(self, ssh: object, job_id: str) -> bool:
+        """Cancel a job. Return ``True`` on success."""
+
+    @abstractmethod
+    def generate_script(
+        self,
+        command: str,
+        job_name: str,
+        working_dir: str,
+        time_limit: str = "8:00:00",
+        env: dict[str, str] | None = None,
+    ) -> str:
+        """Generate the contents of a submission script."""

researchloop/schedulers/local.py

@@ -0,0 +1,188 @@
+"""Local subprocess scheduler for testing without a real cluster."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import signal
+import textwrap
+
+from researchloop.schedulers.base import BaseScheduler
+
+logger = logging.getLogger(__name__)
+
+
+class LocalScheduler(BaseScheduler):
+    """Scheduler that runs commands as local subprocesses.
+
+    Intended for development and testing.  Job IDs are OS process IDs (PIDs)
+    stored as strings for consistency with the ``BaseScheduler`` interface.
+    The *ssh* parameter is accepted but ignored -- all commands run locally.
+    """
+
+    def __init__(self) -> None:
+        # Track running processes keyed by PID string.
+        self._processes: dict[str, asyncio.subprocess.Process] = {}
+
+    # ------------------------------------------------------------------
+    # submit
+    # ------------------------------------------------------------------
+
+    async def submit(
+        self,
+        ssh: object,
+        script: str,
+        job_name: str,
+        working_dir: str,
+        env: dict[str, str] | None = None,
+    ) -> str:
+        """Run *script* as a local background subprocess.
+
+        Returns the PID of the spawned process as a string.
+        """
+        # Build the environment for the subprocess.
+        proc_env = os.environ.copy()
+        if env:
+            proc_env.update(env)
+
+        # Ensure the working directory exists.
+        os.makedirs(working_dir, exist_ok=True)
+
+        # Write the script to a temporary file in working_dir so it can
+        # be executed as a proper bash script.
+        script_path = os.path.join(working_dir, f".researchloop_local_{job_name}.sh")
+        with open(script_path, "w") as f:
+            f.write(script)
+        os.chmod(script_path, 0o755)
+
+        stdout_path = os.path.join(working_dir, f"{job_name}.out")
+        stderr_path = os.path.join(working_dir, f"{job_name}.err")
+
+        stdout_file = open(stdout_path, "w")
+        stderr_file = open(stderr_path, "w")
+
+        proc = await asyncio.create_subprocess_exec(
+            "bash",
+            script_path,
+            cwd=working_dir,
+            env=proc_env,
+            stdout=stdout_file,
+            stderr=stderr_file,
+            start_new_session=True,
+        )
+
+        stdout_file.close()
+        stderr_file.close()
+
+        job_id = str(proc.pid)
+        self._processes[job_id] = proc
+
+        logger.info(
+            "Local job submitted: pid=%s name=%s dir=%s", job_id, job_name, working_dir
+        )
+        return job_id
+
+    # ------------------------------------------------------------------
+    # status
+    # ------------------------------------------------------------------
+
+    async def status(self, ssh: object, job_id: str) -> str:
+        """Check whether the process identified by *job_id* (PID) is alive."""
+        proc = self._processes.get(job_id)
+
+        if proc is not None:
+            # We have a handle -- check if it has terminated.
+            if proc.returncode is None:
+                # Process still running (or hasn't been awaited yet).
+                # Use a non-blocking poll.
+                try:
+                    # wait with timeout=0 raises TimeoutError if still running.
+                    await asyncio.wait_for(asyncio.shield(proc.wait()), timeout=0.0)
+                except (asyncio.TimeoutError, asyncio.CancelledError):
+                    return "running"
+
+            # Process has finished.
+            if proc.returncode == 0:
+                return "completed"
+            return "failed"
+
+        # No tracked handle -- try the OS.
+        try:
+            pid = int(job_id)
+            os.kill(pid, 0)  # Signal 0: check existence without killing.
+            return "running"
+        except (OSError, ValueError):
+            return "unknown"
+
+    # ------------------------------------------------------------------
+    # cancel
+    # ------------------------------------------------------------------
+
+    async def cancel(self, ssh: object, job_id: str) -> bool:
+        """Kill the local process identified by *job_id* (PID)."""
+        proc = self._processes.get(job_id)
+
+        if proc is not None:
+            try:
+                proc.terminate()
+                # Give the process a moment to exit gracefully.
+                try:
+                    await asyncio.wait_for(proc.wait(), timeout=5.0)
+                except asyncio.TimeoutError:
+                    proc.kill()
+                logger.info("Cancelled local job pid=%s", job_id)
+                return True
+            except ProcessLookupError:
+                logger.warning("Process pid=%s already exited", job_id)
+                return True
+
+        # Fall back to OS signal.
+        try:
+            pid = int(job_id)
+            os.kill(pid, signal.SIGTERM)
+            logger.info("Sent SIGTERM to pid=%s", job_id)
+            return True
+        except (OSError, ValueError) as exc:
+            logger.error("Failed to cancel local job pid=%s: %s", job_id, exc)
+            return False
+
+    # ------------------------------------------------------------------
+    # generate_script
+    # ------------------------------------------------------------------
+
+    def generate_script(
+        self,
+        command: str,
+        job_name: str,
+        working_dir: str,
+        time_limit: str = "8:00:00",
+        env: dict[str, str] | None = None,
+    ) -> str:
+        """Generate a simple bash script for local execution."""
+        env_exports = ""
+        if env:
+            lines = [
+                f"export {key}={_shell_quote(value)}" for key, value in env.items()
+            ]
+            env_exports = "\n".join(lines) + "\n"
+
+        script = textwrap.dedent(f"""\
+            #!/usr/bin/env bash
+            set -euo pipefail
+
+            echo "=== Local job '{job_name}' started at $(date -u) ==="
+            echo "Host: $(hostname)"
+            echo "Working directory: $(pwd)"
+
+            {env_exports}# --- Run the command ---
+            {command}
+
+            echo "=== Local job '{job_name}' finished at $(date -u) ==="
+        """)
+        return script
+
+
+def _shell_quote(value: str) -> str:
+    """Wrap *value* in single quotes, escaping embedded single quotes."""
+    return "'" + value.replace("'", "'\\''") + "'"

researchloop/schedulers/sge.py

@@ -0,0 +1,163 @@
+"""Sun Grid Engine (SGE) / Open Grid Scheduler implementation."""
+
+from __future__ import annotations
+
+import logging
+import re
+import textwrap
+
+from researchloop.schedulers.base import BaseScheduler
+
+logger = logging.getLogger(__name__)
+
+# Maps SGE state codes to normalised status strings.
+_SGE_STATE_MAP: dict[str, str] = {
+    "qw": "pending",  # queued waiting
+    "hqw": "pending",  # hold queued waiting
+    "r": "running",
+    "t": "running",  # transferring
+    "s": "running",  # suspended
+    "S": "running",  # suspended by queue
+    "T": "running",  # threshold
+    "Eqw": "failed",  # error
+    "dr": "failed",  # deleting/running
+    "dt": "failed",  # deleting/transferring
+}
+
+
+class SGEScheduler(BaseScheduler):
+    """Job scheduler for SGE/Grid Engine clusters."""
+
+    async def submit(
+        self,
+        ssh: object,
+        script: str,
+        job_name: str,
+        working_dir: str,
+        env: dict[str, str] | None = None,
+    ) -> str:
+        # Submit via qsub — script is a remote file path
+        submit_cmd = f"cd {working_dir} && qsub {script}"
+        stdout, stderr, rc = await ssh.run(  # type: ignore[attr-defined]
+            submit_cmd, timeout=60
+        )
+        if rc != 0:
+            raise RuntimeError(f"qsub failed (exit {rc}): {stderr}")
+
+        # Parse job ID from
+        # "Your job XXXXX ("name") has been submitted"
+        match = re.search(r"Your job\s+(\d+)", stdout)
+        if not match:
+            raise RuntimeError(f"Could not parse job ID from qsub: {stdout!r}")
+
+        job_id = match.group(1)
+        logger.info(
+            "Submitted SGE job %s (name=%s)",
+            job_id,
+            job_name,
+        )
+
+        return job_id
+
+    async def status(self, ssh: object, job_id: str) -> str:
+        # Try qstat first
+        stdout, stderr, rc = await ssh.run(  # type: ignore[attr-defined]
+            f"qstat -j {job_id}", timeout=30
+        )
+        if rc == 0 and stdout.strip():
+            # Parse state from qstat output
+            for line in stdout.splitlines():
+                if line.strip().startswith("job_state"):
+                    # Format: job_state   1:  r
+                    parts = line.split(":")
+                    if len(parts) >= 2:
+                        state = parts[-1].strip()
+                        return _SGE_STATE_MAP.get(state, "unknown")
+            # Job exists but couldn't parse state
+            return "running"
+
+        # Try qstat listing format
+        stdout, stderr, rc = await ssh.run(  # type: ignore[attr-defined]
+            f"qstat | grep {job_id}", timeout=30
+        )
+        if rc == 0 and stdout.strip():
+            parts = stdout.split()
+            if len(parts) >= 5:
+                state = parts[4]
+                return _SGE_STATE_MAP.get(state, "unknown")
+
+        # Job not in queue -- try qacct for finished jobs
+        stdout, stderr, rc = await ssh.run(  # type: ignore[attr-defined]
+            f"qacct -j {job_id}", timeout=30
+        )
+        if rc == 0 and stdout.strip():
+            for line in stdout.splitlines():
+                if "exit_status" in line:
+                    parts = line.split()
+                    if len(parts) >= 2:
+                        code = parts[-1].strip()
+                        if code == "0":
+                            return "completed"
+                        return "failed"
+            return "completed"
+
+        logger.warning(
+            "Could not determine status for SGE job %s",
+            job_id,
+        )
+        return "unknown"
+
+    async def cancel(self, ssh: object, job_id: str) -> bool:
+        stdout, stderr, rc = await ssh.run(  # type: ignore[attr-defined]
+            f"qdel {job_id}", timeout=30
+        )
+        if rc == 0:
+            logger.info("Cancelled SGE job %s", job_id)
+            return True
+        logger.error(
+            "qdel failed for job %s (exit %d): %s",
+            job_id,
+            rc,
+            stderr,
+        )
+        return False
+
+    def generate_script(
+        self,
+        command: str,
+        job_name: str,
+        working_dir: str,
+        time_limit: str = "8:00:00",
+        env: dict[str, str] | None = None,
+    ) -> str:
+        env_exports = ""
+        if env:
+            lines = [f"export {k}={_shell_quote(v)}" for k, v in env.items()]
+            env_exports = "\n".join(lines) + "\n"
+
+        script = textwrap.dedent(f"""\
+            #!/usr/bin/env bash
+            #$ -N {job_name}
+            #$ -o {working_dir}/{job_name}_$JOB_ID.out
+            #$ -e {working_dir}/{job_name}_$JOB_ID.err
+            #$ -l h_rt={time_limit}
+            #$ -cwd
+            #$ -S /bin/bash
+
+            set -euo pipefail
+
+            echo "=== SGE Job $JOB_ID started ==="
+            echo "Host: $(hostname)"
+
+            {env_exports}# --- Run the command ---
+            {command}
+
+            echo "=== SGE Job $JOB_ID finished ==="
+        """)
+        return script
+
+
+def _shell_quote(value: str) -> str:
+    """Wrap *value* in single quotes, escaping embedded
+    single quotes."""
+    return "'" + value.replace("'", "'\\''") + "'"