"""
HPC cluster job management tools for SLURM job submission and monitoring.
This module provides tools for interacting with HPC clusters via SSH,
including job submission, status monitoring, cancellation, and output retrieval.
"""
import logging
import os
import re
import time
from collections.abc import Generator
from contextlib import contextmanager
from pathlib import Path
from typing import Any
from cryptography.fernet import Fernet
from langchain_core.tools import tool
from config import config
from src.tools.connection import HPCBrokerConnection, HPCConnection, HPCConnectionLike
from src.tools.engibench import get_current_session_id
logger = logging.getLogger(__name__)
# Session credentials storage - keyed by session_id for multi-user isolation
_session_credentials_store: dict[str, dict[str, Any]] = {}
# Current session ID - set by Streamlit before tool invocation
_current_session_id: str | None = None
# Progress callback - set by Streamlit UI to receive progress updates
_progress_callback: Any | None = None
# Per-trial cap on full training-job submissions (self-correction budget for the
# training_code_* benchmark). ``None`` disables the cap, so interactive use and all
# non-training_code problems are unaffected. ``_submit_counts`` tracks successful
# full-run submits keyed by session id (the eval harness sets a unique session id
# per trial); the local 1-epoch dryrun and the cluster ``dryrun_*`` job never count.
# The session id is engibench's harness-set contextvar (get_current_session_id), NOT
# the local _current_session_id below (SSH-credential isolation for the Streamlit UI).
# ``_submit_job_ids`` records the job ids of those full-run submits (same keying) so
# the harness can scancel a dead trial's in-flight jobs at teardown -- see
# ``cancel_session_jobs`` (a trial that hits the agent recursion limit must not leave
# PENDING/RUNNING SLURM jobs running unwatched on the cluster).
_submit_budget: int | None = None
_submit_counts: dict[str, int] = {}
_submit_job_ids: dict[str, list[str]] = {}
# Prefix of the message ``submit_slurm_job`` returns once the per-trial cap is hit.
# ``HPCAgent._should_force_end`` greps the agent's tool messages for this marker to
# stop the loop churning past the budget -- keep the two in sync (single source).
SUBMISSION_BUDGET_EXHAUSTED_MARKER = "Submission budget exhausted"
# Credential expiration time in seconds (20 minutes for security)
CREDENTIAL_EXPIRATION_SECONDS = 1200
# Encryption key for password storage (generated once per application lifetime)
# This provides defense-in-depth - passwords are encrypted in memory
_encryption_key = Fernet.generate_key()
_cipher = Fernet(_encryption_key)
def _encrypt_password(password: str) -> str:
"""Encrypt a password for secure in-memory storage."""
return _cipher.encrypt(password.encode()).decode()
def _decrypt_password(encrypted_password: str) -> str:
"""Decrypt a password from secure storage."""
return _cipher.decrypt(encrypted_password.encode()).decode()
# Hostname validation constants
MAX_HOSTNAME_LENGTH = 253
MAX_LABEL_LENGTH = 63
MAX_USERNAME_LENGTH = 64
MAX_PORT = 65535
[docs]
def set_current_session_id(session_id: str | None) -> None:
"""
Set the current session ID for credential isolation.
This should be called by Streamlit before invoking any HPC tools.
Args:
session_id: Unique session identifier
"""
global _current_session_id # noqa: PLW0603
_current_session_id = session_id
[docs]
def set_submit_budget(k: int | None) -> None:
"""Set the per-trial cap on full training-job submissions (``None`` disables).
Called once by the evaluation harness for ``training_code_*`` runs; the cap is
a cost guardrail enforced inside ``submit_slurm_job``. Dryrun submissions (the
cluster ``dryrun_*`` job) are never counted.
"""
global _submit_budget # noqa: PLW0603
_submit_budget = k
[docs]
def reset_submit_count(session_id: str) -> None:
"""Reset (and free) the submitted-jobs counter and job-id list for a session.
The harness calls this at trial start and teardown so each trial gets a fresh
budget and ``_submit_counts`` / ``_submit_job_ids`` do not grow across a long grid
run. Call ``cancel_session_jobs`` *before* this if a dead trial's jobs need
cancelling -- this clears the recorded ids.
"""
_submit_counts.pop(session_id, None)
_submit_job_ids.pop(session_id, None)
[docs]
def cancel_session_jobs(session_id: str, host_alias: str | None = None) -> list[str]:
"""Best-effort ``scancel`` of every full-run job submitted in ``session_id``.
Called by the eval harness at trial teardown: a trial that died (e.g. hit the
agent recursion limit) must not leave PENDING/RUNNING SLURM jobs running
unwatched on the cluster, wasting GPU and clobbering the moving ``seed_N`` wandb
alias. ``scancel`` on an already-finished job is a harmless no-op, so this is
safe to call after *every* trial. Never raises -- a teardown must not be
derailed by an SSH hiccup.
Returns the list of job ids it attempted to cancel (for logging / tests).
"""
job_ids = _submit_job_ids.get(session_id, [])
if not job_ids:
return []
if host_alias is None:
host_alias = config.hpc_host_alias
try:
hpc = create_hpc_connection(host_alias)
hpc.cancel_jobs(job_ids)
logger.info(
"cancel_session_jobs: scancelled %d in-flight job(s) for session %s: %s",
len(job_ids),
session_id,
", ".join(job_ids),
)
except Exception as e: # best-effort teardown, must never raise
logger.warning(
"cancel_session_jobs: failed to scancel %s for session %s: %s",
job_ids,
session_id,
e,
)
return job_ids
[docs]
def set_hpc_progress_callback(callback: Any | None) -> None:
"""
Set the progress callback for HPC operations.
This should be called by Streamlit before invoking HPC tools.
Args:
callback: Callback function(step: str, message: str) or None to clear
"""
global _progress_callback # noqa: PLW0603
_progress_callback = callback
[docs]
def get_hpc_progress_callback() -> Any | None:
"""Get the current HPC progress callback."""
return _progress_callback
[docs]
def validate_hostname(hostname: str) -> bool:
"""
Validate hostname format to prevent injection attacks.
Args:
hostname: The hostname to validate
Returns:
True if valid, False otherwise
"""
# Allow alphanumeric, dots, and hyphens only
pattern = r"^[a-zA-Z0-9]([a-zA-Z0-9\-\.]*[a-zA-Z0-9])?$"
if not re.match(pattern, hostname):
return False
# Check length
if len(hostname) > MAX_HOSTNAME_LENGTH:
return False
# Check each label
labels = hostname.split(".")
return all(len(label) <= MAX_LABEL_LENGTH for label in labels)
[docs]
def validate_username(username: str) -> bool:
"""
Validate username format.
Args:
username: The username to validate
Returns:
True if valid, False otherwise
"""
# Allow alphanumeric, underscores, and hyphens
pattern = r"^[a-zA-Z0-9_\-]+$"
return bool(re.match(pattern, username)) and len(username) <= MAX_USERNAME_LENGTH
[docs]
def set_ssh_credentials(
host: str | None = None,
user: str | None = None,
password: str | None = None,
port: int = 22,
session_id: str | None = None,
) -> tuple[bool, str]:
"""
Set SSH credentials for a specific session.
Args:
host: Hostname (e.g., "hpc.example.com")
user: Username
password: Password
port: SSH port (default: 22)
session_id: Session ID for isolation (uses current if not provided)
Returns:
Tuple of (success, message)
"""
sid = session_id or _current_session_id
if not sid:
return False, "No session ID available"
if not (host and user and password):
# Clear credentials
if sid in _session_credentials_store:
del _session_credentials_store[sid]
logger.info(f"Cleared SSH credentials for session {sid[:8]}...")
return True, "Credentials cleared"
# Validate inputs
if not validate_hostname(host):
return False, f"Invalid hostname format: {host}"
if not validate_username(user):
return False, f"Invalid username format: {user}"
if not (1 <= port <= MAX_PORT):
return False, f"Invalid port number: {port}"
# Store credentials with expiration (password encrypted for security)
_session_credentials_store[sid] = {
"host": host,
"user": user,
"password": _encrypt_password(password),
"port": port,
"created_at": time.time(),
"expires_at": time.time() + CREDENTIAL_EXPIRATION_SECONDS,
}
logger.info(
f"SSH credentials set for session {sid[:8]}... (expires in {CREDENTIAL_EXPIRATION_SECONDS // 60} min)"
)
return (
True,
f"Credentials saved (expires in {CREDENTIAL_EXPIRATION_SECONDS // 60} minutes)",
)
[docs]
def clear_ssh_credentials(session_id: str | None = None) -> None:
"""
Clear stored SSH credentials for a session.
Args:
session_id: Session ID (uses current if not provided)
"""
sid = session_id or _current_session_id
if sid and sid in _session_credentials_store:
del _session_credentials_store[sid]
logger.info(f"Cleared SSH credentials for session {sid[:8]}...")
[docs]
def get_ssh_credentials(session_id: str | None = None) -> dict[str, Any] | None:
"""
Get SSH credentials for a session.
Args:
session_id: Session ID (uses current if not provided)
Returns:
Credentials dict or None if not set/expired
"""
sid = session_id or _current_session_id
if not sid:
return None
creds = _session_credentials_store.get(sid)
if not creds:
return None
# Check expiration
if creds.get("expires_at", 0) < time.time():
del _session_credentials_store[sid]
logger.info(f"SSH credentials expired for session {sid[:8]}...")
return None
# Return with decrypted password (without internal fields)
return {
"host": creds["host"],
"user": creds["user"],
"password": _decrypt_password(creds["password"]),
"port": creds["port"],
}
[docs]
@contextmanager
def get_ssh_credentials_secure(
session_id: str | None = None,
) -> Generator[dict[str, Any] | None]:
"""
Securely get SSH credentials and clear from memory after use.
This context manager ensures the decrypted password is cleared
from the returned dictionary after the block exits.
Args:
session_id: Session ID (uses current if not provided)
Yields:
dict or None: Credentials dict or None if not set/expired
Example::
with get_ssh_credentials_secure() as creds:
if creds:
connection = HPCConnection(..., password=creds["password"])
# Password is cleared from creds after this point
"""
creds = get_ssh_credentials(session_id)
try:
yield creds
finally:
# Clear the decrypted password from memory
if creds and "password" in creds:
# Overwrite with zeros before deletion
creds["password"] = "\x00" * len(creds["password"])
del creds["password"]
[docs]
def cleanup_expired_credentials() -> int:
"""
Remove all expired credentials from storage.
Returns:
Number of expired sessions cleaned up
"""
current_time = time.time()
expired = [
sid
for sid, creds in _session_credentials_store.items()
if creds.get("expires_at", 0) < current_time
]
for sid in expired:
del _session_credentials_store[sid]
if expired:
logger.info(f"Cleaned up {len(expired)} expired credential sessions")
return len(expired)
def _broker_token() -> str:
"""Read the broker capability from an environment value or mounted file."""
token_file = os.getenv("HPC_BROKER_TOKEN_FILE", "").strip()
if token_file:
try:
return Path(token_file).read_text(encoding="utf-8").strip()
except OSError as exc:
raise RuntimeError(f"Cannot read HPC_BROKER_TOKEN_FILE: {exc}") from exc
return os.getenv("HPC_BROKER_TOKEN", "").strip()
[docs]
def create_hpc_connection(host_alias: str) -> HPCConnectionLike:
"""
Create an HPC connection using session credentials or SSH config.
Args:
host_alias: Host alias for SSH config fallback
Returns:
Direct SSH or managed-broker connection instance
"""
# Clean up expired credentials periodically
cleanup_expired_credentials()
# Get progress callback if set
progress_callback = get_hpc_progress_callback()
# Production deployments keep the technical account's private key in a
# native host service. Broker mode deliberately takes precedence over
# browser-supplied passwords so the shared UI never handles that credential.
broker_url = os.getenv("HPC_BROKER_URL", "").strip()
if broker_url:
token = _broker_token()
if not token:
raise RuntimeError(
"HPC_BROKER_URL is configured but no HPC_BROKER_TOKEN or "
"HPC_BROKER_TOKEN_FILE is available"
)
logger.info("Creating managed HPC broker connection for %s", host_alias)
return HPCBrokerConnection(
broker_url=broker_url,
token=token,
host_alias=host_alias,
progress_callback=progress_callback,
)
# Use secure context manager to ensure password is cleared after use
with get_ssh_credentials_secure() as creds:
if creds:
# Use password authentication
sid = _current_session_id or "unknown"
logger.info(
f"Creating HPC connection with password auth (session {sid[:8]}...)"
)
return HPCConnection(
host_alias=host_alias,
host=creds["host"],
user=creds["user"],
password=creds["password"],
port=creds["port"],
progress_callback=progress_callback,
)
else:
# Use SSH config (original behavior)
logger.info(f"Creating HPC connection with SSH config for {host_alias}")
return HPCConnection(
host_alias=host_alias, progress_callback=progress_callback
)
@tool
def test_hpc_connection(host_alias: str | None = None) -> dict[str, Any]:
"""
Test SSH connection to HPC cluster.
This tool verifies that the SSH connection to the HPC cluster is working
by running a simple 'pwd' command on the remote server.
Args:
host_alias: Host alias from ~/.ssh/config. If not provided, uses HPC_HOST_ALIAS from .env
Returns:
Dictionary with connection status and remote directory
"""
if host_alias is None:
host_alias = config.hpc_host_alias
try:
hpc = create_hpc_connection(host_alias)
output = hpc.get_remote_directory()
except Exception as e:
return {
"status": "failed",
"message": f"SSH connection failed: {e}",
"error": str(e),
}
else:
return {
"status": "success",
"message": (f"SSH connection successful to {host_alias} ({hpc.auth_mode})"),
"remote_directory": output,
}
@tool
def submit_slurm_job(
slurm_file: str,
host_alias: str | None = None,
remote_dir: str = "~/slurm_jobs",
) -> dict[str, Any]:
"""
Submit a SLURM job to HPC cluster.
This tool transfers a SLURM script to the HPC cluster and submits it
using the sbatch command.
Args:
slurm_file: Path to local SLURM script file
host_alias: Host alias from ~/.ssh/config. If not provided, uses HPC_HOST_ALIAS from .env
remote_dir: Remote directory to store scripts (default: ~/slurm_jobs)
Returns:
Dictionary with submission status and job ID
"""
if host_alias is None:
host_alias = config.hpc_host_alias
slurm_path = Path(slurm_file)
if not slurm_path.exists():
return {
"status": "failed",
"message": f"SLURM file not found: {slurm_file}",
}
# Per-trial submission budget (training_code_* cost guardrail). The cluster
# dryrun_* job is excluded -- same definition the scorer's _full_run_job_ids
# uses -- so only real full training runs count against the budget.
budget = _submit_budget
is_full_run = not slurm_path.name.startswith("dryrun_")
sid = get_current_session_id() if (budget is not None and is_full_run) else ""
if budget is not None and is_full_run and _submit_counts.get(sid, 0) >= budget:
return {
"status": "failed",
"message": (
f"{SUBMISSION_BUDGET_EXHAUSTED_MARKER}: "
f"{_submit_counts.get(sid, 0)}/{budget} full training jobs already "
"submitted for this task. Do NOT submit again -- finalize by calling "
"evaluate_model on your best completed job."
),
}
try:
hpc = create_hpc_connection(host_alias)
job_id = hpc.submit_job(slurm_file, remote_dir=remote_dir)
except Exception as e:
return {
"status": "failed",
"message": f"Failed to submit job: {e}",
"error": str(e),
}
else:
if budget is not None and is_full_run:
_submit_counts[sid] = _submit_counts.get(sid, 0) + 1
if job_id:
# Record for teardown cleanup (cancel_session_jobs) -- same keying
# and full-run gating as the budget counter above.
_submit_job_ids.setdefault(sid, []).append(job_id)
return {
"status": "success",
"message": "Job submitted successfully",
"job_id": job_id,
"remote_file": f"{remote_dir}/{slurm_path.name}",
}
@tool
def get_slurm_job_status(
job_id: str,
host_alias: str | None = None,
) -> dict[str, Any]:
"""
Get the status of a SLURM job.
This tool checks the status of a submitted SLURM job using the squeue command.
Args:
job_id: SLURM job ID
host_alias: Host alias from ~/.ssh/config. If not provided, uses HPC_HOST_ALIAS from .env
Returns:
Dictionary with job status information
"""
if host_alias is None:
host_alias = config.hpc_host_alias
try:
hpc = create_hpc_connection(host_alias)
status = hpc.get_job_status(job_id)
except Exception as e:
return {
"status": "failed",
"job_id": job_id,
"message": f"Failed to get job status: {e}",
"error": str(e),
}
else:
return {
"status": "success",
"job_id": job_id,
"status_output": status,
}
@tool
def cancel_slurm_job(
job_id: str,
host_alias: str | None = None,
) -> dict[str, Any]:
"""
Cancel a running SLURM job.
This tool cancels a submitted SLURM job using the scancel command.
Args:
job_id: SLURM job ID
host_alias: Host alias from ~/.ssh/config. If not provided, uses HPC_HOST_ALIAS from .env
Returns:
Dictionary with cancellation status
"""
if host_alias is None:
host_alias = config.hpc_host_alias
try:
hpc = create_hpc_connection(host_alias)
hpc.cancel_job(job_id)
except Exception as e:
return {
"status": "failed",
"job_id": job_id,
"message": f"Failed to cancel job: {e}",
"error": str(e),
}
else:
return {
"status": "success",
"message": f"Job {job_id} cancelled successfully",
"job_id": job_id,
}
@tool
def download_job_outputs(
job_id: str,
host_alias: str | None = None,
remote_dir: str = "~/slurm_jobs",
local_dir: str = "outputs",
) -> dict[str, Any]:
"""
Download job output files from HPC cluster.
This tool downloads the .out and .err output files from a completed
SLURM job to the local machine.
Args:
job_id: SLURM job ID
host_alias: Host alias from ~/.ssh/config. If not provided, uses HPC_HOST_ALIAS from .env
remote_dir: Remote directory where outputs are stored (default: ~/slurm_jobs)
local_dir: Local directory to download to (default: outputs)
Returns:
Dictionary with download status and file locations
"""
if host_alias is None:
host_alias = config.hpc_host_alias
try:
hpc = create_hpc_connection(host_alias)
hpc.get_job_output(job_id, remote_dir=remote_dir, local_dir=local_dir)
except Exception as e:
return {
"status": "failed",
"job_id": job_id,
"message": f"Failed to download outputs: {e}",
"error": str(e),
}
else:
# Build list of downloaded files
output_dir = Path(local_dir)
err_files = list(output_dir.glob(f"*{job_id}.err"))
out_files = list(output_dir.glob(f"*{job_id}.out"))
file_list = [str(f) for f in err_files + out_files]
message = f"Job outputs downloaded successfully to {local_dir}/"
if file_list:
message += "\n\nDownloaded files:\n" + "\n".join(
f"- {Path(f).name}" for f in file_list
)
return {
"status": "success",
"message": message,
"job_id": job_id,
"local_directory": local_dir,
"downloaded_files": file_list,
}
