SLURM compatible commands
This commit is contained in:
Sami Khan
@@ -0,0 +1,32 @@
|
||||
"""Exo CLI - SLURM-compatible job management commands."""
|
||||
|
||||
|
||||
def run_subcommand(command: str, args: list[str]) -> int:
|
||||
"""Route to the appropriate subcommand handler.
|
||||
|
||||
Args:
|
||||
command: The subcommand name (sbatch, squeue, scancel, salloc)
|
||||
args: Command line arguments for the subcommand
|
||||
|
||||
Returns:
|
||||
Exit code from the subcommand
|
||||
"""
|
||||
if command == "sbatch":
|
||||
from exo.cli.sbatch import main
|
||||
|
||||
return main(args)
|
||||
elif command == "squeue":
|
||||
from exo.cli.squeue import main
|
||||
|
||||
return main(args)
|
||||
elif command == "scancel":
|
||||
from exo.cli.scancel import main
|
||||
|
||||
return main(args)
|
||||
elif command == "salloc":
|
||||
from exo.cli.salloc import main
|
||||
|
||||
return main(args)
|
||||
else:
|
||||
print(f"Unknown subcommand: {command}")
|
||||
return 1
|
||||
@@ -0,0 +1,118 @@
|
||||
"""Common utilities for Exo CLI commands."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import urllib.request
|
||||
from typing import Any
|
||||
from urllib.error import HTTPError, URLError
|
||||
|
||||
# Default API endpoint
|
||||
DEFAULT_API_HOST = "localhost"
|
||||
DEFAULT_API_PORT = 52415
|
||||
|
||||
|
||||
def get_api_base() -> str:
|
||||
"""Get the API base URL from environment or defaults."""
|
||||
host = os.environ.get("EXO_API_HOST", DEFAULT_API_HOST)
|
||||
port = os.environ.get("EXO_API_PORT", str(DEFAULT_API_PORT))
|
||||
return f"http://{host}:{port}"
|
||||
|
||||
|
||||
def api_request(
|
||||
method: str,
|
||||
path: str,
|
||||
data: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any] | list[Any]:
|
||||
"""Make an API request to the Exo server.
|
||||
|
||||
Args:
|
||||
method: HTTP method (GET, POST, DELETE, etc.)
|
||||
path: API path (e.g., "/flash/instances")
|
||||
data: Optional JSON data for POST/PUT requests
|
||||
|
||||
Returns:
|
||||
Parsed JSON response
|
||||
|
||||
Raises:
|
||||
SystemExit: On connection or HTTP errors
|
||||
"""
|
||||
url = f"{get_api_base()}{path}"
|
||||
|
||||
request_data = None
|
||||
if data is not None:
|
||||
request_data = json.dumps(data).encode("utf-8")
|
||||
|
||||
req = urllib.request.Request(
|
||||
url,
|
||||
data=request_data,
|
||||
method=method,
|
||||
)
|
||||
req.add_header("Content-Type", "application/json")
|
||||
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=30) as response: # pyright: ignore[reportAny]
|
||||
body: str = response.read().decode("utf-8") # pyright: ignore[reportAny]
|
||||
if body:
|
||||
return json.loads(body) # pyright: ignore[reportAny]
|
||||
return {}
|
||||
except HTTPError as e:
|
||||
error_body = e.read().decode("utf-8") if e.fp else ""
|
||||
print(f"API error: {e.code} {e.reason}")
|
||||
if error_body:
|
||||
try:
|
||||
error_json: dict[str, str] = json.loads(error_body) # pyright: ignore[reportAny]
|
||||
if "detail" in error_json:
|
||||
print(f" {error_json['detail']}")
|
||||
except json.JSONDecodeError:
|
||||
print(f" {error_body}")
|
||||
raise SystemExit(1)
|
||||
except URLError as e:
|
||||
print(f"Connection error: {e.reason}")
|
||||
print(f"Is Exo running at {get_api_base()}?")
|
||||
raise SystemExit(1)
|
||||
|
||||
|
||||
def truncate_id(instance_id: str, length: int = 8) -> str:
|
||||
"""Truncate a UUID for display.
|
||||
|
||||
Args:
|
||||
instance_id: Full UUID string
|
||||
length: Number of characters to keep
|
||||
|
||||
Returns:
|
||||
Truncated ID without hyphens
|
||||
"""
|
||||
return instance_id.replace("-", "")[:length]
|
||||
|
||||
|
||||
def format_table(headers: list[str], rows: list[list[str]]) -> str:
|
||||
"""Format data as a simple text table.
|
||||
|
||||
Args:
|
||||
headers: Column headers
|
||||
rows: List of rows, each row is a list of column values
|
||||
|
||||
Returns:
|
||||
Formatted table string
|
||||
"""
|
||||
if not rows:
|
||||
return " ".join(f"{h:<10}" for h in headers)
|
||||
|
||||
# Calculate column widths
|
||||
widths = [len(h) for h in headers]
|
||||
for row in rows:
|
||||
for i, cell in enumerate(row):
|
||||
if i < len(widths):
|
||||
widths[i] = max(widths[i], len(cell))
|
||||
|
||||
# Build format string
|
||||
fmt = " ".join(f"{{:<{w}}}" for w in widths)
|
||||
|
||||
# Format output
|
||||
lines = [fmt.format(*headers)]
|
||||
for row in rows:
|
||||
# Pad row if needed
|
||||
padded = row + [""] * (len(headers) - len(row))
|
||||
lines.append(fmt.format(*padded[: len(headers)]))
|
||||
|
||||
return "\n".join(lines)
|
||||
@@ -0,0 +1,100 @@
|
||||
"""salloc - Allocate nodes for interactive use.
|
||||
|
||||
Usage:
|
||||
exo salloc [options] [-- command [args...]]
|
||||
|
||||
Options:
|
||||
-N, --nodes N Number of nodes to allocate (default: 1)
|
||||
--hosts HOSTS Comma-separated list of hostnames
|
||||
|
||||
If a command is provided after --, it will be executed with
|
||||
SLURM-like environment variables set:
|
||||
SLURM_JOB_NODELIST - Comma-separated list of allocated nodes
|
||||
SLURM_NNODES - Number of allocated nodes
|
||||
|
||||
Examples:
|
||||
exo salloc --nodes=2 --hosts=node1,node2 -- mpirun ./my_program
|
||||
exo salloc --hosts=localhost -- bash
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
|
||||
def main(args: list[str]) -> int:
|
||||
"""Main entry point for salloc command."""
|
||||
# Split args at -- if present
|
||||
cmd_args: list[str] = []
|
||||
salloc_args = args
|
||||
|
||||
if "--" in args:
|
||||
idx = args.index("--")
|
||||
salloc_args = args[:idx]
|
||||
cmd_args = args[idx + 1 :]
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="exo salloc",
|
||||
description="Allocate nodes for interactive use",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-N",
|
||||
"--nodes",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of nodes to allocate (default: 1)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hosts",
|
||||
help="Comma-separated list of hostnames (required)",
|
||||
)
|
||||
|
||||
parsed = parser.parse_args(salloc_args)
|
||||
|
||||
nodes: int = parsed.nodes # pyright: ignore[reportAny]
|
||||
hosts: str | None = parsed.hosts # pyright: ignore[reportAny]
|
||||
|
||||
# Require explicit hosts since we can't discover them from topology
|
||||
if not hosts:
|
||||
print("Error: --hosts is required (e.g., --hosts=node1,node2)", file=sys.stderr)
|
||||
print(" The Exo topology doesn't expose hostnames.", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
host_list = [h.strip() for h in hosts.split(",") if h.strip()]
|
||||
|
||||
if len(host_list) < nodes:
|
||||
print(
|
||||
f"Error: Requested {nodes} nodes but only {len(host_list)} hosts provided",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return 1
|
||||
|
||||
# Use first N hosts
|
||||
allocated_hosts = host_list[:nodes]
|
||||
nodelist = ",".join(allocated_hosts)
|
||||
|
||||
# Set environment variables
|
||||
env = os.environ.copy()
|
||||
env["SLURM_JOB_NODELIST"] = nodelist
|
||||
env["SLURM_NNODES"] = str(nodes)
|
||||
|
||||
print(f"salloc: Granted job allocation on {nodes} node(s)")
|
||||
print(f"salloc: Nodes: {nodelist}")
|
||||
|
||||
if cmd_args:
|
||||
# Run the command
|
||||
print(f"salloc: Running: {' '.join(cmd_args)}")
|
||||
result = subprocess.run(cmd_args, env=env)
|
||||
return result.returncode
|
||||
else:
|
||||
# Start interactive shell
|
||||
shell = os.environ.get("SHELL", "/bin/bash")
|
||||
print(f"salloc: Starting shell {shell}")
|
||||
print("salloc: Use 'exit' to release allocation")
|
||||
result = subprocess.run([shell], env=env)
|
||||
return result.returncode
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main(sys.argv[1:]))
|
||||
@@ -0,0 +1,233 @@
|
||||
"""sbatch - Submit a batch job to Exo.
|
||||
|
||||
Usage:
|
||||
exo sbatch [options] <script|executable>
|
||||
exo sbatch --job-name=NAME --nodes=N <executable>
|
||||
|
||||
Options:
|
||||
-J, --job-name NAME Job name
|
||||
-N, --nodes N Number of nodes (default: 1)
|
||||
--ntasks-per-node N Tasks per node (default: 1)
|
||||
-D, --chdir DIR Working directory
|
||||
--hosts HOSTS Comma-separated list of hostnames
|
||||
|
||||
Job scripts can contain #SBATCH directives:
|
||||
#!/bin/bash
|
||||
#SBATCH --job-name=Sod2D
|
||||
#SBATCH --nodes=2
|
||||
#SBATCH --chdir=/path/to/workdir
|
||||
|
||||
/path/to/flash4
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
|
||||
from exo.cli.common import api_request, truncate_id
|
||||
|
||||
|
||||
def parse_job_script(script_path: str) -> tuple[dict[str, str], str | None]:
|
||||
"""Parse a job script for #SBATCH directives and executable.
|
||||
|
||||
Args:
|
||||
script_path: Path to the job script
|
||||
|
||||
Returns:
|
||||
Tuple of (directives dict, executable path or None)
|
||||
"""
|
||||
directives: dict[str, str] = {}
|
||||
executable: str | None = None
|
||||
|
||||
with open(script_path, "r") as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
|
||||
# Parse #SBATCH directives
|
||||
if line.startswith("#SBATCH"):
|
||||
# Handle both --option=value and --option value formats
|
||||
match = re.match(r"#SBATCH\s+(-\w|--[\w-]+)(?:=|\s+)(.+)", line)
|
||||
if match:
|
||||
opt, val = match.groups()
|
||||
directives[opt.lstrip("-")] = val.strip()
|
||||
continue
|
||||
|
||||
# Skip comments and empty lines
|
||||
if line.startswith("#") or not line:
|
||||
continue
|
||||
|
||||
# First non-comment, non-directive line is the executable
|
||||
if executable is None:
|
||||
# Handle lines like "/path/to/flash4" or "srun /path/to/flash4"
|
||||
parts = line.split()
|
||||
if parts:
|
||||
# Skip srun/mpirun prefixes if present
|
||||
for part in parts:
|
||||
if not part.startswith("-") and "/" in part:
|
||||
executable = part
|
||||
break
|
||||
if executable is None and parts:
|
||||
executable = parts[-1] # Last token
|
||||
|
||||
return directives, executable
|
||||
|
||||
|
||||
def main(args: list[str]) -> int:
|
||||
"""Main entry point for sbatch command."""
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="exo sbatch",
|
||||
description="Submit a batch job to Exo",
|
||||
)
|
||||
parser.add_argument(
|
||||
"script",
|
||||
help="Job script or executable path",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-J",
|
||||
"--job-name",
|
||||
dest="job_name",
|
||||
help="Job name",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-N",
|
||||
"--nodes",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of nodes (default: 1)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ntasks-per-node",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Tasks per node (default: 1)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-D",
|
||||
"--chdir",
|
||||
help="Working directory",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hosts",
|
||||
help="Comma-separated list of hostnames",
|
||||
)
|
||||
|
||||
parsed = parser.parse_args(args)
|
||||
|
||||
# Extract typed values from namespace
|
||||
script_path: str = parsed.script # pyright: ignore[reportAny]
|
||||
arg_job_name: str | None = parsed.job_name # pyright: ignore[reportAny]
|
||||
arg_nodes: int = parsed.nodes # pyright: ignore[reportAny]
|
||||
arg_ntasks: int = parsed.ntasks_per_node # pyright: ignore[reportAny]
|
||||
arg_chdir: str | None = parsed.chdir # pyright: ignore[reportAny]
|
||||
arg_hosts: str | None = parsed.hosts # pyright: ignore[reportAny]
|
||||
|
||||
# Determine if input is a script or direct executable
|
||||
executable: str | None = None
|
||||
directives: dict[str, str] = {}
|
||||
|
||||
if os.path.isfile(script_path):
|
||||
# Check if it's a binary file (executable) or text script
|
||||
is_binary = False
|
||||
try:
|
||||
with open(script_path, "rb") as f:
|
||||
chunk = f.read(512)
|
||||
# Binary files typically contain null bytes
|
||||
is_binary = b"\x00" in chunk
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
if is_binary:
|
||||
# It's a binary executable
|
||||
executable = script_path
|
||||
else:
|
||||
# Try to read as text
|
||||
try:
|
||||
with open(script_path, "r") as f:
|
||||
first_line = f.readline()
|
||||
f.seek(0)
|
||||
content = f.read(1024)
|
||||
|
||||
if first_line.startswith("#!") or "#SBATCH" in content:
|
||||
# It's a job script - parse it
|
||||
directives, executable = parse_job_script(script_path)
|
||||
else:
|
||||
# It's an executable (text but no shebang/directives)
|
||||
executable = script_path
|
||||
except UnicodeDecodeError:
|
||||
# Can't read as text - treat as binary executable
|
||||
executable = script_path
|
||||
else:
|
||||
# Not a file - treat as executable path
|
||||
executable = script_path
|
||||
|
||||
if executable is None:
|
||||
print("Error: No executable found in job script", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
# Build job parameters - CLI args override script directives
|
||||
job_name = arg_job_name or directives.get("job-name") or directives.get("J")
|
||||
if not job_name:
|
||||
# Generate name from executable
|
||||
job_name = os.path.basename(executable).replace(".", "_")
|
||||
|
||||
nodes = arg_nodes
|
||||
if "nodes" in directives:
|
||||
nodes = int(directives["nodes"])
|
||||
if "N" in directives:
|
||||
nodes = int(directives["N"])
|
||||
if arg_nodes != 1: # CLI override
|
||||
nodes = arg_nodes
|
||||
|
||||
ntasks = arg_ntasks
|
||||
if "ntasks-per-node" in directives:
|
||||
ntasks = int(directives["ntasks-per-node"])
|
||||
if arg_ntasks != 1: # CLI override
|
||||
ntasks = arg_ntasks
|
||||
|
||||
workdir = arg_chdir or directives.get("chdir") or directives.get("D")
|
||||
if not workdir:
|
||||
workdir = os.getcwd()
|
||||
|
||||
hosts = arg_hosts or directives.get("hosts") or ""
|
||||
|
||||
# Resolve executable to absolute path
|
||||
if not os.path.isabs(executable):
|
||||
executable = os.path.abspath(os.path.join(workdir, executable))
|
||||
|
||||
# Submit job via API using query parameters
|
||||
from urllib.parse import urlencode
|
||||
|
||||
params = {
|
||||
"simulation_name": job_name,
|
||||
"flash_executable_path": executable,
|
||||
"parameter_file_path": "", # FLASH par file - use default
|
||||
"working_directory": workdir,
|
||||
"ranks_per_node": str(ntasks),
|
||||
"min_nodes": str(nodes),
|
||||
"hosts": hosts,
|
||||
}
|
||||
|
||||
query_string = urlencode(params)
|
||||
result = api_request("POST", f"/flash/launch?{query_string}")
|
||||
|
||||
# Print job submission confirmation
|
||||
if isinstance(result, dict):
|
||||
instance_id_val = result.get("instance_id")
|
||||
|
||||
if instance_id_val is not None:
|
||||
job_id = truncate_id(str(instance_id_val)) # pyright: ignore[reportAny]
|
||||
print(f"Submitted batch job {job_id}")
|
||||
else:
|
||||
# Instance created asynchronously - user should check squeue
|
||||
print("Job submitted successfully")
|
||||
print("Use 'exo squeue' to view job ID")
|
||||
else:
|
||||
print("Job submitted successfully")
|
||||
print("Use 'exo squeue' to view job ID")
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main(sys.argv[1:]))
|
||||
@@ -0,0 +1,95 @@
|
||||
"""scancel - Cancel jobs in the Exo queue.
|
||||
|
||||
Usage:
|
||||
exo scancel <jobid> [<jobid>...]
|
||||
|
||||
Arguments:
|
||||
jobid Job ID (or prefix) to cancel. Can specify multiple.
|
||||
|
||||
Examples:
|
||||
exo scancel abc123 # Cancel job starting with abc123
|
||||
exo scancel abc123 def456 # Cancel multiple jobs
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
from typing import Any, cast
|
||||
|
||||
from exo.cli.common import api_request, truncate_id
|
||||
|
||||
|
||||
def main(args: list[str]) -> int:
|
||||
"""Main entry point for scancel command."""
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="exo scancel",
|
||||
description="Cancel jobs in the Exo queue",
|
||||
)
|
||||
parser.add_argument(
|
||||
"jobids",
|
||||
nargs="+",
|
||||
help="Job ID(s) to cancel",
|
||||
)
|
||||
|
||||
parsed = parser.parse_args(args)
|
||||
jobids: list[str] = parsed.jobids # pyright: ignore[reportAny]
|
||||
|
||||
# Fetch current jobs to resolve partial IDs
|
||||
result = api_request("GET", "/flash/instances")
|
||||
if isinstance(result, list):
|
||||
instances = cast(list[dict[str, Any]], result)
|
||||
else:
|
||||
instances = cast(list[dict[str, Any]], result.get("instances", []))
|
||||
|
||||
# Build lookup of full IDs
|
||||
id_map: dict[str, str] = {}
|
||||
for inst in instances:
|
||||
iid = inst.get("instance_id", "") # pyright: ignore[reportAny]
|
||||
full_id = str(iid) if iid else "" # pyright: ignore[reportAny]
|
||||
if full_id:
|
||||
# Map both full ID and truncated versions
|
||||
normalized = full_id.replace("-", "").lower()
|
||||
id_map[normalized] = full_id
|
||||
# Also map prefixes
|
||||
for length in range(4, len(normalized) + 1):
|
||||
prefix = normalized[:length]
|
||||
if prefix not in id_map:
|
||||
id_map[prefix] = full_id
|
||||
|
||||
cancelled = 0
|
||||
errors = 0
|
||||
|
||||
for jobid in jobids:
|
||||
search = jobid.lower().replace("-", "")
|
||||
|
||||
# Find matching full ID
|
||||
full_id = id_map.get(search)
|
||||
if not full_id:
|
||||
# Try prefix match
|
||||
matches = [fid for key, fid in id_map.items() if key.startswith(search)]
|
||||
if len(matches) == 1:
|
||||
full_id = matches[0]
|
||||
elif len(matches) > 1:
|
||||
print(f"Ambiguous job ID: {jobid} matches multiple jobs")
|
||||
errors += 1
|
||||
continue
|
||||
else:
|
||||
print(f"Job not found: {jobid}")
|
||||
errors += 1
|
||||
continue
|
||||
|
||||
# Cancel the job
|
||||
try:
|
||||
api_request("DELETE", f"/flash/{full_id}")
|
||||
print(f"Job {truncate_id(full_id)} cancelled")
|
||||
cancelled += 1
|
||||
except SystemExit:
|
||||
print(f"Failed to cancel job {truncate_id(full_id)}")
|
||||
errors += 1
|
||||
|
||||
if errors > 0 and cancelled == 0:
|
||||
return 1
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main(sys.argv[1:]))
|
||||
@@ -0,0 +1,165 @@
|
||||
"""squeue - View the Exo job queue.
|
||||
|
||||
Usage:
|
||||
exo squeue [options]
|
||||
|
||||
Options:
|
||||
-l, --long Show detailed output
|
||||
-j, --job ID Show only this job
|
||||
|
||||
Output columns:
|
||||
JOBID - Job identifier (truncated UUID)
|
||||
NAME - Job name
|
||||
NODES - Number of nodes
|
||||
STATE - Job state (PENDING, RUNNING, FAILED, etc.)
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
from typing import Any, cast
|
||||
|
||||
from exo.cli.common import api_request, format_table, truncate_id
|
||||
|
||||
# Map Exo runner statuses to SLURM-like states
|
||||
STATUS_MAP: dict[str, str] = {
|
||||
"RunnerIdle": "PENDING",
|
||||
"RunnerConnecting": "CONFIGURING",
|
||||
"RunnerConnected": "CONFIGURING",
|
||||
"RunnerLoading": "CONFIGURING",
|
||||
"RunnerLoaded": "CONFIGURING",
|
||||
"RunnerWarmingUp": "CONFIGURING",
|
||||
"RunnerReady": "COMPLETING",
|
||||
"RunnerRunning": "RUNNING",
|
||||
"RunnerShuttingDown": "COMPLETING",
|
||||
"RunnerShutdown": "COMPLETED",
|
||||
"RunnerFailed": "FAILED",
|
||||
}
|
||||
|
||||
|
||||
def get_job_state(runner_statuses: dict[str, Any]) -> str:
|
||||
"""Determine overall job state from runner statuses."""
|
||||
if not runner_statuses:
|
||||
return "PENDING"
|
||||
|
||||
states: set[str] = set()
|
||||
for status_val in runner_statuses.values(): # pyright: ignore[reportAny]
|
||||
if isinstance(status_val, dict):
|
||||
# Extract status type from discriminated union
|
||||
type_val = status_val.get("type", "RunnerIdle") # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
|
||||
status_type = str(type_val) if type_val else "RunnerIdle" # pyright: ignore[reportUnknownArgumentType]
|
||||
elif isinstance(status_val, str):
|
||||
status_type = status_val
|
||||
else:
|
||||
status_type = "RunnerIdle"
|
||||
# Strip parentheses from status strings like "RunnerRunning()"
|
||||
if status_type.endswith("()"):
|
||||
status_type = status_type[:-2]
|
||||
states.add(STATUS_MAP.get(status_type, "UNKNOWN"))
|
||||
|
||||
# Priority order for overall state
|
||||
if "FAILED" in states:
|
||||
return "FAILED"
|
||||
if "RUNNING" in states:
|
||||
return "RUNNING"
|
||||
if "CONFIGURING" in states:
|
||||
return "CONFIGURING"
|
||||
if "COMPLETING" in states:
|
||||
return "COMPLETING"
|
||||
if "COMPLETED" in states:
|
||||
return "COMPLETED"
|
||||
if "PENDING" in states:
|
||||
return "PENDING"
|
||||
return "UNKNOWN"
|
||||
|
||||
|
||||
def main(args: list[str]) -> int:
|
||||
"""Main entry point for squeue command."""
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="exo squeue",
|
||||
description="View the Exo job queue",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-l",
|
||||
"--long",
|
||||
action="store_true",
|
||||
help="Show detailed output",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-j",
|
||||
"--job",
|
||||
help="Show only this job ID",
|
||||
)
|
||||
|
||||
parsed = parser.parse_args(args)
|
||||
|
||||
# Extract typed values
|
||||
long_format: bool = parsed.long # pyright: ignore[reportAny]
|
||||
job_filter: str | None = parsed.job # pyright: ignore[reportAny]
|
||||
|
||||
# Fetch jobs from API - returns list directly
|
||||
result = api_request("GET", "/flash/instances")
|
||||
# API returns list directly, not {"instances": [...]}
|
||||
if isinstance(result, list):
|
||||
instances = cast(list[dict[str, Any]], result)
|
||||
else:
|
||||
instances = cast(list[dict[str, Any]], result.get("instances", []))
|
||||
|
||||
if not instances:
|
||||
# No jobs - just print header
|
||||
if long_format:
|
||||
print("JOBID NAME NODES RANKS STATE WORKDIR")
|
||||
else:
|
||||
print("JOBID NAME NODES STATE")
|
||||
return 0
|
||||
|
||||
# Filter by job ID if specified
|
||||
if job_filter:
|
||||
search = job_filter.lower()
|
||||
filtered: list[dict[str, Any]] = []
|
||||
for i in instances:
|
||||
iid = i.get("instance_id", "") # pyright: ignore[reportAny]
|
||||
if search in str(iid).lower().replace("-", ""): # pyright: ignore[reportAny]
|
||||
filtered.append(i)
|
||||
instances = filtered
|
||||
|
||||
# Build table
|
||||
rows: list[list[str]] = []
|
||||
|
||||
if long_format:
|
||||
headers = ["JOBID", "NAME", "NODES", "RANKS", "STATE", "WORKDIR"]
|
||||
for inst in instances:
|
||||
iid_val = inst.get("instance_id", "") # pyright: ignore[reportAny]
|
||||
instance_id = str(iid_val) if iid_val else "" # pyright: ignore[reportAny]
|
||||
job_id = truncate_id(instance_id, 12)
|
||||
name_val = inst.get("simulation_name", "") # pyright: ignore[reportAny]
|
||||
name = (str(name_val) if name_val else "")[:15] # pyright: ignore[reportAny]
|
||||
runner_statuses = cast(dict[str, Any], inst.get("runner_statuses", {}))
|
||||
nodes = str(len(runner_statuses))
|
||||
ranks_val = inst.get("total_ranks", 0) # pyright: ignore[reportAny]
|
||||
ranks = str(ranks_val) if ranks_val else "0" # pyright: ignore[reportAny]
|
||||
state = get_job_state(runner_statuses)
|
||||
workdir_val = inst.get("working_directory", "") # pyright: ignore[reportAny]
|
||||
workdir = str(workdir_val) if workdir_val else "" # pyright: ignore[reportAny]
|
||||
# Truncate workdir for display
|
||||
if len(workdir) > 30:
|
||||
workdir = "..." + workdir[-27:]
|
||||
rows.append([job_id, name, nodes, ranks, state, workdir])
|
||||
else:
|
||||
headers = ["JOBID", "NAME", "NODES", "STATE"]
|
||||
for inst in instances:
|
||||
iid_val = inst.get("instance_id", "") # pyright: ignore[reportAny]
|
||||
instance_id = str(iid_val) if iid_val else "" # pyright: ignore[reportAny]
|
||||
job_id = truncate_id(instance_id, 8)
|
||||
name_val = inst.get("simulation_name", "") # pyright: ignore[reportAny]
|
||||
name = (str(name_val) if name_val else "")[:15] # pyright: ignore[reportAny]
|
||||
runner_statuses = cast(dict[str, Any], inst.get("runner_statuses", {}))
|
||||
nodes = str(len(runner_statuses))
|
||||
state = get_job_state(runner_statuses)
|
||||
rows.append([job_id, name, nodes, state])
|
||||
|
||||
print(format_table(headers, rows))
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main(sys.argv[1:]))
|
||||
@@ -195,6 +195,14 @@ class Node:
|
||||
|
||||
|
||||
def main():
|
||||
# Check for SLURM-compatible subcommands first
|
||||
import sys
|
||||
|
||||
if len(sys.argv) > 1 and sys.argv[1] in ("sbatch", "squeue", "scancel", "salloc"):
|
||||
from exo.cli import run_subcommand
|
||||
|
||||
sys.exit(run_subcommand(sys.argv[1], sys.argv[2:]))
|
||||
|
||||
args = Args.parse()
|
||||
soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
|
||||
resource.setrlimit(resource.RLIMIT_NOFILE, (max(soft, 65535), hard))
|
||||
|
||||
Reference in New Issue
Block a user