diff --git a/src/exo/cli/__init__.py b/src/exo/cli/__init__.py new file mode 100644 index 00000000..27d81cd8 --- /dev/null +++ b/src/exo/cli/__init__.py @@ -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 diff --git a/src/exo/cli/common.py b/src/exo/cli/common.py new file mode 100644 index 00000000..ca448be3 --- /dev/null +++ b/src/exo/cli/common.py @@ -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) diff --git a/src/exo/cli/salloc.py b/src/exo/cli/salloc.py new file mode 100644 index 00000000..e4845599 --- /dev/null +++ b/src/exo/cli/salloc.py @@ -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:])) diff --git a/src/exo/cli/sbatch.py b/src/exo/cli/sbatch.py new file mode 100644 index 00000000..ad653f8e --- /dev/null +++ b/src/exo/cli/sbatch.py @@ -0,0 +1,233 @@ +"""sbatch - Submit a batch job to Exo. + +Usage: + exo sbatch [options] + exo sbatch --job-name=NAME --nodes=N + +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:])) diff --git a/src/exo/cli/scancel.py b/src/exo/cli/scancel.py new file mode 100644 index 00000000..56cf1557 --- /dev/null +++ b/src/exo/cli/scancel.py @@ -0,0 +1,95 @@ +"""scancel - Cancel jobs in the Exo queue. + +Usage: + exo scancel [...] + +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:])) diff --git a/src/exo/cli/squeue.py b/src/exo/cli/squeue.py new file mode 100644 index 00000000..1b59ae6f --- /dev/null +++ b/src/exo/cli/squeue.py @@ -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:])) diff --git a/src/exo/main.py b/src/exo/main.py index 85bc095b..296c9435 100644 --- a/src/exo/main.py +++ b/src/exo/main.py @@ -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))