hpc-runner 0.1.0__py3-none-any.whl

hpc_runner/schedulers/local/scheduler.py

@@ -0,0 +1,237 @@
+"""Local scheduler - executes jobs as subprocesses."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import tempfile
+from datetime import datetime
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from hpc_runner.core.result import ArrayJobResult, JobResult, JobStatus
+from hpc_runner.schedulers.base import BaseScheduler
+from hpc_runner.templates import render_template
+
+if TYPE_CHECKING:
+    from hpc_runner.core.job import Job
+    from hpc_runner.core.job_array import JobArray
+
+
+class LocalScheduler(BaseScheduler):
+    """Execute jobs locally (for development/testing)."""
+
+    name = "local"
+
+    _job_counter: int = 0
+    _processes: dict[str, subprocess.Popen] = {}  # type: ignore[type-arg]
+    _exit_codes: dict[str, int] = {}
+    _output_paths: dict[str, dict[str, Path]] = {}
+
+    def submit(self, job: "Job", interactive: bool = False) -> JobResult:
+        """Run job as local subprocess."""
+        LocalScheduler._job_counter += 1
+        job_id = f"local_{LocalScheduler._job_counter}_{datetime.now().strftime('%Y%m%d%H%M%S')}"
+
+        # Set up environment with modules (modules not actually loaded locally)
+        env = os.environ.copy() if job.inherit_env else {}
+
+        # Generate and write script
+        script = self.generate_script(job)
+        script_path = Path(tempfile.gettempdir()) / f".hpc_local_{job_id}.sh"
+        script_path.write_text(script)
+        script_path.chmod(0o755)
+
+        workdir = Path(job.workdir) if job.workdir else Path.cwd()
+
+        # Determine output paths
+        stdout_file = job.stdout or f"{job.name}.{job_id}.out"
+        stdout_path = workdir / stdout_file
+        if job.merge_output:
+            stderr_path = stdout_path  # Merge stderr into stdout
+        else:
+            stderr_file = job.stderr or f"{job.name}.{job_id}.err"
+            stderr_path = workdir / stderr_file
+
+        # Store output paths
+        LocalScheduler._output_paths[job_id] = {
+            "stdout": stdout_path,
+            "stderr": stderr_path,
+        }
+
+        if interactive:
+            # Blocking execution
+            with open(stdout_path, "w") as stdout_f:
+                if job.merge_output:
+                    result = subprocess.run(
+                        [str(script_path)],
+                        cwd=workdir,
+                        env=env,
+                        stdout=stdout_f,
+                        stderr=subprocess.STDOUT,
+                    )
+                else:
+                    with open(stderr_path, "w") as stderr_f:
+                        result = subprocess.run(
+                            [str(script_path)],
+                            cwd=workdir,
+                            env=env,
+                            stdout=stdout_f,
+                            stderr=stderr_f,
+                        )
+            LocalScheduler._exit_codes[job_id] = result.returncode
+            script_path.unlink(missing_ok=True)
+        else:
+            # Background execution
+            stdout_f = open(stdout_path, "w")
+            if job.merge_output:
+                proc = subprocess.Popen(
+                    [str(script_path)],
+                    cwd=workdir,
+                    env=env,
+                    stdout=stdout_f,
+                    stderr=subprocess.STDOUT,
+                )
+            else:
+                stderr_f = open(stderr_path, "w")
+                proc = subprocess.Popen(
+                    [str(script_path)],
+                    cwd=workdir,
+                    env=env,
+                    stdout=stdout_f,
+                    stderr=stderr_f,
+                )
+            LocalScheduler._processes[job_id] = proc
+            # Store script path for cleanup
+            proc._script_path = script_path  # type: ignore[attr-defined]
+            proc._stdout_file = stdout_f  # type: ignore[attr-defined]
+            if not job.merge_output:
+                proc._stderr_file = stderr_f  # type: ignore[attr-defined]
+
+        return JobResult(job_id=job_id, scheduler=self, job=job)
+
+    def submit_array(self, array: "JobArray") -> ArrayJobResult:
+        """Simulate array job by submitting multiple jobs."""
+        # For local scheduler, we just run one job
+        # and return an ArrayJobResult pointing to it
+        LocalScheduler._job_counter += 1
+        base_job_id = f"local_array_{LocalScheduler._job_counter}"
+
+        # Run jobs sequentially (or could be parallel)
+        for idx in array.indices:
+            # Set array index environment variable
+            os.environ["HPC_ARRAY_TASK_ID"] = str(idx)
+            os.environ["SGE_TASK_ID"] = str(idx)  # SGE compat
+            os.environ["SLURM_ARRAY_TASK_ID"] = str(idx)  # Slurm compat
+
+            # Create a job ID for this task
+            task_job_id = f"{base_job_id}.{idx}"
+            self._submit_array_task(array.job, task_job_id, idx)
+
+        return ArrayJobResult(base_job_id=base_job_id, scheduler=self, array=array)
+
+    def _submit_array_task(self, job: "Job", job_id: str, index: int) -> None:
+        """Submit a single array task."""
+        env = os.environ.copy() if job.inherit_env else {}
+        env["HPC_ARRAY_TASK_ID"] = str(index)
+
+        script = self.generate_script(job)
+        script_path = Path(tempfile.gettempdir()) / f".hpc_local_{job_id}.sh"
+        script_path.write_text(script)
+        script_path.chmod(0o755)
+
+        workdir = Path(job.workdir) if job.workdir else Path.cwd()
+        stdout_path = workdir / f"{job.name}.{job_id}.out"
+
+        LocalScheduler._output_paths[job_id] = {"stdout": stdout_path, "stderr": stdout_path}
+
+        stdout_f = open(stdout_path, "w")
+        proc = subprocess.Popen(
+            [str(script_path)],
+            cwd=workdir,
+            env=env,
+            stdout=stdout_f,
+            stderr=subprocess.STDOUT,
+        )
+        LocalScheduler._processes[job_id] = proc
+        proc._script_path = script_path  # type: ignore[attr-defined]
+        proc._stdout_file = stdout_f  # type: ignore[attr-defined]
+
+    def cancel(self, job_id: str) -> bool:
+        """Cancel a local job."""
+        if job_id in LocalScheduler._processes:
+            proc = LocalScheduler._processes[job_id]
+            proc.terminate()
+            proc.wait()
+            self._cleanup_process(job_id)
+            return True
+        return False
+
+    def get_status(self, job_id: str) -> JobStatus:
+        """Get job status."""
+        if job_id in LocalScheduler._exit_codes:
+            # Already completed
+            return JobStatus.COMPLETED if LocalScheduler._exit_codes[job_id] == 0 else JobStatus.FAILED
+
+        if job_id not in LocalScheduler._processes:
+            return JobStatus.UNKNOWN
+
+        proc = LocalScheduler._processes[job_id]
+        poll = proc.poll()
+
+        if poll is None:
+            return JobStatus.RUNNING
+
+        # Process completed
+        LocalScheduler._exit_codes[job_id] = poll
+        self._cleanup_process(job_id)
+
+        return JobStatus.COMPLETED if poll == 0 else JobStatus.FAILED
+
+    def _cleanup_process(self, job_id: str) -> None:
+        """Clean up process resources."""
+        if job_id in LocalScheduler._processes:
+            proc = LocalScheduler._processes[job_id]
+            # Close file handles
+            if hasattr(proc, "_stdout_file"):
+                proc._stdout_file.close()  # type: ignore[attr-defined]
+            if hasattr(proc, "_stderr_file"):
+                proc._stderr_file.close()  # type: ignore[attr-defined]
+            # Remove script
+            if hasattr(proc, "_script_path"):
+                proc._script_path.unlink(missing_ok=True)  # type: ignore[attr-defined]
+            del LocalScheduler._processes[job_id]
+
+    def get_exit_code(self, job_id: str) -> int | None:
+        """Get exit code."""
+        # First check if we have a cached exit code
+        if job_id in LocalScheduler._exit_codes:
+            return LocalScheduler._exit_codes[job_id]
+
+        # Check if process is done
+        if job_id in LocalScheduler._processes:
+            proc = LocalScheduler._processes[job_id]
+            poll = proc.poll()
+            if poll is not None:
+                LocalScheduler._exit_codes[job_id] = poll
+                return poll
+
+        return None
+
+    def get_output_path(self, job_id: str, stream: str) -> Path | None:
+        """Get output file path."""
+        if job_id in LocalScheduler._output_paths:
+            return LocalScheduler._output_paths[job_id].get(stream)
+        return None
+
+    def generate_script(self, job: "Job") -> str:
+        """Generate local execution script."""
+        return render_template(
+            "local/templates/job.sh.j2",
+            job=job,
+            scheduler=self,
+        )
+
+    def build_submit_command(self, job: "Job") -> list[str]:
+        """Build command - for local, just bash."""
+        return ["bash", "-c", job.command if isinstance(job.command, str) else " ".join(job.command)]

hpc_runner/schedulers/local/templates/job.sh.j2

@@ -0,0 +1,28 @@
+#!/bin/bash
+# Generated by hpc-tools (local scheduler)
+
+# Exit on error
+set -e
+
+{% if job.modules_path %}
+# Additional module paths (simulated for local)
+{% for path in job.modules_path %}
+# module use {{ path }}
+{% endfor %}
+{% endif %}
+
+{% if job.modules %}
+# Modules (simulated for local - not actually loaded)
+{% for mod in job.modules %}
+# module load {{ mod }}
+{% endfor %}
+{% endif %}
+
+{% if job.workdir %}
+# Change to working directory
+cd {{ job.workdir }}
+{% endif %}
+
+# Execute command
+{{ job.command }}
+exit $?

hpc_runner/schedulers/sge/__init__.py

@@ -0,0 +1,5 @@
+"""SGE (Sun Grid Engine) scheduler implementation."""
+
+from hpc_runner.schedulers.sge.scheduler import SGEScheduler
+
+__all__ = ["SGEScheduler"]

hpc_runner/schedulers/sge/args.py

@@ -0,0 +1,165 @@
+"""SGE-specific argument descriptors."""
+
+from hpc_runner.core.descriptors import SchedulerArg
+
+
+class SGEArg(SchedulerArg):
+    """Base SGE argument.
+
+    SGE uses #$ -flag value format for directives.
+    """
+
+    def to_args(self, value) -> list[str]:
+        if value is None:
+            return []
+        return [f"-{self.flag}", str(self.converter(value))]
+
+    def to_directive(self, value) -> str | None:
+        if value is None:
+            return None
+        return f"#$ -{self.flag} {self.converter(value)}"
+
+
+class SGECpuArg(SGEArg):
+    """CPU/slots argument using parallel environment.
+
+    Note: The PE name is configurable via config.
+    """
+
+    def __init__(self, pe_name: str = "smp"):
+        super().__init__("pe", converter=lambda v: f"{pe_name} {v}", doc="Parallel environment")
+        self.pe_name = pe_name
+
+    def to_args(self, value, pe_name: str | None = None) -> list[str]:
+        if value is None:
+            return []
+        pe = pe_name or self.pe_name
+        return ["-pe", f"{pe} {value}"]
+
+    def to_directive(self, value, pe_name: str | None = None) -> str | None:
+        if value is None:
+            return None
+        pe = pe_name or self.pe_name
+        return f"#$ -pe {pe} {value}"
+
+
+class SGEMemArg(SGEArg):
+    """Memory argument.
+
+    Uses -l resource=value format. Resource name is configurable.
+    """
+
+    def __init__(self, resource_name: str = "mem_free"):
+        super().__init__("l", doc="Memory requirement")
+        self.resource_name = resource_name
+
+    def to_args(self, value, resource_name: str | None = None) -> list[str]:
+        if value is None:
+            return []
+        res = resource_name or self.resource_name
+        return ["-l", f"{res}={value}"]
+
+    def to_directive(self, value, resource_name: str | None = None) -> str | None:
+        if value is None:
+            return None
+        res = resource_name or self.resource_name
+        return f"#$ -l {res}={value}"
+
+
+class SGETimeArg(SGEArg):
+    """Time limit argument.
+
+    Uses -l h_rt=HH:MM:SS format. Resource name is configurable.
+    """
+
+    def __init__(self, resource_name: str = "h_rt"):
+        super().__init__("l", doc="Hard runtime limit")
+        self.resource_name = resource_name
+
+    def to_args(self, value, resource_name: str | None = None) -> list[str]:
+        if value is None:
+            return []
+        res = resource_name or self.resource_name
+        return ["-l", f"{res}={value}"]
+
+    def to_directive(self, value, resource_name: str | None = None) -> str | None:
+        if value is None:
+            return None
+        res = resource_name or self.resource_name
+        return f"#$ -l {res}={value}"
+
+
+class SGEQueueArg(SGEArg):
+    """Queue argument."""
+
+    def __init__(self):
+        super().__init__("q", doc="Queue name")
+
+
+class SGEJobNameArg(SGEArg):
+    """Job name argument."""
+
+    def __init__(self):
+        super().__init__("N", doc="Job name")
+
+
+class SGEOutputArg(SGEArg):
+    """Stdout path argument."""
+
+    def __init__(self):
+        super().__init__("o", doc="Stdout file path")
+
+
+class SGEErrorArg(SGEArg):
+    """Stderr path argument."""
+
+    def __init__(self):
+        super().__init__("e", doc="Stderr file path")
+
+
+class SGEArrayArg(SGEArg):
+    """Array job argument."""
+
+    def __init__(self):
+        super().__init__("t", doc="Array job range (e.g., 1-100, 1-100:10)")
+
+
+class SGEJoinOutputArg(SGEArg):
+    """Join stdout and stderr."""
+
+    def __init__(self):
+        super().__init__("j", doc="Join stdout and stderr")
+
+    def to_args(self, value) -> list[str]:
+        if value:
+            return ["-j", "y"]
+        return []
+
+    def to_directive(self, value) -> str | None:
+        if value:
+            return "#$ -j y"
+        return None
+
+
+class SGECwdArg(SGEArg):
+    """Use current working directory."""
+
+    def __init__(self):
+        super().__init__("cwd", doc="Use current working directory")
+
+    def to_args(self, value) -> list[str]:
+        if value:
+            return ["-cwd"]
+        return []
+
+    def to_directive(self, value) -> str | None:
+        if value:
+            return "#$ -cwd"
+        return None
+
+
+class SGEShellArg(SGEArg):
+    """Shell to use for the job."""
+
+    def __init__(self):
+        super().__init__("S", doc="Shell path")

hpc_runner/schedulers/sge/parser.py

@@ -0,0 +1,194 @@
+"""SGE output parsing utilities."""
+
+import re
+import xml.etree.ElementTree as ET
+from typing import Any
+
+from hpc_runner.core.result import JobStatus
+
+
+def parse_qstat_xml(xml_output: str) -> dict[str, Any]:
+    """Parse qstat -xml output.
+
+    Returns dict with job_id -> job_info mappings.
+    """
+    jobs: dict[str, Any] = {}
+
+    try:
+        root = ET.fromstring(xml_output)
+
+        # Parse queue_info (running jobs)
+        for job_list in root.findall(".//job_list"):
+            job_info = _parse_job_element(job_list)
+            if job_info:
+                jobs[job_info["job_id"]] = job_info
+
+        # Parse job_info (pending jobs)
+        for job_list in root.findall(".//job_info/job_list"):
+            job_info = _parse_job_element(job_list)
+            if job_info:
+                jobs[job_info["job_id"]] = job_info
+
+    except ET.ParseError:
+        pass
+
+    return jobs
+
+
+def _parse_job_element(elem: ET.Element) -> dict[str, Any] | None:
+    """Parse a single job_list element."""
+    job_id_elem = elem.find("JB_job_number")
+    if job_id_elem is None or job_id_elem.text is None:
+        return None
+
+    job_info: dict[str, Any] = {
+        "job_id": job_id_elem.text,
+    }
+
+    # Job name
+    name_elem = elem.find("JB_name")
+    if name_elem is not None and name_elem.text:
+        job_info["name"] = name_elem.text
+
+    # State
+    state_elem = elem.find("state")
+    if state_elem is not None and state_elem.text:
+        job_info["state"] = state_elem.text
+
+    # Queue
+    queue_elem = elem.find("queue_name")
+    if queue_elem is not None and queue_elem.text:
+        job_info["queue"] = queue_elem.text
+
+    # Slots
+    slots_elem = elem.find("slots")
+    if slots_elem is not None and slots_elem.text:
+        job_info["slots"] = int(slots_elem.text)
+
+    return job_info
+
+
+def parse_qstat_plain(output: str) -> dict[str, Any]:
+    """Parse plain qstat output.
+
+    Format:
+    job-ID  prior   name       user         state submit/start at     queue                          slots ja-task-ID
+    --------------------------------------------------------------------------------
+    12345   0.55500 myjob      user         r     01/01/2024 10:00:00 all.q@node1                    1
+    """
+    jobs: dict[str, Any] = {}
+
+    lines = output.strip().split("\n")
+
+    # Skip header lines
+    data_started = False
+    for line in lines:
+        if line.startswith("-"):
+            data_started = True
+            continue
+        if not data_started:
+            continue
+
+        parts = line.split()
+        if len(parts) >= 5:
+            job_id = parts[0]
+            jobs[job_id] = {
+                "job_id": job_id,
+                "priority": parts[1],
+                "name": parts[2],
+                "user": parts[3],
+                "state": parts[4],
+            }
+
+            # Parse queue if present
+            if len(parts) >= 8:
+                jobs[job_id]["queue"] = parts[7]
+
+            # Parse slots if present
+            if len(parts) >= 9:
+                try:
+                    jobs[job_id]["slots"] = int(parts[8])
+                except ValueError:
+                    pass
+
+    return jobs
+
+
+def parse_qacct_output(output: str) -> dict[str, Any]:
+    """Parse qacct output for job accounting info.
+
+    Format:
+    ==============================================================
+    qname        all.q
+    hostname     node1
+    group        users
+    owner        user
+    jobname      myjob
+    jobnumber    12345
+    ...
+    exit_status  0
+    """
+    info: dict[str, Any] = {}
+
+    for line in output.strip().split("\n"):
+        if line.startswith("="):
+            continue
+
+        parts = line.split(None, 1)
+        if len(parts) == 2:
+            key, value = parts
+            info[key] = value.strip()
+
+    return info
+
+
+def state_to_status(state: str) -> JobStatus:
+    """Convert SGE state code to JobStatus.
+
+    SGE states:
+    - qw: pending (queued, waiting)
+    - hqw: hold (on hold)
+    - r: running
+    - t: transferring
+    - Rr, Rt: restarted
+    - s, ts: suspended
+    - S, tS: queue suspended
+    - T, tT: threshold
+    - Eqw: error (waiting)
+    - dr: deleting (running)
+    - dt: deleting (transferring)
+    """
+    state = state.lower()
+
+    if state in ("r", "t", "rr", "rt"):
+        return JobStatus.RUNNING
+    elif state in ("qw", "hqw"):
+        return JobStatus.PENDING
+    elif state in ("eqw",):
+        return JobStatus.FAILED
+    elif state in ("dr", "dt"):
+        return JobStatus.CANCELLED
+    elif state in ("s", "ts", "ss", "ts"):
+        return JobStatus.PENDING  # Suspended, treat as pending
+
+    return JobStatus.UNKNOWN
+
+
+def parse_qsub_output(output: str) -> str | None:
+    """Parse qsub output to extract job ID.
+
+    Expected format:
+    Your job 12345 ("jobname") has been submitted
+    Your job-array 12345.1-10:1 ("jobname") has been submitted
+    """
+    # Standard job
+    match = re.search(r"Your job (\d+)", output)
+    if match:
+        return match.group(1)
+
+    # Array job
+    match = re.search(r"Your job-array (\d+)", output)
+    if match:
+        return match.group(1)
+
+    return None