"""
HPC cluster agent for SLURM job management and training-code workflows.
In addition to the original SLURM job-management tools (submit / monitor / cancel /
download / status), the agent owns the end-to-end training-code-generation
workflow for benchmarks like ``training_code_beams2d``: authoring the candidate
training file, validating it locally, transferring it to the cluster, and running
it through dryrun / submit / monitor / diagnose / evaluate.
The training-code tools are imported from their submodules rather than from
``src.tools`` (its package ``__init__`` deliberately does NOT re-export them --
see the note there -- because they pull in the benchmark ``contract_validator``
and we don't want unrelated agents to drag that in via the tools package init).
"""
from langchain_core.messages import AIMessage, AnyMessage, ToolMessage
from src.agents.base_agent import BaseAgent
from src.models.state import MessagesState
from src.tools.engiopt import evaluate_model
from src.tools.hpc import (
SUBMISSION_BUDGET_EXHAUSTED_MARKER,
cancel_slurm_job,
download_job_outputs,
get_slurm_job_status,
submit_slurm_job,
test_hpc_connection,
)
from src.tools.job_monitor import (
check_job_status_change,
get_active_jobs_summary,
monitor_job_until_complete,
)
from src.tools.training_code_gen import (
validate_training_script,
write_training_script,
)
from src.tools.training_code_hpc import (
diagnose_job_failure,
dryrun_training_script,
generate_training_slurm,
validate_slurm_script,
)
from src.utils.prompts import get_hpc_agent_system_prompt
[docs]
class HPCAgent(BaseAgent):
"""Agent specialized in HPC job management and training-code workflows."""
[docs]
def __init__(
self,
model_name: str | None = None,
temperature: float | None = None,
seed: int | None = None,
):
"""Initialize the HPC agent.
Args:
model_name: Name of the LLM model to use (defaults to config.llm_model)
temperature: Model temperature (defaults to config.llm_temperature)
seed: Random seed for model (defaults to config.llm_seed)
"""
super().__init__(model_name=model_name, temperature=temperature, seed=seed)
def _create_tools(self) -> list:
"""Create the list of HPC tools.
Returns:
List of LangChain tools for HPC operations
"""
return [
# Generic SLURM job management.
test_hpc_connection,
submit_slurm_job,
get_slurm_job_status,
cancel_slurm_job,
download_job_outputs,
monitor_job_until_complete,
check_job_status_change,
get_active_jobs_summary,
# Training-code workflow (write -> validate -> generate SLURM ->
# validate SLURM -> dryrun -> submit -> monitor -> diagnose -> eval).
# Every step lives on this agent (submit / monitor reuse the generic
# block above), so the efficiency scorer sees the full sequence on one
# agent rather than a fragmented multi-agent trace.
write_training_script,
validate_training_script,
generate_training_slurm,
validate_slurm_script,
dryrun_training_script,
diagnose_job_failure,
evaluate_model,
]
def _get_system_prompt(self) -> str:
"""Get the system prompt for the HPC agent.
Returns:
System prompt string
"""
return get_hpc_agent_system_prompt()
# Churn turns allowed after the submission budget is exhausted before the loop is
# hard-stopped, so the agent still has room to finalize (see _FINALIZE_TOOLS).
_POST_BUDGET_TURNS = 3
# Tool calls that *finalize* a trial rather than churn it. A turn issuing one of
# these is the agent wrapping up: it is always let through and never consumes the
# post-budget window -- otherwise the loop could END exactly as the agent emits its
# evaluate_model call, scoring a genuinely completed job as a failure.
_FINALIZE_TOOLS = frozenset({"evaluate_model", "download_job_outputs"})
def _is_finalize_turn(self, msg: AnyMessage) -> bool:
"""True if ``msg`` is an AIMessage issuing a finalize tool call."""
return (
isinstance(msg, AIMessage)
and bool(msg.tool_calls)
and any(tc.get("name") in self._FINALIZE_TOOLS for tc in msg.tool_calls)
)
def _should_force_end(self, state: MessagesState) -> bool:
"""End the loop once the submission budget is exhausted and the agent has
churned past its grace window.
``submit_slurm_job`` returns a ToolMessage containing
``SUBMISSION_BUDGET_EXHAUSTED_MARKER`` once the per-trial cap is hit. Without
this stop a self-correcting agent keeps churning (rewrite -> validate ->
blocked submit -> diagnose ...) until the LangGraph recursion limit kills the
whole trial. After the marker, allow ``_POST_BUDGET_TURNS`` *churn* turns to
wrap up, then force END -- but a finalize turn (evaluate_model / output
download) is always let through and never counts against the window, so the
agent can always run its scoring call on a completed job.
"""
messages = state["messages"]
budget_hit_index: int | None = None
for i, msg in enumerate(messages):
if (
isinstance(msg, ToolMessage)
and isinstance(msg.content, str)
and SUBMISSION_BUDGET_EXHAUSTED_MARKER in msg.content
):
budget_hit_index = i
break
if budget_hit_index is None:
return False
# Never cut off the agent's own finalize call -- if it is wrapping up right
# now, let that tool run regardless of how many churn turns preceded it.
if self._is_finalize_turn(messages[-1]):
return False
# Count only churn turns after the marker; finalize turns don't consume budget.
churn_turns = sum(
1
for msg in messages[budget_hit_index + 1 :]
if isinstance(msg, AIMessage) and not self._is_finalize_turn(msg)
)
return churn_turns >= self._POST_BUDGET_TURNS
