hpc-runner 0.1.0__py3-none-any.whl

hpc_runner/schedulers/sge/scheduler.py

@@ -0,0 +1,325 @@
+"""SGE scheduler implementation."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import tempfile
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from hpc_runner.core.config import get_config
+from hpc_runner.core.result import ArrayJobResult, JobResult, JobStatus
+from hpc_runner.schedulers.base import BaseScheduler
+from hpc_runner.schedulers.sge.args import (
+    SGECpuArg,
+    SGECwdArg,
+    SGEErrorArg,
+    SGEJobNameArg,
+    SGEJoinOutputArg,
+    SGEMemArg,
+    SGEOutputArg,
+    SGEQueueArg,
+    SGETimeArg,
+)
+from hpc_runner.schedulers.sge.parser import (
+    parse_qacct_output,
+    parse_qstat_plain,
+    parse_qsub_output,
+    state_to_status,
+)
+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 SGEScheduler(BaseScheduler):
+    """Sun Grid Engine scheduler implementation."""
+
+    name = "sge"
+
+    # Descriptor-based argument definitions
+    cpu_arg = SGECpuArg()
+    mem_arg = SGEMemArg()
+    time_arg = SGETimeArg()
+    queue_arg = SGEQueueArg()
+    job_name_arg = SGEJobNameArg()
+    stdout_arg = SGEOutputArg()
+    stderr_arg = SGEErrorArg()
+    join_output_arg = SGEJoinOutputArg()
+    cwd_arg = SGECwdArg()
+
+    def __init__(self) -> None:
+        # Load scheduler-specific config
+        config = get_config()
+        sge_config = config.get_scheduler_config("sge")
+
+        self.pe_name = sge_config.get("parallel_environment", "smp")
+        self.mem_resource = sge_config.get("memory_resource", "mem_free")
+        self.time_resource = sge_config.get("time_resource", "h_rt")
+        self.merge_output_default = sge_config.get("merge_output", True)
+
+    def submit(self, job: "Job", interactive: bool = False) -> JobResult:
+        """Submit a job to SGE."""
+        if interactive:
+            return self._submit_interactive(job)
+        return self._submit_batch(job)
+
+    def _submit_batch(self, job: "Job") -> JobResult:
+        """Submit via qsub."""
+        script = self.generate_script(job)
+
+        with tempfile.NamedTemporaryFile(
+            mode="w", suffix=".sh", delete=False, prefix="hpc_"
+        ) as f:
+            f.write(script)
+            script_path = f.name
+
+        try:
+            cmd = ["qsub", script_path]
+            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+            job_id = parse_qsub_output(result.stdout)
+
+            if job_id is None:
+                raise RuntimeError(f"Failed to parse job ID from qsub output: {result.stdout}")
+
+            return JobResult(job_id=job_id, scheduler=self, job=job)
+        finally:
+            Path(script_path).unlink(missing_ok=True)
+
+    def _submit_interactive(self, job: "Job") -> JobResult:
+        """Submit via qrsh for interactive execution."""
+        cmd = self.build_interactive_command(job)
+        result = subprocess.run(cmd, check=False)
+        # For interactive jobs, we don't have a real job ID
+        return JobResult(job_id="interactive", scheduler=self, job=job)
+
+    def submit_array(self, array: "JobArray") -> ArrayJobResult:
+        """Submit array job."""
+        job = array.job
+        script = self.generate_script(job, array_range=array.range_str)
+
+        with tempfile.NamedTemporaryFile(
+            mode="w", suffix=".sh", delete=False, prefix="hpc_"
+        ) as f:
+            f.write(script)
+            script_path = f.name
+
+        try:
+            cmd = ["qsub", script_path]
+            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+            job_id = parse_qsub_output(result.stdout)
+
+            if job_id is None:
+                raise RuntimeError(f"Failed to parse job ID from qsub output: {result.stdout}")
+
+            return ArrayJobResult(base_job_id=job_id, scheduler=self, array=array)
+        finally:
+            Path(script_path).unlink(missing_ok=True)
+
+    def cancel(self, job_id: str) -> bool:
+        """Cancel a job via qdel."""
+        try:
+            subprocess.run(["qdel", job_id], check=True, capture_output=True)
+            return True
+        except subprocess.CalledProcessError:
+            return False
+
+    def get_status(self, job_id: str) -> JobStatus:
+        """Get job status via qstat."""
+        # Try qstat first (running/pending jobs)
+        try:
+            result = subprocess.run(
+                ["qstat", "-j", job_id],
+                capture_output=True,
+                text=True,
+            )
+            if result.returncode == 0:
+                # Job exists, check state from regular qstat
+                result2 = subprocess.run(
+                    ["qstat"],
+                    capture_output=True,
+                    text=True,
+                )
+                if result2.returncode == 0:
+                    jobs = parse_qstat_plain(result2.stdout)
+                    # Handle array job task IDs (e.g., 12345.1)
+                    base_id = job_id.split(".")[0]
+                    if base_id in jobs:
+                        state = jobs[base_id].get("state", "")
+                        return state_to_status(state)
+                    # Check if full ID matches
+                    if job_id in jobs:
+                        state = jobs[job_id].get("state", "")
+                        return state_to_status(state)
+
+                # Job exists but not in qstat output - likely running
+                return JobStatus.RUNNING
+        except subprocess.CalledProcessError:
+            pass
+
+        # Job not in qstat, check qacct for completed jobs
+        try:
+            result = subprocess.run(
+                ["qacct", "-j", job_id],
+                capture_output=True,
+                text=True,
+            )
+            if result.returncode == 0:
+                info = parse_qacct_output(result.stdout)
+                exit_status = info.get("exit_status", "")
+                if exit_status == "0":
+                    return JobStatus.COMPLETED
+                else:
+                    return JobStatus.FAILED
+        except subprocess.CalledProcessError:
+            pass
+
+        return JobStatus.UNKNOWN
+
+    def get_exit_code(self, job_id: str) -> int | None:
+        """Get exit code from qacct."""
+        try:
+            result = subprocess.run(
+                ["qacct", "-j", job_id],
+                capture_output=True,
+                text=True,
+            )
+            if result.returncode == 0:
+                info = parse_qacct_output(result.stdout)
+                exit_status = info.get("exit_status")
+                if exit_status is not None:
+                    return int(exit_status)
+        except (subprocess.CalledProcessError, ValueError):
+            pass
+        return None
+
+    def get_output_path(self, job_id: str, stream: str) -> Path | None:
+        """Determine output path.
+
+        SGE uses patterns that need to be resolved.
+        """
+        # This is tricky with SGE as paths can use $JOB_ID, etc.
+        # For now, return None and let user check
+        return None
+
+    def generate_script(self, job: "Job", array_range: str | None = None) -> str:
+        """Generate qsub script using template."""
+        directives = self._build_directives(job, array_range)
+        return render_template(
+            "sge/templates/job.sh.j2",
+            job=job,
+            scheduler=self,
+            directives=directives,
+        )
+
+    def _build_directives(self, job: "Job", array_range: str | None = None) -> list[str]:
+        """Build #$ directives."""
+        directives: list[str] = []
+
+        # Shell
+        directives.append("#$ -S /bin/bash")
+
+        # Use current working directory
+        if job.workdir is None:
+            directives.append("#$ -cwd")
+
+        # Job name
+        if job.name:
+            directives.append(f"#$ -N {job.name}")
+
+        # CPU/slots via parallel environment
+        if job.cpu:
+            directives.append(f"#$ -pe {self.pe_name} {job.cpu}")
+
+        # Memory
+        if job.mem:
+            directives.append(f"#$ -l {self.mem_resource}={job.mem}")
+
+        # Time
+        if job.time:
+            directives.append(f"#$ -l {self.time_resource}={job.time}")
+
+        # Queue
+        if job.queue:
+            directives.append(f"#$ -q {job.queue}")
+
+        # Output handling - merge by default
+        if job.merge_output:
+            directives.append("#$ -j y")
+            if job.stdout:
+                directives.append(f"#$ -o {job.stdout}")
+        else:
+            if job.stdout:
+                directives.append(f"#$ -o {job.stdout}")
+            if job.stderr:
+                directives.append(f"#$ -e {job.stderr}")
+
+        # Array job
+        if array_range:
+            directives.append(f"#$ -t {array_range}")
+
+        # Resources (GRES-style)
+        for resource in job.resources:
+            directives.append(f"#$ -l {resource.name}={resource.value}")
+
+        # Dependencies
+        if job.dependencies:
+            dep_ids = ",".join(dep.job_id for dep in job.dependencies)
+            # SGE uses -hold_jid for dependencies
+            directives.append(f"#$ -hold_jid {dep_ids}")
+
+        # Raw args
+        for arg in job.raw_args + job.sge_args:
+            if arg.startswith("-"):
+                directives.append(f"#$ {arg}")
+            else:
+                directives.append(f"#$ -{arg}")
+
+        return directives
+
+    def build_submit_command(self, job: "Job") -> list[str]:
+        """Build qsub command line."""
+        cmd = ["qsub"]
+
+        if job.name:
+            cmd.extend(["-N", job.name])
+        if job.cpu:
+            cmd.extend(["-pe", self.pe_name, str(job.cpu)])
+        if job.mem:
+            cmd.extend(["-l", f"{self.mem_resource}={job.mem}"])
+        if job.time:
+            cmd.extend(["-l", f"{self.time_resource}={job.time}"])
+        if job.queue:
+            cmd.extend(["-q", job.queue])
+
+        cmd.extend(job.raw_args)
+        cmd.extend(job.sge_args)
+
+        return cmd
+
+    def build_interactive_command(self, job: "Job") -> list[str]:
+        """Build qrsh command for interactive jobs."""
+        cmd = ["qrsh"]
+
+        if job.cpu:
+            cmd.extend(["-pe", self.pe_name, str(job.cpu)])
+        if job.mem:
+            cmd.extend(["-l", f"{self.mem_resource}={job.mem}"])
+        if job.time:
+            cmd.extend(["-l", f"{self.time_resource}={job.time}"])
+        if job.queue:
+            cmd.extend(["-q", job.queue])
+
+        cmd.extend(job.raw_args)
+        cmd.extend(job.sge_args)
+
+        # Add the command
+        if isinstance(job.command, str):
+            cmd.append(job.command)
+        else:
+            cmd.extend(job.command)
+
+        return cmd

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

@@ -0,0 +1,39 @@
+#!/bin/bash
+# Generated by hpc-tools (SGE scheduler)
+
+{% for directive in directives %}
+{{ directive }}
+{% endfor %}
+
+# Exit on error
+set -e
+
+# Source module system if available
+if [ -f /etc/profile.d/modules.sh ]; then
+    . /etc/profile.d/modules.sh
+elif [ -f /usr/share/Modules/init/bash ]; then
+    . /usr/share/Modules/init/bash
+fi
+
+{% if job.modules_path %}
+# Additional module paths
+{% for path in job.modules_path %}
+module use {{ path }}
+{% endfor %}
+{% endif %}
+
+{% if job.modules %}
+# Load modules
+{% 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/templates/__init__.py

@@ -0,0 +1,5 @@
+"""Template engine for job script generation."""
+
+from hpc_runner.templates.engine import render_template
+
+__all__ = ["render_template"]

hpc_runner/templates/engine.py

@@ -0,0 +1,55 @@
+"""Jinja2 template engine for job scripts."""
+
+from pathlib import Path
+from typing import Any
+
+import jinja2
+
+# Template directories - schedulers and package templates
+_SCHEDULERS_DIR = Path(__file__).parent.parent / "schedulers"
+_PACKAGE_DIR = Path(__file__).parent
+
+_env: jinja2.Environment | None = None
+
+
+def _get_env() -> jinja2.Environment:
+    """Get or create the Jinja2 environment."""
+    global _env
+    if _env is None:
+        _env = jinja2.Environment(
+            loader=jinja2.FileSystemLoader([str(_SCHEDULERS_DIR), str(_PACKAGE_DIR)]),
+            trim_blocks=True,
+            lstrip_blocks=True,
+            keep_trailing_newline=True,
+        )
+    return _env
+
+
+def render_template(name: str, **context: Any) -> str:
+    """Render a template.
+
+    Args:
+        name: Template name (e.g., "sge/templates/job.sh.j2")
+        **context: Template context variables
+
+    Returns:
+        Rendered template content
+    """
+    env = _get_env()
+    template = env.get_template(name)
+    return template.render(**context)
+
+
+def render_string(template_str: str, **context: Any) -> str:
+    """Render a template string.
+
+    Args:
+        template_str: Template content as a string
+        **context: Template context variables
+
+    Returns:
+        Rendered content
+    """
+    env = _get_env()
+    template = env.from_string(template_str)
+    return template.render(**context)

hpc_runner/workflow/__init__.py

@@ -0,0 +1,6 @@
+"""Workflow support for job pipelines."""
+
+from hpc_runner.workflow.dependency import DependencyType
+from hpc_runner.workflow.pipeline import Pipeline, PipelineJob
+
+__all__ = ["Pipeline", "PipelineJob", "DependencyType"]

hpc_runner/workflow/dependency.py

@@ -0,0 +1,20 @@
+"""Dependency type definitions."""
+
+from enum import Enum
+
+
+class DependencyType(str, Enum):
+    """Job dependency types.
+
+    These map to scheduler-specific dependency modes:
+    - SGE: -hold_jid (basic), -hold_jid_ad (array)
+    - Slurm: --dependency=afterok, afterany, after, afternotok
+    """
+
+    AFTEROK = "afterok"  # Run after all dependencies complete successfully
+    AFTERANY = "afterany"  # Run after all dependencies complete (success or failure)
+    AFTER = "after"  # Run after all dependencies start
+    AFTERNOTOK = "afternotok"  # Run after any dependency fails
+
+    def __str__(self) -> str:
+        return self.value

hpc_runner/workflow/pipeline.py

@@ -0,0 +1,180 @@
+"""Pipeline API for job workflows with dependencies."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+from hpc_runner.core.job import Job
+from hpc_runner.workflow.dependency import DependencyType
+
+if TYPE_CHECKING:
+    from hpc_runner.core.result import JobResult
+    from hpc_runner.schedulers.base import BaseScheduler
+
+
+@dataclass
+class PipelineJob:
+    """A job within a pipeline."""
+
+    job: Job
+    name: str
+    depends_on: list[PipelineJob] = field(default_factory=list)
+    result: JobResult | None = None
+
+
+class Pipeline:
+    """Workflow pipeline with job dependencies.
+
+    Example:
+        with Pipeline("build-test") as p:
+            build = p.add("make build", name="build", cpu=4)
+            test = p.add("make test", name="test", depends_on=["build"])
+            package = p.add("make package", name="package", depends_on=["test"])
+
+        results = p.submit()
+        p.wait()
+    """
+
+    def __init__(self, name: str = "pipeline") -> None:
+        self.name = name
+        self.jobs: list[PipelineJob] = []
+        self._name_map: dict[str, PipelineJob] = {}
+        self._submitted = False
+
+    def add(
+        self,
+        command: str,
+        name: str | None = None,
+        depends_on: list[str | PipelineJob] | None = None,
+        **job_kwargs: Any,
+    ) -> PipelineJob:
+        """Add a job to the pipeline.
+
+        Args:
+            command: Command to execute
+            name: Job name (auto-generated if None)
+            depends_on: List of job names or PipelineJob objects
+            **job_kwargs: Additional Job parameters
+        """
+        if name is None:
+            name = f"step_{len(self.jobs) + 1}"
+
+        if name in self._name_map:
+            raise ValueError(f"Job name '{name}' already exists in pipeline")
+
+        job = Job(command=command, name=f"{self.name}_{name}", **job_kwargs)
+
+        dependencies: list[PipelineJob] = []
+        if depends_on:
+            for dep in depends_on:
+                if isinstance(dep, str):
+                    if dep not in self._name_map:
+                        raise ValueError(f"Unknown dependency: {dep}")
+                    dependencies.append(self._name_map[dep])
+                else:
+                    dependencies.append(dep)
+
+        pipeline_job = PipelineJob(job=job, name=name, depends_on=dependencies)
+        self.jobs.append(pipeline_job)
+        self._name_map[name] = pipeline_job
+
+        return pipeline_job
+
+    def submit(
+        self,
+        scheduler: BaseScheduler | None = None,
+        dependency_type: DependencyType = DependencyType.AFTEROK,
+    ) -> dict[str, JobResult]:
+        """Submit all jobs respecting dependencies.
+
+        Args:
+            scheduler: Scheduler to use (auto-detect if None)
+            dependency_type: Type of dependency to use
+
+        Returns:
+            Dict mapping job names to results
+        """
+        from hpc_runner.schedulers import get_scheduler
+
+        if self._submitted:
+            raise RuntimeError("Pipeline has already been submitted")
+
+        if scheduler is None:
+            scheduler = get_scheduler()
+
+        results: dict[str, JobResult] = {}
+
+        for pjob in self._topological_sort():
+            # Set up dependencies
+            if pjob.depends_on:
+                dep_results = [results[d.name] for d in pjob.depends_on]
+                pjob.job.dependencies = dep_results
+                pjob.job.dependency_type = str(dependency_type)
+
+            # Submit
+            result = scheduler.submit(pjob.job)
+            pjob.result = result
+            results[pjob.name] = result
+
+        self._submitted = True
+        return results
+
+    def _topological_sort(self) -> list[PipelineJob]:
+        """Sort jobs by dependency order (Kahn's algorithm)."""
+        # Build in-degree map
+        in_degree: dict[str, int] = {pj.name: 0 for pj in self.jobs}
+        for pj in self.jobs:
+            for dep in pj.depends_on:
+                in_degree[pj.name] += 1
+
+        # Find all jobs with no dependencies
+        queue = [pj for pj in self.jobs if in_degree[pj.name] == 0]
+        result: list[PipelineJob] = []
+
+        while queue:
+            pj = queue.pop(0)
+            result.append(pj)
+
+            # Reduce in-degree for dependent jobs
+            for other_pj in self.jobs:
+                if pj in other_pj.depends_on:
+                    in_degree[other_pj.name] -= 1
+                    if in_degree[other_pj.name] == 0:
+                        queue.append(other_pj)
+
+        if len(result) != len(self.jobs):
+            raise ValueError("Circular dependency detected in pipeline")
+
+        return result
+
+    def wait(self, poll_interval: float = 5.0) -> dict[str, JobResult]:
+        """Wait for all jobs to complete.
+
+        Returns:
+            Dict mapping job names to results
+        """
+        if not self._submitted:
+            raise RuntimeError("Pipeline has not been submitted")
+
+        for pjob in self.jobs:
+            if pjob.result:
+                pjob.result.wait(poll_interval=poll_interval)
+
+        return {pj.name: pj.result for pj in self.jobs if pj.result}
+
+    def get_job(self, name: str) -> PipelineJob | None:
+        """Get a job by name."""
+        return self._name_map.get(name)
+
+    def __enter__(self) -> Pipeline:
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        pass
+
+    def __len__(self) -> int:
+        return len(self.jobs)
+
+    def __iter__(self):
+        return iter(self.jobs)

hpc_runner-0.1.0.dist-info/METADATA

@@ -0,0 +1,46 @@
+Metadata-Version: 2.4
+Name: hpc-runner
+Version: 0.1.0
+Summary: Unified HPC job submission across multiple schedulers
+Project-URL: Homepage, https://github.com/shareefj/hpc-tools
+Project-URL: Repository, https://github.com/shareefj/hpc-tools
+Author: Shareef Jalloq
+License-Expression: MIT
+Keywords: cluster,hpc,job-submission,pbs,sge,slurm
+Classifier: Development Status :: 3 - Alpha
+Classifier: Environment :: Console
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: System :: Clustering
+Classifier: Topic :: System :: Distributed Computing
+Requires-Python: >=3.10
+Requires-Dist: jinja2>=3.0
+Requires-Dist: rich-click>=1.7
+Requires-Dist: tomli>=2.0; python_version < '3.11'
+Provides-Extra: all
+Requires-Dist: build; extra == 'all'
+Requires-Dist: hatch-vcs; extra == 'all'
+Requires-Dist: mypy; extra == 'all'
+Requires-Dist: pytest-cov; extra == 'all'
+Requires-Dist: pytest>=7.0; extra == 'all'
+Requires-Dist: ruff; extra == 'all'
+Requires-Dist: twine; extra == 'all'
+Provides-Extra: dev
+Requires-Dist: build; extra == 'dev'
+Requires-Dist: hatch-vcs; extra == 'dev'
+Requires-Dist: mypy; extra == 'dev'
+Requires-Dist: pytest-cov; extra == 'dev'
+Requires-Dist: pytest>=7.0; extra == 'dev'
+Requires-Dist: ruff; extra == 'dev'
+Requires-Dist: twine; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# hpc-tools
+
+A collection of tools aimed at abstracting the intricacies of HPC job schedulers from the user.  Having writen tools like this in every job I've had, I thought it about time to write one to rule them all.  With some help from my mate Claude.