Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 23 additions & 1 deletion src/opensmi/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -1540,9 +1540,27 @@ def _prompt_slurm_cluster(existing: Optional[dict] = None) -> dict:
input(_ob_prompt("SSH user for login node", "", default_user)).strip()
or default_user
)
# Slurm binary prefix (for clusters where sinfo/squeue/scontrol aren't in PATH)
default_prefix = str((existing or {}).get("slurm_bin_prefix") or "")
print(
f"\n {_OB_DIM}If Slurm commands (sinfo, squeue, scontrol) are not in the"
f" remote PATH,{_OB_RESET}"
)
print(
f" {_OB_DIM}enter the full directory path (e.g. /opt/slurm-21.08/bin).{_OB_RESET}"
)
slurm_prefix = (
input(
_ob_prompt("Slurm bin path", "leave empty if default", default_prefix)
).strip()
or default_prefix
)

result: dict = {"name": name, "login_node": login_node, "user": user}
if login_port is not None:
result["port"] = login_port
if slurm_prefix:
result["slurm_bin_prefix"] = slurm_prefix
return result

def _print_review(ssh_clusters: list[dict], slurm_clusters: list[dict]) -> None:
Expand All @@ -1564,9 +1582,11 @@ def _print_review(ssh_clusters: list[dict], slurm_clusters: list[dict]) -> None:
for idx, sc in enumerate(slurm_clusters, start=1):
_sc_port = sc.get("port")
_port_suffix = f":{_sc_port}" if _sc_port else ""
_sc_prefix = sc.get("slurm_bin_prefix", "")
_prefix_suffix = f" slurm: {_sc_prefix}" if _sc_prefix else ""
print(
f" {idx}. {sc.get('name', 'Slurm Cluster')} → "
f"{sc.get('login_node', '')}{_port_suffix} (user: {sc.get('user', '')})"
f"{sc.get('login_node', '')}{_port_suffix} (user: {sc.get('user', '')}){_prefix_suffix}"
)

print(f"\n {_OB_GREEN}╭{line}╮{_OB_RESET}")
Expand Down Expand Up @@ -2882,6 +2902,7 @@ def _cmd_slurm(args: argparse.Namespace) -> int:
login_node=str(sc["login_node"]),
user=str(sc.get("user", "")),
port=int(sc.get("port", 22)),
slurm_bin_prefix=str(sc.get("slurm_bin_prefix", "")),
)
for sc in raw_slurm
]
Expand All @@ -2895,6 +2916,7 @@ def _cmd_slurm(args: argparse.Namespace) -> int:
ssh_user=sc.user,
ssh_port=sc.port,
cluster_name=sc.name,
slurm_bin_prefix=sc.slurm_bin_prefix,
)
results.append(snap)
if args.output_json:
Expand Down
1 change: 1 addition & 0 deletions src/opensmi/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@ def _parse_cluster_config(
login_node=str(raw["login_node"]),
user=str(raw.get("user", "")),
port=int(raw.get("port", 22)),
slurm_bin_prefix=str(raw.get("slurm_bin_prefix", "")),
)
)

Expand Down
1 change: 1 addition & 0 deletions src/opensmi/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ class SlurmClusterConfig:
login_node: str # SSH alias or address for the login node
user: str = "" # SSH user (optional, defaults to current user)
port: int = 22
slurm_bin_prefix: str = "" # e.g. "/opt/slurm-21.08/bin" for non-standard installs


@dataclass
Expand Down
8 changes: 6 additions & 2 deletions src/opensmi/slurm.py
Original file line number Diff line number Diff line change
Expand Up @@ -201,6 +201,7 @@ def collect_slurm_snapshot(
ssh_user: str = "",
ssh_port: int = 22,
cluster_name: str = "Slurm Cluster",
slurm_bin_prefix: str = "",
) -> SlurmClusterSnapshot:
"""Collect a full Slurm cluster snapshot using only Slurm CLI tools.

Expand All @@ -218,11 +219,14 @@ def collect_slurm_snapshot(
)

def run_cmd(cmd: List[str]) -> str:
resolved = list(cmd)
if slurm_bin_prefix and resolved and resolved[0] in ("sinfo", "squeue", "scontrol"):
resolved[0] = f"{slurm_bin_prefix.rstrip('/')}/{resolved[0]}"
if login_node:
return _run_remote(
cmd, login_node, user=ssh_user, port=ssh_port, timeout=timeout
resolved, login_node, user=ssh_user, port=ssh_port, timeout=timeout
)
return _run(cmd, timeout=timeout)
return _run(resolved, timeout=timeout)

# 1. sinfo — node list
try:
Expand Down
Loading