"""
SSH connection helper for transferring SLURM scripts to HPC and submitting jobs.
Uses Fabric (https://www.fabfile.org/) for robust SSH connection management.
This module reads your SSH config file and provides utilities to:
1. Transfer generated SLURM scripts to a remote HPC cluster
2. Submit jobs via sbatch
3. Monitor job status
4. Download results
Usage:
from src.tools.connection import HPCConnection
# Initialize connection (reads ~/.ssh/config)
hpc = HPCConnection(host_alias="mycluster")
# Copy SLURM script and submit job
job_id = hpc.submit_job(slurm_file="outputs/train_cgan_cnn_2d_beams2d_seed1.slurm")
print(f"Job submitted with ID: {job_id}")
# Check job status
status = hpc.get_job_status(job_id)
print(status)
"""
import base64
import logging
import os
import re
import sys
from collections.abc import Callable
from pathlib import Path
from typing import Any, Protocol
import requests
from fabric import Connection
logger = logging.getLogger(__name__)
# Type alias for progress callback
ProgressCallback = Callable[[str, str], None]
[docs]
class HPCConnectionLike(Protocol):
"""Operations used by EngiAI regardless of the SSH transport."""
host_alias: str
auth_mode: str
[docs]
def get_remote_directory(self) -> str: ...
[docs]
def ensure_remote_dir(self, remote_dir: str) -> None: ...
[docs]
def put_file(self, local_file: str, remote_path: str) -> None: ...
[docs]
def get_file(self, remote_file: str, local_path: str) -> None: ...
[docs]
def test_slurm_script(self, remote_path: str) -> str: ...
[docs]
def remove_remote_file(self, remote_path: str) -> None: ...
[docs]
def submit_job(self, slurm_file: str, remote_dir: str = "~/slurm_jobs") -> str: ...
[docs]
def get_job_status(self, job_id: str) -> str: ...
[docs]
def get_job_final_state(self, job_id: str) -> tuple[str, str]: ...
[docs]
def cancel_job(self, job_id: str) -> None: ...
[docs]
def cancel_jobs(self, job_ids: list[str]) -> None: ...
self, job_id: str, remote_dir: str = "~/slurm_jobs", local_dir: str = "outputs"
) -> None: ...
[docs]
def get_job_diagnostics(
self, job_id: str, remote_dir: str = "~/slurm_jobs"
) -> tuple[str, str]: ...
[docs]
def list_active_jobs(self) -> str: ...
[docs]
def close(self) -> None: ...
[docs]
class HPCBrokerConnection:
"""HPC connection backed by the native, key-holding host broker."""
[docs]
def __init__(
self,
broker_url: str,
token: str,
host_alias: str = "euler",
progress_callback: ProgressCallback | None = None,
timeout: int = 120,
) -> None:
if not broker_url or not token:
raise RuntimeError("HPC broker URL and token are required")
self.base_url = broker_url.rstrip("/")
self.token = token
self.host_alias = host_alias or "euler"
self.progress_callback = progress_callback
self.timeout = timeout
self.auth_mode = "managed broker"
def _emit_progress(self, step: str, message: str) -> None:
if self.progress_callback:
self.progress_callback(step, message)
logger.info("[%s] %s", step, message)
def _request(
self,
path: str,
payload: dict[str, Any] | None = None,
) -> dict[str, Any]:
try:
response = requests.post(
f"{self.base_url}{path}",
json=payload or {},
headers={"Authorization": f"Bearer {self.token}"},
timeout=self.timeout,
)
response.raise_for_status()
data = response.json()
except requests.RequestException as exc:
detail = ""
if exc.response is not None:
try:
detail = str(exc.response.json().get("error", ""))
except (ValueError, AttributeError):
detail = exc.response.text[-500:]
message = detail or type(exc).__name__
raise RuntimeError(f"HPC broker request failed: {message}") from exc
except ValueError as exc:
raise RuntimeError("HPC broker returned invalid JSON") from exc
if not isinstance(data, dict):
raise TypeError("HPC broker returned an invalid response")
return data
@staticmethod
def _encode_file(path: Path) -> str:
return base64.b64encode(path.read_bytes()).decode("ascii")
[docs]
def get_remote_directory(self) -> str:
return str(self._request("/v1/test")["remote_directory"])
[docs]
def ensure_remote_dir(self, remote_dir: str) -> None:
self._request("/v1/directories", {"remote_dir": remote_dir})
[docs]
def put_file(self, local_file: str, remote_path: str) -> None:
local_path = Path(local_file)
if not local_path.exists():
raise FileNotFoundError(f"Local file not found: {local_file}")
self._request(
"/v1/files/upload",
{
"remote_path": remote_path,
"content_base64": self._encode_file(local_path),
},
)
[docs]
def get_file(self, remote_file: str, local_path: str) -> None:
data = self._request("/v1/files/download", {"remote_path": remote_file})
try:
content = base64.b64decode(str(data["content_base64"]), validate=True)
except (KeyError, ValueError) as exc:
raise RuntimeError("HPC broker returned invalid file content") from exc
destination = Path(local_path)
destination.parent.mkdir(parents=True, exist_ok=True)
destination.write_bytes(content)
[docs]
def test_slurm_script(self, remote_path: str) -> str:
return str(
self._request("/v1/slurm/validate", {"remote_path": remote_path})["output"]
)
[docs]
def remove_remote_file(self, remote_path: str) -> None:
self._request("/v1/files/remove", {"remote_path": remote_path})
[docs]
def submit_job(
self,
slurm_file: str,
remote_dir: str = "~/slurm_jobs",
) -> str:
local_path = Path(slurm_file)
if not local_path.exists():
raise FileNotFoundError(f"SLURM script not found: {slurm_file}")
self._emit_progress("transfer", f"Transferring {local_path.name} to Euler")
data = self._request(
"/v1/jobs/submit",
{
"filename": local_path.name,
"remote_dir": remote_dir,
"content_base64": self._encode_file(local_path),
},
)
job_id = str(data["job_id"])
self._emit_progress("complete", f"Job submitted successfully! Job ID: {job_id}")
return job_id
[docs]
def get_job_status(self, job_id: str) -> str:
return str(self._request(f"/v1/jobs/{job_id}/status")["output"])
[docs]
def get_job_final_state(self, job_id: str) -> tuple[str, str]:
data = self._request(f"/v1/jobs/{job_id}/final")
return str(data.get("state", "")), str(data.get("exit_code", ""))
[docs]
def cancel_job(self, job_id: str) -> None:
self.cancel_jobs([job_id])
[docs]
def cancel_jobs(self, job_ids: list[str]) -> None:
self._request("/v1/jobs/cancel", {"job_ids": job_ids})
[docs]
def get_job_output(
self,
job_id: str,
remote_dir: str = "~/slurm_jobs",
local_dir: str = "outputs",
) -> None:
data = self._request(f"/v1/jobs/{job_id}/outputs", {"remote_dir": remote_dir})
destination = Path(local_dir)
destination.mkdir(parents=True, exist_ok=True)
for item in data.get("files", []):
name = Path(str(item["name"])).name
content = base64.b64decode(str(item["content_base64"]), validate=True)
(destination / name).write_bytes(content)
[docs]
def get_job_diagnostics(
self,
job_id: str,
remote_dir: str = "~/slurm_jobs",
) -> tuple[str, str]:
data = self._request(
f"/v1/jobs/{job_id}/diagnostics", {"remote_dir": remote_dir}
)
return str(data.get("sacct", "")), str(data.get("logs", ""))
[docs]
def list_active_jobs(self) -> str:
return str(self._request("/v1/jobs/active")["output"])
[docs]
def close(self) -> None:
"""The broker owns SSH process lifetimes; there is no client socket to close."""
[docs]
class HPCConnection:
"""SSH connection to HPC cluster for job submission using Fabric."""
[docs]
def __init__( # noqa: PLR0913
self,
host_alias: str = "",
host: str | None = None,
user: str | None = None,
password: str | None = None,
port: int = 22,
progress_callback: ProgressCallback | None = None,
) -> None:
"""
Initialize HPC connection.
Supports two authentication modes:
1. SSH config mode (default): Uses ~/.ssh/config with SSH keys
2. Password mode: Uses explicit username/password credentials
Args:
host_alias: Host alias from ~/.ssh/config (e.g., "mycluster")
host: Explicit hostname (e.g., "hpc.example.com"). If provided, uses password auth.
user: Username for password authentication
password: Password for authentication
port: SSH port (default: 22)
progress_callback: Optional callback for progress updates (step, message)
"""
self.host_alias = host_alias
self.progress_callback = progress_callback
self.auth_mode = "password" if host and user and password else "SSH config"
# Determine authentication mode
if host and user and password:
# Password authentication mode (don't log credentials for security)
logger.info("Establishing SSH connection (password auth)")
try:
self.connection = Connection(
host=host,
user=user,
port=port,
connect_kwargs={
"password": password,
"look_for_keys": False,
"allow_agent": False,
},
)
self.host_alias = host
except Exception as e:
# Don't expose credentials in error messages
msg = "SSH connection failed. Please verify:\n"
msg += " 1. Hostname is correct\n"
msg += " 2. Username and password are correct\n"
msg += " 3. SSH access is enabled on the server\n"
msg += " 4. Network connectivity to the host"
logger.exception(f"SSH connection failed: {type(e).__name__}")
raise RuntimeError(msg) from e
else:
# SSH config mode (original behavior)
if not host_alias:
msg = (
"No HPC host configured. Set HPC_HOST_ALIAS in your .env file "
"(e.g., HPC_HOST_ALIAS=mycluster) or provide host/user/password "
"for direct authentication."
)
raise RuntimeError(msg)
# Fabric automatically uses ~/.ssh/config for connection details
logger.info(f"Connecting to {host_alias} (SSH config mode)")
logger.info(
"Using host alias: %s (from HPC_HOST_ALIAS in .env)", host_alias
)
try:
# Enable SSH agent to handle encrypted private keys
self.connection = Connection(
host_alias,
connect_kwargs={
"allow_agent": True, # Use SSH agent for encrypted keys
"look_for_keys": True, # Look for keys in standard locations
},
)
logger.info(f"SSH config connection established for {host_alias}")
except Exception as e:
logger.exception(f"SSH config connection failed for {host_alias}")
msg = f"Failed to connect to '{host_alias}': {e}\n"
msg += "Make sure:\n"
msg += f" 1. Host {host_alias} is configured in ~/.ssh/config\n"
msg += " 2. SSH key is set up correctly\n"
msg += (
" 3. SSH agent is running (if using Docker, mount SSH_AUTH_SOCK)\n"
)
msg += f" 4. Test manually: ssh {host_alias}"
raise RuntimeError(msg) from e
def _emit_progress(self, step: str, message: str) -> None:
"""
Emit progress update if callback is provided.
Args:
step: Step identifier (e.g., "connect", "transfer", "submit")
message: Human-readable message
"""
if self.progress_callback:
self.progress_callback(step, message)
logger.info(f"[{step}] {message}")
[docs]
def run_command(self, command: str) -> str:
"""
Run a command over SSH.
Args:
command: Command to execute on remote
Returns:
Command output
"""
result = self.connection.run(command, hide=False, warn=False)
if result.failed:
msg = f"SSH command failed: {command}\nError: {result.stderr}"
raise RuntimeError(msg)
return result.stdout.strip()
[docs]
def get_remote_directory(self) -> str:
"""Return the Euler working directory for a connection smoke test."""
return self.run_command("pwd")
[docs]
def ensure_remote_dir(self, remote_dir: str) -> None:
"""Create an allowed working directory on the remote host."""
self.run_command(f"mkdir -p {remote_dir}")
[docs]
def test_slurm_script(self, remote_path: str) -> str:
"""Ask SLURM to validate a staged script without submitting it."""
return self.run_command(f"sbatch --test-only {remote_path}")
[docs]
def remove_remote_file(self, remote_path: str) -> None:
"""Remove a staged remote file."""
self.run_command(f"rm -f {remote_path}")
[docs]
def put_file(self, local_file: str, remote_path: str) -> None:
"""
Copy a file to remote using SFTP.
Args:
local_file: Path to local file
remote_path: Destination path on remote (can be relative to home)
"""
local_path = Path(local_file)
if not local_path.exists():
msg = f"Local file not found: {local_file}"
raise FileNotFoundError(msg)
# Expand tilde in remote path using remote home directory
if remote_path.startswith("~"):
home_dir = self.run_command("echo $HOME").strip()
expanded_remote = remote_path.replace("~", home_dir, 1)
else:
expanded_remote = remote_path
# Fabric's put method handles SFTP file transfer
try:
self.connection.put(str(local_path.absolute()), expanded_remote)
except Exception as e:
msg = f"SFTP transfer failed: {e}"
raise RuntimeError(msg) from e
[docs]
def get_file(self, remote_file: str, local_path: str) -> None:
"""
Download a file from remote using SFTP.
Args:
remote_file: Path to remote file
local_path: Destination path on local
"""
self.connection.get(remote_file, local_path)
[docs]
def submit_job(
self,
slurm_file: str,
remote_dir: str = "~/slurm_jobs",
) -> str:
"""
Transfer SLURM script and submit job to HPC.
Args:
slurm_file: Path to local SLURM script
remote_dir: Directory on remote to store scripts (default: ~/slurm_jobs)
Returns:
Job ID returned by sbatch
"""
local_path = Path(slurm_file)
if not local_path.exists():
msg = f"SLURM script not found: {slurm_file}"
raise FileNotFoundError(msg)
filename = local_path.name
# Ensure remote directory exists (silently)
try:
self.ensure_remote_dir(remote_dir)
except RuntimeError as e:
msg = f"Failed to create remote directory {remote_dir}: {e}"
raise RuntimeError(msg) from e
# Emit progress: Transferring file
remote_path = f"{remote_dir}/{filename}"
# Capitalize Euler for display
display_host = (
self.host_alias.capitalize()
if self.host_alias.lower() == "euler"
else self.host_alias
)
self._emit_progress("transfer", f"Transferring {filename} to {display_host}")
logger.info(f"Copying {filename} to {self.host_alias}:{remote_path}")
self.put_file(slurm_file, remote_path)
self._emit_progress("transfer", "File transferred successfully")
# Emit progress: Submitting job
self._emit_progress("submit", "Submitting SLURM job...")
sbatch_cmd = f"cd {remote_dir} && sbatch {filename}"
output = self.run_command(sbatch_cmd)
# Parse job ID from sbatch output
# Expected format: "Submitted batch job 12345"
match = re.search(r"Submitted batch job (\d+)", output)
if not match:
msg = f"Could not parse job ID from sbatch output: {output}"
raise RuntimeError(msg)
job_id = match.group(1)
self._emit_progress("complete", f"Job submitted successfully! Job ID: {job_id}")
logger.info(f"Job submitted with ID: {job_id}")
return job_id
[docs]
def close(self) -> None:
"""Close the underlying SSH connection, releasing the socket."""
try:
self.connection.close()
except Exception as e:
logger.warning(f"Error closing SSH connection to {self.host_alias}: {e}")
[docs]
def get_job_status(self, job_id: str) -> str:
"""
Get status of a submitted job.
Args:
job_id: SLURM job ID
Returns:
Job status output
"""
try:
output = self.run_command(f"squeue -j {job_id}")
except RuntimeError as e:
return f"Error getting job status: {e}"
else:
if output:
return output
return f"Job {job_id} not found or completed"
[docs]
def get_job_final_state(self, job_id: str) -> tuple[str, str]:
"""Return a finished job's terminal ``(state, exit_code)`` via ``sacct``.
``squeue`` only lists pending/running jobs, so a finished job -- whether
it COMPLETED or FAILED -- simply vanishes from it; callers that only poll
``get_job_status`` cannot tell success from failure. ``sacct`` records the
real final ``State`` and ``ExitCode`` (e.g. a crashed training run that
exits non-zero shows ``FAILED``), which is what failure detection needs.
``state`` is ``""`` when sacct has no record yet -- accounting can lag a
few seconds behind ``squeue`` -- so callers should retry briefly on empty.
"""
try:
# -n no header, -P parsable, -X main allocation only (skip .batch/.extern)
output = self.run_command(
f"sacct -j {job_id} -n -P -X --format=State,ExitCode"
)
except RuntimeError as e:
logger.warning(f"sacct query failed for job {job_id}: {e}")
return "", ""
line = next((ln for ln in output.strip().splitlines() if ln.strip()), "")
if not line:
return "", ""
parts = line.split("|")
# sacct decorates some states, e.g. "CANCELLED by 12345" -> keep the first token
state = parts[0].strip().split()[0] if parts and parts[0].strip() else ""
exit_code = parts[1].strip() if len(parts) > 1 else ""
return state, exit_code
[docs]
def cancel_job(self, job_id: str) -> None:
"""
Cancel a running job.
Args:
job_id: SLURM job ID
"""
logger.info(f"Cancelling job {job_id}")
self.run_command(f"scancel {job_id}")
logger.info(f"Job {job_id} cancelled")
[docs]
def cancel_jobs(self, job_ids: list[str]) -> None:
"""Cancel one or more jobs in one SSH round-trip."""
if not job_ids:
return
self.run_command("scancel " + " ".join(job_ids))
[docs]
def list_active_jobs(self) -> str:
"""List active jobs owned by the SSH account."""
return self.run_command("squeue -u $USER")
[docs]
def get_job_diagnostics(
self,
job_id: str,
remote_dir: str = "~/slurm_jobs",
) -> tuple[str, str]:
"""Return the accounting line and recent stdout/stderr for a job."""
sacct = self.run_command(
f"sacct -j {job_id} --format=State,ExitCode --noheader --parsable2 "
f"2>/dev/null | head -1 || true"
)
logs = self.run_command(
f"{{ tail -n 400 {remote_dir}/engiopt_*_{job_id}.err 2>/dev/null; "
f"echo '----- STDOUT -----'; "
f"tail -n 100 {remote_dir}/engiopt_*_{job_id}.out 2>/dev/null; }} || true"
)
return sacct, logs
[docs]
def get_job_output(
self,
job_id: str,
remote_dir: str = "~/slurm_jobs",
local_dir: str = "outputs",
) -> None:
"""
Download job output files.
Args:
job_id: SLURM job ID
remote_dir: Remote directory where SLURM outputs are stored
local_dir: Local directory to download to
"""
local_path = Path(local_dir)
local_path.mkdir(parents=True, exist_ok=True)
# Download output and error files
for file_type in ["out", "err"]:
# List files matching pattern in remote directory
list_cmd = (
f"ls {remote_dir}/engiopt_*_{job_id}.{file_type} 2>/dev/null || true"
)
try:
file_list_output = self.run_command(list_cmd)
if not file_list_output:
logger.warning(f"No .{file_type} files found for job {job_id}")
continue
# Download each file
for remote_file in file_list_output.split("\n"):
if remote_file.strip():
local_dest = str(local_path / Path(remote_file).name)
self.get_file(remote_file, local_dest)
logger.info(
f"Downloaded {Path(remote_file).name} to {local_dir}/"
)
except RuntimeError as e:
if "No such file" not in str(e):
logger.warning(f"Warning downloading .{file_type} files: {e}")
if __name__ == "__main__":
min_args = 2
if len(sys.argv) < min_args:
print("Usage:")
print(
" python connection.py test # Test SSH connection"
)
print(" python connection.py submit <slurm_file>")
print(" python connection.py status <job_id>")
print(" python connection.py cancel <job_id>")
print(" python connection.py download <job_id>")
sys.exit(1)
command = sys.argv[1]
try:
alias = os.getenv("HPC_HOST_ALIAS", "")
hpc = HPCConnection(host_alias=alias)
if command == "test":
output = hpc.run_command("pwd")
print("✅ SSH connection successful!")
print(f" Remote directory: {output}")
elif command == "submit" and len(sys.argv) > min_args:
job_id = hpc.submit_job(sys.argv[2])
print(f"Job ID: {job_id}")
elif command == "status" and len(sys.argv) > min_args:
status = hpc.get_job_status(sys.argv[2])
print(status)
elif command == "cancel" and len(sys.argv) > min_args:
hpc.cancel_job(sys.argv[2])
elif command == "download" and len(sys.argv) > min_args:
hpc.get_job_output(sys.argv[2])
else:
msg = f"Unknown command: {command}"
print(msg)
sys.exit(1)
except Exception as e:
print(f"❌ Error: {e}")
sys.exit(1)
