diff --git a/CLAUDE.md b/CLAUDE.md index f0a9631..7d1ef36 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -21,6 +21,7 @@ Nexus is a GPU job management system with a client-server architecture. It sched - **CLI Client** (`nexus.cli`): Command-line interface for users to submit jobs, view status, attach to sessions - **FastAPI Server** (`nexus.server`): Backend service that manages jobs, schedules work, and monitors system health - **Communication**: Client talks to server via REST API on localhost (default port configurable) +- **Code Separation**: CLI and server maintain separate implementations of shared utilities (e.g., `ids.py`) to ensure CLI can operate independently without server dependencies **Server Components:** - **Core** (`nexus.server.core`): Job management, database operations, configuration, context diff --git a/pyproject.toml b/pyproject.toml index e10e23c..4f18ea2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "nexusai" -version = "0.5.31" +version = "0.5.32" description = "Add your description here" readme = "README.md" requires-python = ">=3.9" diff --git a/src/nexus/cli/api_client.py b/src/nexus/cli/api_client.py index 1d234e0..9ceebf2 100644 --- a/src/nexus/cli/api_client.py +++ b/src/nexus/cli/api_client.py @@ -1,5 +1,6 @@ import functools import json +import time import typing as tp import requests @@ -42,15 +43,37 @@ def _print_error_response(response): def handle_api_errors(func): @functools.wraps(func) def wrapper(*args, **kwargs): - try: - return func(*args, **kwargs) - except SSHTunnelError as e: - print(colored("\nSSH Tunnel Error:", "red", attrs=["bold"])) - print(str(e)) - raise - except requests.exceptions.HTTPError as e: - _print_error_response(e.response) - raise + max_retries = 2 + + for attempt in range(max_retries): + try: + return func(*args, **kwargs) + except SSHTunnelError as e: + print(colored("\nSSH Tunnel Error:", "red", attrs=["bold"])) + print(str(e)) + raise + except (requests.exceptions.ConnectionError, requests.exceptions.Timeout) as e: + if attempt < max_retries - 1: + target_name = kwargs.get("target_name") + active_name, target_cfg = config.get_active_target(target_name) + if target_cfg and target_cfg.host not in ("localhost", "127.0.0.1"): + print( + colored( + f"\nConnection failed, recreating tunnel (attempt {attempt + 1}/{max_retries})...", + "yellow", + ) + ) + tunnel_manager._stop_control_master(active_name) + time.sleep(0.5) + continue + print(colored("\nConnection Error:", "red", attrs=["bold"])) + print(f"Failed to connect after {max_retries} attempts: {e}") + raise + except requests.exceptions.HTTPError as e: + _print_error_response(e.response) + raise + + return None return wrapper @@ -78,7 +101,7 @@ def check_api_connection(target_name: str | None = None) -> bool: @handle_api_errors def get_gpus(target_name: str | None = None) -> list[dict]: api_url = get_api_base_url(target_name) - response = requests.get(f"{api_url}/gpus") + response = requests.get(f"{api_url}/gpus", timeout=5) response.raise_for_status() return response.json() @@ -87,7 +110,7 @@ def get_gpus(target_name: str | None = None) -> list[dict]: def get_jobs(status: str | None = None, target_name: str | None = None) -> list[dict]: params = {"status": status} if status else {} api_url = get_api_base_url(target_name) - response = requests.get(f"{api_url}/jobs", params=params) + response = requests.get(f"{api_url}/jobs", params=params, timeout=5) response.raise_for_status() return response.json() @@ -95,7 +118,7 @@ def get_jobs(status: str | None = None, target_name: str | None = None) -> list[ @handle_api_errors def get_job(job_id: str, target_name: str | None = None) -> dict: api_url = get_api_base_url(target_name) - response = requests.get(f"{api_url}/jobs/{job_id}") + response = requests.get(f"{api_url}/jobs/{job_id}", timeout=5) response.raise_for_status() return response.json() @@ -106,7 +129,7 @@ def get_job_logs(job_id: str, last_n_lines: int | None = None, target_name: str if last_n_lines is not None: params["last_n_lines"] = last_n_lines api_url = get_api_base_url(target_name) - response = requests.get(f"{api_url}/jobs/{job_id}/logs", params=params) + response = requests.get(f"{api_url}/jobs/{job_id}/logs", params=params, timeout=10) response.raise_for_status() data = response.json() if "data" not in data: @@ -117,7 +140,7 @@ def get_job_logs(job_id: str, last_n_lines: int | None = None, target_name: str @handle_api_errors def get_server_status(target_name: str | None = None) -> dict: api_url = get_api_base_url(target_name) - response = requests.get(f"{api_url}/server/status") + response = requests.get(f"{api_url}/server/status", timeout=10) response.raise_for_status() return response.json() @@ -128,7 +151,7 @@ def get_detailed_health(refresh: bool = False, target_name: str | None = None) - if refresh: params["refresh"] = True api_url = get_api_base_url(target_name) - response = requests.get(f"{api_url}/health", params=params) + response = requests.get(f"{api_url}/health", params=params, timeout=10) response.raise_for_status() return response.json() @@ -136,7 +159,7 @@ def get_detailed_health(refresh: bool = False, target_name: str | None = None) - @handle_api_errors def check_artifact_by_sha(git_sha: str, target_name: str | None = None) -> tuple[bool, str | None]: api_url = get_api_base_url(target_name) - response = requests.get(f"{api_url}/artifacts/by-sha/{git_sha}") + response = requests.get(f"{api_url}/artifacts/by-sha/{git_sha}", timeout=5) response.raise_for_status() result = response.json() return result["exists"], result.get("artifact_id") @@ -149,7 +172,7 @@ def upload_artifact(data: bytes, git_sha: str | None = None, target_name: str | if git_sha: params["git_sha"] = git_sha api_url = get_api_base_url(target_name) - response = requests.post(f"{api_url}/artifacts", data=data, params=params, headers=headers) + response = requests.post(f"{api_url}/artifacts", data=data, params=params, headers=headers, timeout=30) response.raise_for_status() result = response.json() if "data" not in result: @@ -160,18 +183,24 @@ def upload_artifact(data: bytes, git_sha: str | None = None, target_name: str | @handle_api_errors def add_job(job_request: dict, target_name: str | None = None) -> dict: api_url = get_api_base_url(target_name) - response = requests.post(f"{api_url}/jobs", json=job_request) + response = requests.post(f"{api_url}/jobs", json=job_request, timeout=10) response.raise_for_status() return response.json() -def _process_job_batch(job_ids: list[str], method: str, endpoint_suffix: str, success_key: str, target_name: str | None) -> dict: +def _process_job_batch( + job_ids: list[str], + method: tp.Literal["POST", "DELETE"], + endpoint_suffix: str, + success_key: str, + target_name: str | None, +) -> dict: results = {success_key: [], "failed": []} api_url = get_api_base_url(target_name) request_fn = requests.post if method == "POST" else requests.delete for job_id in job_ids: try: - response = request_fn(f"{api_url}/jobs/{job_id}{endpoint_suffix}") + response = request_fn(f"{api_url}/jobs/{job_id}{endpoint_suffix}", timeout=5) if response.status_code == 204: results[success_key].append(job_id) else: @@ -211,7 +240,7 @@ def edit_job( update_data["git_tag"] = git_tag api_url = get_api_base_url(target_name) - response = requests.patch(f"{api_url}/jobs/{job_id}", json=update_data) + response = requests.patch(f"{api_url}/jobs/{job_id}", json=update_data, timeout=5) response.raise_for_status() return response.json() @@ -226,11 +255,11 @@ def manage_blacklist( for gpu_idx in gpu_indices: try: if action == "add": - response = requests.put(f"{api_url}/gpus/{gpu_idx}/blacklist") + response = requests.put(f"{api_url}/gpus/{gpu_idx}/blacklist", timeout=5) if response.ok: results["blacklisted"].append(gpu_idx) else: - response = requests.delete(f"{api_url}/gpus/{gpu_idx}/blacklist") + response = requests.delete(f"{api_url}/gpus/{gpu_idx}/blacklist", timeout=5) if response.ok: results["removed"].append(gpu_idx) diff --git a/src/nexus/cli/constants.py b/src/nexus/cli/constants.py index c71ad5b..11b5847 100644 --- a/src/nexus/cli/constants.py +++ b/src/nexus/cli/constants.py @@ -6,6 +6,12 @@ TERMINAL_STATUSES: frozenset[str] = frozenset([STATUS_COMPLETED, STATUS_FAILED, STATUS_KILLED]) +STATUS_ICONS: dict[str, str] = { + STATUS_COMPLETED: "✓", + STATUS_FAILED: "✗", + STATUS_KILLED: "🛑", +} + JOB_INIT_MAX_ATTEMPTS = 10 COMPLETED_JOB_LOG_TAIL_LINES = 5000 HISTORY_MAX_DISPLAY = 25 diff --git a/src/nexus/cli/jobs.py b/src/nexus/cli/jobs.py index 1d106fe..14d7383 100644 --- a/src/nexus/cli/jobs.py +++ b/src/nexus/cli/jobs.py @@ -1,6 +1,7 @@ import re import sys import time +import typing as tp from termcolor import colored @@ -8,7 +9,6 @@ from nexus.cli.config import IntegrationType, NotificationType from nexus.cli.constants import ( ATTACH_LOG_TAIL_LINES, - COMMAND_TRUNCATE_DEFAULT, COMMAND_TRUNCATE_QUEUE, COMMAND_TRUNCATE_SHORT, COMPLETED_JOB_LOG_TAIL_LINES, @@ -19,6 +19,7 @@ QUEUE_PREVIEW_COUNT, STATUS_COMPLETED, STATUS_FAILED, + STATUS_ICONS, STATUS_KILLED, STATUS_QUEUED, STATUS_RUNNING, @@ -35,7 +36,9 @@ def _build_job_info(job: dict, **extras) -> dict: return {**base_info, **extras} -def _validate_notifications(notifications: list[NotificationType], env_vars: dict[str, str]) -> list[NotificationType]: +def _validate_notifications( + notifications: list[NotificationType], env_vars: dict[str, str] +) -> list[NotificationType] | None: invalid = [n for n in notifications if any(env_vars.get(v) is None for v in config.REQUIRED_ENV_VARS.get(n, []))] if invalid: print(colored("\nWarning: Some notification types are missing required configuration:", "yellow")) @@ -43,7 +46,7 @@ def _validate_notifications(notifications: list[NotificationType], env_vars: dic print(f" {colored('•', 'yellow')} {notification_type}") if not utils.ask_yes_no("Continue with remaining notification types?"): utils.print_cancellation() - return [] + return None return [n for n in notifications if n not in invalid] return notifications @@ -100,19 +103,18 @@ def _resolve_job_target( jobs.extend(api_client.get_jobs(status, target_name=target_name)) if not jobs: - status_msg = "running " if require_running else "" - print(colored(f"No {status_msg}jobs found".replace(" ", " "), "yellow")) + if require_running: + print(colored("No running jobs found", "yellow")) + else: + print(colored("No jobs found", "yellow")) return None latest_job = utils.get_latest_user_job(jobs, user) if not latest_job: - status_msg = "running " if require_running else "" - print( - colored( - f"No {status_msg}jobs with valid start times found for user '{user}'".replace(" ", " "), - "yellow", - ) - ) + if require_running: + print(colored(f"No running jobs with valid start times found for user '{user}'", "yellow")) + else: + print(colored(f"No jobs with valid start times found for user '{user}'", "yellow")) return None return latest_job["id"] @@ -136,7 +138,9 @@ def _resolve_job_target( return target -def _format_gpu_status_part(gpus: list[dict], label: str, color: str, filter_fn) -> str | None: +def _format_gpu_status_part( + gpus: list[dict], label: str, color: str, filter_fn: tp.Callable[[dict], bool] +) -> str | None: gpu_list = [str(g["index"]) for g in gpus if filter_fn(g)] if not gpu_list: return None @@ -145,6 +149,67 @@ def _format_gpu_status_part(gpus: list[dict], label: str, color: str, filter_fn) return f"{count} {label} {gpu_str}" +def _build_notification_lists( + cfg: config.NexusCliConfig, + notification_types: list[NotificationType] | None, + integration_types: list[IntegrationType] | None, + silent: bool, +) -> tuple[list[NotificationType], list[IntegrationType]]: + notifications = [] if silent else list(cfg.default_notifications) + integrations = list(cfg.default_integrations) + if notification_types: + notifications.extend(n for n in notification_types if n not in notifications) + if integration_types: + integrations.extend(i for i in integration_types if i not in integrations) + return notifications, integrations + + +def _load_jobrc() -> str | None: + jobrc_path = setup.get_jobrc_path() + if jobrc_path.exists(): + with open(jobrc_path) as f: + return f.read() + return None + + +def _build_job_request( + job_id: str, + command: str, + user: str, + git_ctx: utils.GitArtifactContext, + num_gpus: int, + gpu_idxs: list[int] | None, + priority: int, + notifications: list[NotificationType], + integrations: list[IntegrationType], + env_vars: dict[str, str], + jobrc_content: str | None, + run_immediately: bool, + force: bool, + git_tag_pushed: bool, +) -> dict: + gpus_count = len(gpu_idxs) if gpu_idxs else num_gpus + return { + "job_id": job_id, + "command": command, + "user": user, + "artifact_id": git_ctx.artifact_id, + "git_repo_url": git_ctx.git_repo_url, + "git_branch": git_ctx.branch_name, + "git_tag": git_ctx.git_tag, + "num_gpus": gpus_count, + "priority": priority, + "integrations": integrations, + "notifications": notifications, + "env": env_vars, + "jobrc": jobrc_content, + "gpu_idxs": gpu_idxs, + "run_immediately": run_immediately, + "ignore_blacklist": force, + "git_tag_pushed": git_tag_pushed, + } + + def run_job( cfg: config.NexusCliConfig, commands: list[str], @@ -194,19 +259,7 @@ def run_job( return user = cfg.user or "anonymous" - - notifications = [] if silent else list(cfg.default_notifications) - integrations = list(cfg.default_integrations) - - if notification_types: - for notification_type in notification_types: - if notification_type not in notifications: - notifications.append(notification_type) - - if integration_types: - for integration_type in integration_types: - if integration_type not in integrations: - integrations.append(integration_type) + notifications, integrations = _build_notification_lists(cfg, notification_types, integration_types, silent) git_ctx = None try: @@ -214,36 +267,27 @@ def run_job( job_env_vars = _load_and_merge_env() if notification_types or cfg.default_notifications: notifications = _validate_notifications(notifications, job_env_vars) - if not notifications: + if notifications is None: return - gpus_count = len(gpu_idxs) if gpu_idxs else num_gpus - - jobrc_content = None - jobrc_path = setup.get_jobrc_path() - if jobrc_path.exists(): - with open(jobrc_path) as f: - jobrc_content = f.read() - - job_request = { - "job_id": git_ctx.job_id, - "command": command, - "user": user, - "artifact_id": git_ctx.artifact_id, - "git_repo_url": git_ctx.git_repo_url, - "git_branch": git_ctx.branch_name, - "git_tag": git_ctx.git_tag, - "num_gpus": gpus_count, - "priority": 0, - "integrations": integrations, - "notifications": notifications, - "env": job_env_vars, - "jobrc": jobrc_content, - "gpu_idxs": gpu_idxs, - "run_immediately": True, - "ignore_blacklist": force, - "git_tag_pushed": bool(cfg.enable_git_tag_push and not local), - } + jobrc_content = _load_jobrc() + + job_request = _build_job_request( + job_id=git_ctx.job_id, + command=command, + user=user, + git_ctx=git_ctx, + num_gpus=num_gpus, + gpu_idxs=gpu_idxs, + priority=0, + notifications=notifications, + integrations=integrations, + env_vars=job_env_vars, + jobrc_content=jobrc_content, + run_immediately=True, + force=force, + git_tag_pushed=bool(cfg.enable_git_tag_push and not local), + ) result = api_client.add_job(job_request, target_name=target_name) if "id" not in result: @@ -263,11 +307,8 @@ def run_job( try: job = api_client.get_job(job_id, target_name=target_name) if job["status"] in TERMINAL_STATUSES: - print( - colored( - f"\nJob {job_id} {job['status']}", "red" if job["status"] != STATUS_COMPLETED else "green" - ) - ) + status_color = "green" if job["status"] == STATUS_COMPLETED else "red" + print(colored(f"\nJob {job_id} {job['status']}", status_color)) view_logs(cfg, target=job_id, target_name=target_name) return if job["status"] == STATUS_RUNNING and job.get("screen_session_name"): @@ -346,60 +387,38 @@ def add_jobs( return user = cfg.user or "anonymous" - - notifications = [] if silent else list(cfg.default_notifications) - integrations = list(cfg.default_integrations) - - if notification_types: - for notification_type in notification_types: - if notification_type not in notifications: - notifications.append(notification_type) - - if integration_types: - for integration_type in integration_types: - if integration_type not in integrations: - integrations.append(integration_type) + notifications, integrations = _build_notification_lists(cfg, notification_types, integration_types, silent) env_vars = _load_and_merge_env() if notification_types or cfg.default_notifications: notifications = _validate_notifications(notifications, env_vars) - if not notifications: + if notifications is None: return git_ctx = None try: git_ctx = utils.prepare_git_artifact(enable_git_tag_push=False, target_name=target_name) - jobrc_content = None - jobrc_path = setup.get_jobrc_path() - if jobrc_path.exists(): - with open(jobrc_path) as f: - jobrc_content = f.read() + jobrc_content = _load_jobrc() created_jobs = [] - job_env_vars = dict(env_vars) - gpus_count = len(gpu_idxs) if gpu_idxs else num_gpus for cmd in expanded_commands: queued_job_id = utils.generate_job_id() - job_request = { - "job_id": queued_job_id, - "command": cmd, - "user": user, - "artifact_id": git_ctx.artifact_id, - "git_repo_url": git_ctx.git_repo_url, - "git_branch": git_ctx.branch_name, - "git_tag": git_ctx.git_tag, - "num_gpus": gpus_count, - "priority": priority, - "integrations": integrations, - "notifications": notifications, - "env": job_env_vars, - "jobrc": jobrc_content, - "gpu_idxs": gpu_idxs, - "run_immediately": False, - "ignore_blacklist": force, - "git_tag_pushed": False, - } - + job_request = _build_job_request( + job_id=queued_job_id, + command=cmd, + user=user, + git_ctx=git_ctx, + num_gpus=num_gpus, + gpu_idxs=gpu_idxs, + priority=priority, + notifications=notifications, + integrations=integrations, + env_vars=env_vars, + jobrc_content=jobrc_content, + run_immediately=False, + force=force, + git_tag_pushed=False, + ) result = api_client.add_job(job_request, target_name=target_name) created_jobs.append(result) @@ -407,9 +426,8 @@ def add_jobs( for job in created_jobs: priority_str = utils.format_priority_str(priority) gpus_str = utils.format_gpu_info(gpu_idxs, num_gpus, style="parens") if num_gpus > 0 or cpu else "" - print( - f" {colored('•', 'green')} Job {colored(job['id'], 'magenta')}: {job['command']}{priority_str}{gpus_str}" - ) + job_id_colored = colored(job["id"], "magenta") + print(f" {colored('•', 'green')} Job {job_id_colored}: {job['command']}{priority_str}{gpus_str}") finally: if git_ctx: @@ -490,15 +508,7 @@ def get_sort_timestamp(job: dict) -> float: runtime = utils.calculate_runtime(job) started_time = utils.format_timestamp(job.get("started_at")) status_color = utils.get_status_color(job["status"]) - status_icon = ( - "✓" - if job["status"] == STATUS_COMPLETED - else "✗" - if job["status"] == STATUS_FAILED - else "🛑" - if job["status"] == STATUS_KILLED - else "?" - ) + status_icon = STATUS_ICONS.get(job["status"], "?") status_str = colored(f"{status_icon} {job['status'].upper()}", status_color) command = utils.truncate_command(job["command"]) @@ -512,9 +522,9 @@ def get_sort_timestamp(job: dict) -> float: total_jobs = len(jobs) if total_jobs > HISTORY_MAX_DISPLAY: - print( - f"\n{colored(f'Showing most recent {HISTORY_MAX_DISPLAY} of', 'blue', attrs=['bold'])} {colored(str(total_jobs), 'cyan')}" - ) + msg_part1 = colored(f"Showing most recent {HISTORY_MAX_DISPLAY} of", "blue", attrs=["bold"]) + msg_part2 = colored(str(total_jobs), "cyan") + print(f"\n{msg_part1} {msg_part2}") completed_count = sum(1 for j in jobs if j["status"] == STATUS_COMPLETED) failed_count = sum(1 for j in jobs if j["status"] == STATUS_FAILED) @@ -559,9 +569,7 @@ def kill_jobs(targets: list[str] | None = None, bypass_confirm: bool = False, ta runtime_str = utils.format_runtime(runtime) if runtime else "N/A" print(colored(f"Latest job found: {job_id}", "blue")) - print( - f" {colored('•', 'blue')} Command: {utils.truncate_command(latest_job['command'])}" - ) + print(f" {colored('•', 'blue')} Command: {utils.truncate_command(latest_job['command'])}") print(f" {colored('•', 'blue')} Runtime: {colored(runtime_str, 'cyan')}") if latest_job.get("user"): print(f" {colored('•', 'blue')} User: {colored(latest_job['user'], 'cyan')}") @@ -586,7 +594,8 @@ def kill_jobs(targets: list[str] | None = None, bypass_confirm: bool = False, ta job_match = next((j for j in running_jobs if j["id"] == job_id), None) runtime = utils.calculate_runtime(job_match) if job_match else "" job_info = {"id": job_id, "command": "", "user": ""} if not job_match else job_match - jobs_info.append(_build_job_info(job_info, gpu_idx=gpu_idx, runtime=utils.format_runtime(runtime) if runtime else "")) + runtime_str = utils.format_runtime(runtime) if runtime else "" + jobs_info.append(_build_job_info(job_info, gpu_idx=gpu_idx, runtime=runtime_str)) if job_ids: running_jobs = api_client.get_jobs(STATUS_RUNNING, target_name=target_name) @@ -596,7 +605,8 @@ def kill_jobs(targets: list[str] | None = None, bypass_confirm: bool = False, ta j = next(j for j in running_jobs if j["id"] == pattern) jobs_to_kill.add(j["id"]) runtime = utils.calculate_runtime(j) - jobs_info.append(_build_job_info(j, runtime=utils.format_runtime(runtime), gpu_idx=j.get("gpu_idx"))) + runtime_str = utils.format_runtime(runtime) + jobs_info.append(_build_job_info(j, runtime=runtime_str, gpu_idx=j.get("gpu_idx"))) else: try: regex = re.compile(pattern) @@ -604,7 +614,8 @@ def kill_jobs(targets: list[str] | None = None, bypass_confirm: bool = False, ta for m in matched: jobs_to_kill.add(m["id"]) runtime = utils.calculate_runtime(m) - jobs_info.append(_build_job_info(m, runtime=utils.format_runtime(runtime), gpu_idx=m.get("gpu_idx"))) + runtime_str = utils.format_runtime(runtime) + jobs_info.append(_build_job_info(m, runtime=runtime_str, gpu_idx=m.get("gpu_idx"))) except re.error as e: print(colored(f"Invalid regex pattern '{pattern}': {e}", "red")) @@ -628,9 +639,8 @@ def kill_jobs(targets: list[str] | None = None, bypass_confirm: bool = False, ta if info: user_str = f" (User: {info['user']})" if info["user"] else "" runtime_str = f" (Runtime: {info['runtime']})" if info["runtime"] else "" - print( - f" {colored('•', 'green')} Successfully killed job {colored(job_id, 'magenta')}{user_str}{runtime_str}" - ) + job_id_colored = colored(job_id, "magenta") + print(f" {colored('•', 'green')} Successfully killed job {job_id_colored}{user_str}{runtime_str}") else: print(f" {colored('•', 'green')} Successfully killed job {colored(job_id, 'magenta')}") @@ -663,7 +673,8 @@ def remove_jobs(job_ids: list[str], bypass_confirm: bool = False, target_name: s if m["id"] not in jobs_to_remove: jobs_to_remove.add(m["id"]) created_time = utils.format_timestamp(m.get("created_at")) - jobs_info.append(_build_job_info(m, queue_time=created_time, priority=m.get("priority", 0))) + job_priority = m.get("priority", 0) + jobs_info.append(_build_job_info(m, queue_time=created_time, priority=job_priority)) except re.error as e: print(colored(f"Invalid regex pattern '{pattern}': {e}", "red")) @@ -689,9 +700,8 @@ def remove_jobs(job_ids: list[str], bypass_confirm: bool = False, target_name: s if info: user_str = f" (User: {info['user']})" if info["user"] else "" queue_str = f" (Queued: {info['queue_time']})" if info["queue_time"] else "" - print( - f" {colored('•', 'green')} Successfully removed job {colored(job_id, 'magenta')}{user_str}{queue_str}" - ) + job_id_colored = colored(job_id, "magenta") + print(f" {colored('•', 'green')} Successfully removed job {job_id_colored}{user_str}{queue_str}") else: print(f" {colored('•', 'green')} Successfully removed job {colored(job_id, 'magenta')}") @@ -809,7 +819,12 @@ def show_health(refresh: bool = False, target_name: str | None = None) -> None: print(f" {colored('•', 'blue')} Download Speed: {colored(f'{download_speed:.1f} Mbps', 'cyan')}") print(f" {colored('•', 'blue')} Upload Speed: {colored(f'{upload_speed:.1f} Mbps', 'cyan')}") - ping_color = "green" if ping < PING_THRESHOLD_GOOD else "yellow" if ping < PING_THRESHOLD_WARNING else "red" + if ping < PING_THRESHOLD_GOOD: + ping_color = "green" + elif ping < PING_THRESHOLD_WARNING: + ping_color = "yellow" + else: + ping_color = "red" print(f" {colored('•', 'blue')} Ping: {colored(f'{ping:.1f} ms', ping_color)}") except Exception as e: @@ -842,10 +857,10 @@ def edit_job_command( print(f" {colored('•', 'blue')} GPUs: {colored(str(job['num_gpus']), 'cyan')}") print(f"\n{colored('Will edit to:', 'blue', attrs=['bold'])}") - print(f" {colored('•', 'blue')} Command: {colored(command if command is not None else 'unchanged', 'white')}") - print( - f" {colored('•', 'blue')} Priority: {colored(str(priority) if priority is not None else 'unchanged', 'cyan')}" - ) + cmd_display = command if command is not None else "unchanged" + print(f" {colored('•', 'blue')} Command: {colored(cmd_display, 'white')}") + priority_display = str(priority) if priority is not None else "unchanged" + print(f" {colored('•', 'blue')} Priority: {colored(priority_display, 'cyan')}") print( f" {colored('•', 'blue')} GPUs: {colored(str(num_gpus) if num_gpus is not None else 'unchanged', 'cyan')}" ) @@ -1071,9 +1086,9 @@ def print_status(target_name: str | None = None) -> None: command = utils.truncate_command(job.get("command", "")) - print( - f" {colored('•', 'white')} {colored(job['id'], 'magenta')} ({resource_str}) - {colored(runtime_str, 'cyan')}" - ) + job_id_colored = colored(job["id"], "magenta") + runtime_colored = colored(runtime_str, "cyan") + print(f" {colored('•', 'white')} {job_id_colored} ({resource_str}) - {runtime_colored}") print(f" {colored(command, 'white', attrs=['bold'])}") if job.get("wandb_url"): print(f" W&B: {colored(job['wandb_url'], 'yellow')}") @@ -1207,7 +1222,9 @@ def attach_to_job(cfg: config.NexusCliConfig, target: str | None = None, target_ runtime_str = utils.format_runtime(runtime) if runtime else "N/A" print(colored(f"Runtime: {runtime_str}", "cyan")) - logs = api_client.get_job_logs(job_id, last_n_lines=ATTACH_LOG_TAIL_LINES, target_name=target_name) or "" + logs = ( + api_client.get_job_logs(job_id, last_n_lines=ATTACH_LOG_TAIL_LINES, target_name=target_name) or "" + ) if logs: print("\n" + logs) else: diff --git a/src/nexus/cli/tunnel_manager.py b/src/nexus/cli/tunnel_manager.py index 17964ee..eb14cce 100644 --- a/src/nexus/cli/tunnel_manager.py +++ b/src/nexus/cli/tunnel_manager.py @@ -109,7 +109,7 @@ def _check_control_socket(target_name: str) -> bool: ["ssh", "-S", str(socket_path), "-O", "check", "dummy"], capture_output=True, text=True, - timeout=5, + timeout=2, ) return result.returncode == 0 except (subprocess.TimeoutExpired, OSError): @@ -165,9 +165,11 @@ def _start_control_master(target_name: str, target_cfg: config.TargetConfig) -> "-o", "ConnectTimeout=10", "-o", - "ServerAliveInterval=60", + "ServerAliveInterval=15", "-o", - "ServerAliveCountMax=3", + "ServerAliveCountMax=2", + "-o", + "TCPKeepAlive=yes", "-o", "ExitOnForwardFailure=yes", "-o", @@ -224,7 +226,7 @@ def _get_tunnel_port(target_name: str) -> int | None: _stop_control_master(target_name) return None - if not _wait_for_tunnel(local_port, timeout=1.0): + if not _wait_for_tunnel(local_port, timeout=0.5): _stop_control_master(target_name) return None diff --git a/src/nexus/cli/utils.py b/src/nexus/cli/utils.py index 55cab5b..3a6fb0c 100644 --- a/src/nexus/cli/utils.py +++ b/src/nexus/cli/utils.py @@ -11,9 +11,9 @@ from termcolor import colored +from nexus.cli.constants import TERMINAL_STATUSES from nexus.cli.ids import generate_job_id -# Types Color = tp.Literal["grey", "red", "green", "yellow", "blue", "magenta", "cyan", "white"] Attribute = tp.Literal["bold", "dark", "underline", "blink", "reverse", "concealed"] @@ -53,15 +53,19 @@ def print_job_field(label: str, value: str | int, value_color: Color = "cyan") - print(f" {colored('•', 'blue')} {label}: {colored(str(value), value_color)}") -def format_gpu_info(gpu_idxs: list[int] | None, num_gpus: int, style: tp.Literal["prefix", "parens", "inline"] = "prefix") -> str: +def format_gpu_info( + gpu_idxs: list[int] | None, + num_gpus: int, + style: tp.Literal["prefix", "parens", "inline"] = "prefix", +) -> str: if num_gpus == 0: return " (CPU)" if style == "parens" else " on CPU" if gpu_idxs: - gpu_str = ','.join(map(str, gpu_idxs)) - plural = 's' if len(gpu_idxs) > 1 else '' + gpu_str = ",".join(map(str, gpu_idxs)) + plural = "s" if len(gpu_idxs) > 1 else "" else: gpu_str = str(num_gpus) - plural = 's' if num_gpus > 1 else '' + plural = "s" if num_gpus > 1 else "" if style == "prefix": return f" on GPU{plural}: {colored(gpu_str, 'cyan')}" elif style == "parens": @@ -81,6 +85,8 @@ def format_gpu_info(gpu_idxs: list[int] | None, num_gpus: int, style: tp.Literal "unhealthy": "red", } +ELLIPSIS_LENGTH = 3 + def get_status_color(status: str) -> Color: return STATUS_COLOR_MAP.get(status, "white") @@ -91,7 +97,7 @@ def format_priority_str(priority: int) -> str: def truncate_command(command: str, max_length: int = 80) -> str: - return command if len(command) <= max_length else command[:max_length - 3] + "..." + return command if len(command) <= max_length else command[: max_length - ELLIPSIS_LENGTH] + "..." def print_cancellation() -> None: @@ -102,7 +108,7 @@ def get_latest_user_job(jobs: list[dict], user: str) -> dict | None: user_jobs = [j for j in jobs if j.get("user") == user and j.get("started_at") is not None] if not user_jobs: return None - return max(user_jobs, key=lambda x: x.get("started_at", 0)) + return max(user_jobs, key=lambda x: x["started_at"] or 0.0) def print_warning(message: str) -> None: @@ -366,12 +372,14 @@ def format_timestamp(timestamp: float | None) -> str: def calculate_runtime(job: dict) -> float: - if not job.get("started_at"): + started_at = job.get("started_at") + if not started_at: return 0.0 - if job.get("status") in ["completed", "failed", "killed"] and job.get("completed_at"): - return job["completed_at"] - job["started_at"] + completed_at = job.get("completed_at") + if job.get("status") in TERMINAL_STATUSES and completed_at: + return completed_at - started_at elif job.get("status") == "running": - return time.time() - job["started_at"] + return time.time() - started_at return 0.0 diff --git a/src/nexus/server/api/router.py b/src/nexus/server/api/router.py index 7f44842..c199877 100644 --- a/src/nexus/server/api/router.py +++ b/src/nexus/server/api/router.py @@ -110,9 +110,7 @@ async def create_job_endpoint( if gpu_idxs_list or (job_request.run_immediately and job_request.num_gpus > 0): running_jobs = db.list_jobs(conn=ctx.db, status=STATUS_RUNNING) blacklisted = db.list_blacklisted_gpus(conn=ctx.db) - all_gpus = gpu.get_gpus( - running_jobs=running_jobs, blacklisted_gpus=blacklisted, mock_gpus=ctx.config.mock_gpus - ) + all_gpus = gpu.get_gpus(running_jobs=running_jobs, blacklisted_gpus=blacklisted, mock_gpus=ctx.config.mock_gpus) if gpu_idxs_list: if len([g for g in all_gpus if g.index in gpu_idxs_list]) != len(gpu_idxs_list): diff --git a/src/nexus/server/api/scheduler.py b/src/nexus/server/api/scheduler.py index f41a0f7..d50b572 100644 --- a/src/nexus/server/api/scheduler.py +++ b/src/nexus/server/api/scheduler.py @@ -1,7 +1,6 @@ import asyncio import dataclasses as dc import datetime as dt -import typing as tp from nexus.server.core import context, db, job, exceptions as exc, schemas from nexus.server.core.schemas import STATUS_FAILED, STATUS_QUEUED, STATUS_RUNNING diff --git a/src/nexus/server/core/job.py b/src/nexus/server/core/job.py index 35bc0b1..4e67f01 100644 --- a/src/nexus/server/core/job.py +++ b/src/nexus/server/core/job.py @@ -439,11 +439,13 @@ async def async_end_job(_job: schemas.Job, killed: bool) -> schemas.Job: elif exit_code is None: updates.update({"status": STATUS_FAILED, "error_message": "Could not find exit code in log"}) else: - updates.update({ - "exit_code": exit_code, - "status": STATUS_COMPLETED if exit_code == 0 else STATUS_FAILED, - "error_message": None if exit_code == 0 else f"Job failed with exit code {exit_code}" - }) + updates.update( + { + "exit_code": exit_code, + "status": STATUS_COMPLETED if exit_code == 0 else STATUS_FAILED, + "error_message": None if exit_code == 0 else f"Job failed with exit code {exit_code}", + } + ) return dc.replace(_job, **updates) diff --git a/src/nexus/server/external/notifications.py b/src/nexus/server/external/notifications.py index b2c33d1..4b3ad47 100644 --- a/src/nexus/server/external/notifications.py +++ b/src/nexus/server/external/notifications.py @@ -29,36 +29,32 @@ class NotificationMessage(pyd.BaseModel): username: str = "Nexus" -def _get_discord_secrets(job: schemas.Job) -> tuple[str, str]: - webhook_url = job.env.get("DISCORD_WEBHOOK_URL") - if not webhook_url: - raise exc.NotificationError("Missing DISCORD_WEBHOOK_URL in job environment") +@tp.overload +def _require_env(job: schemas.Job, __key1: str, __key2: str, /) -> tuple[str, str]: ... - user_id = job.env.get("DISCORD_USER_ID") - if not user_id: - raise exc.NotificationError("Missing DISCORD_USER_ID in job environment") - return webhook_url, user_id +@tp.overload +def _require_env( + job: schemas.Job, __key1: str, __key2: str, __key3: str, __key4: str, / +) -> tuple[str, str, str, str]: ... -def _get_phone_secrets(job: schemas.Job) -> tuple[str, str, str, str]: - phone_number = job.env.get("PHONE_TO_NUMBER") - if not phone_number: - raise exc.NotificationError("Missing PHONE_TO_NUMBER in job environment") +def _require_env(job: schemas.Job, *keys: str) -> tuple[str, ...]: + values = [] + for key in keys: + value = job.env.get(key) + if not value: + raise exc.NotificationError(f"Missing {key} in job environment") + values.append(value) + return tuple(values) - twilio_account_sid = job.env.get("TWILIO_ACCOUNT_SID") - if not twilio_account_sid: - raise exc.NotificationError("Missing TWILIO_ACCOUNT_SID in job environment") - twilio_auth_token = job.env.get("TWILIO_AUTH_TOKEN") - if not twilio_auth_token: - raise exc.NotificationError("Missing TWILIO_AUTH_TOKEN in job environment") +def _get_discord_secrets(job: schemas.Job) -> tuple[str, str]: + return _require_env(job, "DISCORD_WEBHOOK_URL", "DISCORD_USER_ID") - twilio_from_number = job.env.get("TWILIO_FROM_NUMBER") - if not twilio_from_number: - raise exc.NotificationError("Missing TWILIO_FROM_NUMBER in job environment") - return phone_number, twilio_account_sid, twilio_auth_token, twilio_from_number +def _get_phone_secrets(job: schemas.Job) -> tuple[str, str, str, str]: + return _require_env(job, "PHONE_TO_NUMBER", "TWILIO_ACCOUNT_SID", "TWILIO_AUTH_TOKEN", "TWILIO_FROM_NUMBER") def _truncate_field_value(value: str, max_length: int = 1024) -> str: @@ -186,7 +182,7 @@ async def _make_phone_call(to_number: str, from_number: str, account_sid: str, a async def _send_phone_notification(job: schemas.Job, job_action: JobAction) -> None: - if job_action not in ["completed", "failed", "killed"]: + if job_action not in schemas.TERMINAL_STATUSES: return to_number, account_sid, auth_token, from_number = _get_phone_secrets(job) @@ -216,8 +212,8 @@ async def notify_job_action(_job: schemas.Job, action: JobAction) -> schemas.Job message_data = _format_job_message_for_notification(_job, action) webhook_url = _get_discord_secrets(_job)[0] - if action in ["completed", "failed", "killed"] and _job.dir: - if action in ["failed", "killed"]: + if action in schemas.TERMINAL_STATUSES and _job.dir: + if action in [schemas.STATUS_FAILED, schemas.STATUS_KILLED]: job_logs = await job.async_get_job_logs(_job.dir, last_n_lines=20) if job_logs: MAX_FIELD_LENGTH = 1024 diff --git a/src/nexus/server/external/system.py b/src/nexus/server/external/system.py index d8a1063..bacd8cd 100644 --- a/src/nexus/server/external/system.py +++ b/src/nexus/server/external/system.py @@ -61,9 +61,7 @@ def measure_disk_space(path: str = "/") -> DiskStats: def check_disk_space(path: str = "/", force_refresh: bool = False) -> DiskStats: - if force_refresh and "disk_space" in _cache: - del _cache["disk_space"] - + _clear_cache_if_refresh(force_refresh, "disk_space") return _get_cached(key="disk_space", default_factory=lambda: measure_disk_space(path), ttl=timedelta(minutes=30)) @@ -97,6 +95,12 @@ def measure_network_speed() -> NetworkStats: _cache: dict[str, CachedValue] = {} +def _clear_cache_if_refresh(force_refresh: bool, *keys: str) -> None: + if force_refresh: + for key in keys: + _cache.pop(key, None) + + def _get_cached(key: str, default_factory: tp.Callable[[], tp.Any], ttl: timedelta) -> tp.Any: now = datetime.now() cache_entry = _cache.get(key) @@ -110,9 +114,7 @@ def _get_cached(key: str, default_factory: tp.Callable[[], tp.Any], ttl: timedel def check_network_speed(force_refresh: bool = False) -> NetworkStats: - if force_refresh and "network_speed" in _cache: - del _cache["network_speed"] - + _clear_cache_if_refresh(force_refresh, "network_speed") return _get_cached(key="network_speed", default_factory=measure_network_speed, ttl=timedelta(minutes=60)) @@ -126,9 +128,7 @@ def measure_system_stats() -> SystemStats: def check_system_stats(force_refresh: bool = False) -> SystemStats: - if force_refresh and "system_stats" in _cache: - del _cache["system_stats"] - + _clear_cache_if_refresh(force_refresh, "system_stats") return _get_cached(key="system_stats", default_factory=measure_system_stats, ttl=timedelta(minutes=1)) @@ -143,7 +143,9 @@ def calculate_health_score( return min(30, disk_score) network_score = 0 if network_stats.ping < 9999: - network_score = 15 * max(0, min(1, (200 - network_stats.ping) / 150)) + 15 * min(1, network_stats.download_speed / 100) + network_score = 15 * max(0, min(1, (200 - network_stats.ping) / 150)) + 15 * min( + 1, network_stats.download_speed / 100 + ) system_score = 15 * (2 - (system_stats.cpu_percent + system_stats.memory_percent) / 100) return round(disk_score + network_score + system_score, 1) @@ -169,13 +171,5 @@ def _calculate_health_result() -> HealthCheckResult: def check_health(force_refresh: bool = False) -> HealthCheckResult: - if force_refresh and "health_result" in _cache: - del _cache["health_result"] - if "disk_space" in _cache: - del _cache["disk_space"] - if "network_speed" in _cache: - del _cache["network_speed"] - if "system_stats" in _cache: - del _cache["system_stats"] - + _clear_cache_if_refresh(force_refresh, "health_result", "disk_space", "network_speed", "system_stats") return _get_cached(key="health_result", default_factory=_calculate_health_result, ttl=timedelta(minutes=5)) diff --git a/src/nexus/server/utils/format.py b/src/nexus/server/utils/format.py index cff8070..4de047e 100644 --- a/src/nexus/server/utils/format.py +++ b/src/nexus/server/utils/format.py @@ -24,7 +24,7 @@ def format_timestamp(timestamp: float | None) -> str: def calculate_runtime(job: schemas.Job) -> float: if not job.started_at: return 0.0 - if job.status in ["completed", "failed", "killed"] and job.completed_at: + if job.status in schemas.TERMINAL_STATUSES and job.completed_at: return job.completed_at - job.started_at elif job.status == "running": return dt.datetime.now().timestamp() - job.started_at diff --git a/uv.lock b/uv.lock index 9fb97de..90f1378 100644 --- a/uv.lock +++ b/uv.lock @@ -852,7 +852,7 @@ wheels = [ [[package]] name = "nexusai" -version = "0.5.31" +version = "0.5.32" source = { editable = "." } dependencies = [ { name = "aiohttp" },