import copy import re import time import traceback from datetime import datetime, timezone from datetime import timedelta from typing import Optional, Tuple from config.sku_config import SKU_CONFIG from src.clients import tcms_client from src.clients.tidbcloud_client import ComponentType, get_nodespecs from src.database import db from src.services.cluster_utils import extract_cluster_ids_from_execution from src.services.report.metrics_fetcher import fetch_all_diagrams_with_cache from src.services.result_parser import parse_execution_result from src.services.workload_templates import ( canonicalize_workload_template_or_default, resolve_tcms_workflow_id, resolve_workflow_yaml_path, ) _NODESPEC_KEY_RE = re.compile(r"^\s*(\d+)C(\d+)G\s*$") def _process_submitting_task( task: dict, *, allow_immediate_create: bool, workflow_name_was_preexisting: bool, ) -> None: print( f"Processing SUBMITTING task {task['id']} " f"(Workflow: {task.get('workflow_execution_id')}, " f"Name: {task.get('workflow_execution_name')})..." ) if task.get("workflow_execution_id"): # Submission already recorded; convert to PENDING and let polling handle it. db.update_task_status_if_current( task["id"], "PENDING", expected_current_statuses=["SUBMITTING"], ) return workflow_name = task.get("workflow_execution_name") if not workflow_name: workload_template = canonicalize_workload_template_or_default( task.get("workload_template") ) # Backfill a deterministic name for recovery. workflow_name = ( f"SPU_auto_{workload_template.replace('-','_')}_" f"{task['id']}_{int(time.time())}" ) db.update_task_status_if_current( task["id"], "SUBMITTING", expected_current_statuses=["SUBMITTING"], workflow_name=workflow_name, ) task["workflow_execution_name"] = workflow_name if not allow_immediate_create: updated_at = task.get("updated_at") now = datetime.now(timezone.utc) if isinstance(updated_at, datetime): updated_at_dt = ( updated_at if updated_at.tzinfo is not None else updated_at.replace(tzinfo=timezone.utc) ) if (now - updated_at_dt) < timedelta(minutes=1): # Another scheduler instance may be creating; wait. print( f"Task {task['id']} SUBMITTING recently; wait before retrying create." ) return # Acquire a stale lease to avoid concurrent re-submission. cutoff = datetime.utcnow() - timedelta(minutes=1) if not db.acquire_submitting_lease_if_stale(task["id"], cutoff): return if (not allow_immediate_create) or workflow_name_was_preexisting: found = tcms_client.find_workflow_execution_by_name(workflow_name) if not found: # If the execution already finished quickly, it may not be in active statuses. found = tcms_client.find_workflow_execution_by_name( workflow_name, statuses=( "PENDING", "RUNNING", "CANCELLING", "CANCELLED", "SUCCESS", "FAILURE", "ERROR", ), ) if found and found.get("id") is not None: workflow_id = str(found["id"]) new_status = found.get("status") or "PENDING" updated = db.update_task_status_if_current( task["id"], new_status, workflow_id=workflow_id, expected_current_statuses=["SUBMITTING"], ) if not updated: db.attach_workflow_execution_id_if_missing(task["id"], workflow_id) return # Before creating, re-check DB state to avoid creating after user requested cancel. latest_task = db.get_task_by_id(task["id"]) if not latest_task: return if latest_task.get("workflow_execution_id"): # Another worker already recorded submission; let polling handle it. db.update_task_status_if_current( task["id"], "PENDING", expected_current_statuses=["SUBMITTING"], ) return if latest_task.get("status") != "SUBMITTING": print( f"Skip creating workflow for task {task['id']} because status changed to " f"{latest_task.get('status')}." ) return try: sku_config = _build_sku_config_from_task(task) except Exception as e: error_msg = f"Failed to build sku_config for task {task['id']}: {e}" print(error_msg) traceback.print_exc() db.save_result(task["id"], "FAILURE", error_msg=error_msg, parsed=False) db.update_task_status_if_current( task["id"], "FAILURE", expected_current_statuses=["SUBMITTING"], ) return yaml_path = _select_workflow_yaml(task) try: workflow_id = resolve_tcms_workflow_id(task.get("workload_template")) except Exception: raise ValueError(f"Failed to resolve workflow_id for task {task['id']}") response = tcms_client.create_workflow_execution( workflow_id=workflow_id, name=workflow_name, sku_config=sku_config, yaml_path=yaml_path, cloud_provider=task["cloud_provider"], region=task["region"], ) execution_id = None if response: execution_id = response.get("id") if not execution_id: print(f"Failed to submit task {task['id']} (SUBMITTING).") return workflow_id = str(execution_id) new_status = response.get("status") or "PENDING" updated = db.update_task_status_if_current( task["id"], new_status, workflow_id=workflow_id, expected_current_statuses=["SUBMITTING"], ) if not updated: db.attach_workflow_execution_id_if_missing(task["id"], workflow_id) def _save_cluster_metrics_for_task(task_id, raw_execution_response): """ Extract cluster IDs from execution response and save metrics to database. Uses fetch_all_diagrams_with_cache which automatically checks cache first, so it won't duplicate existing metrics. Safe to call multiple times. Args: task_id: Task ID raw_execution_response: Raw execution response from TCMS """ try: cluster_ids = extract_cluster_ids_from_execution(raw_execution_response) if cluster_ids: print(f"Found {len(cluster_ids)} cluster(s), fetching metrics...") for cluster_key, cluster_id in cluster_ids.items(): print(f" Fetching metrics for cluster {cluster_id} ({cluster_key})...") fetch_all_diagrams_with_cache(task_id, cluster_id) print( f"✓ Metrics saved for task {task_id} " f"clusters: {list(cluster_ids.values())}" ) else: print(f"No clusters found for task {task_id}, skipping metrics") except Exception as e: # Don't fail the task if metrics fetching fails print(f"Warning: Failed to save metrics for task {task_id}: {e}") traceback.print_exc() def _parse_nodespec_cpu_mem(spec: object) -> Optional[Tuple[int, int]]: if not isinstance(spec, str): return None s = spec.strip() if not s: return None m = _NODESPEC_KEY_RE.match(s) if m: return int(m.group(1)), int(m.group(2)) return None def _nodespec_display_name(cpu: int, mem_gib: int) -> str: return f"{cpu} vCPU, {mem_gib} GiB" def _build_sku_config_from_task(task): sku_config = copy.deepcopy(SKU_CONFIG) sku_config["cloud_provider"] = task["cloud_provider"] sku_config["region"] = task["region"] cluster_conf = sku_config.get("cluster", {}) if task.get("cluster_version"): cluster_conf["cluster_version"] = task["cluster_version"] if task.get("org_id"): cluster_conf["org_id"] = task["org_id"] if task.get("project_id"): cluster_conf["project_id"] = task["project_id"] def update_component( comp_key, spu_field, count_field, component_type, storage_type_field=None, storage_size_field=None, ): base_comp = cluster_conf.get(comp_key) spu = task.get(spu_field) count = task.get(count_field) or 0 if not spu or count <= 0: return region_id = task["region"] try: specs = get_nodespecs(region_id, component_type=component_type) except Exception as e: print( f"Warning: get_nodespecs failed for region={region_id!r} " f"component_type={component_type!r}: {e}" ) specs = [] display_name = None for spec in specs: if spec.get("nodeSpecKey") == spu: display_name = spec.get("displayName") break if display_name is None: desired = _parse_nodespec_cpu_mem(spu) if desired is None: raise ValueError( f"Unknown nodespec format for {comp_key}: {spu!r}. " f"Expected like '64C128G'." ) display_name = _nodespec_display_name(*desired) print( f"Warning: nodespec {spu!r} for {comp_key} not found in get_nodespecs; " f"using derived display name {display_name!r}." ) if isinstance(base_comp, dict): comp_conf = copy.deepcopy(base_comp) else: comp_conf = {} comp_conf["nodeProfileName"] = display_name comp_conf["replicas"] = count # Add storage configuration (only for TiKV and TiFlash) if storage_type_field and storage_size_field: storage_type = task.get(storage_type_field) storage_size = task.get(storage_size_field) if storage_type: comp_conf["storageType"] = storage_type if storage_size: comp_conf["storageSizeGi"] = storage_size cluster_conf[comp_key] = comp_conf update_component("tidb", "tidb_spu", "tidb_count", ComponentType.TIDB) update_component( "tikv", "tikv_spu", "tikv_count", ComponentType.TIKV, "tikv_storage_type", "tikv_storage_size_gi", ) update_component( "tiflash", "tiflash_spu", "tiflash_count", ComponentType.TIFLASH, "tiflash_storage_type", "tiflash_storage_size_gi", ) update_component("pd", "pd_spu", "pd_count", ComponentType.PD) update_component("tiproxy", "tiproxy_spu", "tiproxy_count", ComponentType.TIPROXY) sku_config["cluster"] = cluster_conf sku_config["workload_config"] = task.get("workload_config") or {} sku_config["user_email"] = task.get("user_email") return sku_config def _select_workflow_yaml(task): """Select workflow YAML path based on workload_template.""" try: return resolve_workflow_yaml_path(task.get("workload_template")) except Exception: raise ValueError(f"Failed to resolve workflow_yaml_path for task {task['id']}") def process_pending_tasks(): """Fetches and processes pending tasks.""" tasks = db.get_pending_tasks() if not tasks: return print(f"Found {len(tasks)} pending/running tasks.") for task in tasks: # Skip deleted tasks if task.get("deleted"): # TODO: If there's a corresponding result, mark it as deleted continue try: # Refresh latest task status to avoid acting on stale UI/scheduler state. latest_task = db.get_task_by_id(task["id"]) if not latest_task: continue task = latest_task if task["status"] == "CANCELLING": print( f"Processing CANCELLING task {task['id']} " f"(Workflow: {task['workflow_execution_id']})..." ) if not task["workflow_execution_id"]: if task.get("workflow_execution_name"): found = tcms_client.find_workflow_execution_by_name( task["workflow_execution_name"], statuses=( "PENDING", "RUNNING", "CANCELLING", "CANCELLED", "SUCCESS", "FAILURE", "ERROR", ), ) if found and found.get("id") is not None: db.attach_workflow_execution_id_if_missing( task["id"], str(found["id"]) ) # Refresh to continue cancel with workflow id. task = db.get_task_by_id(task["id"]) or task else: updated_at = task.get("updated_at") now = datetime.now(timezone.utc) if isinstance(updated_at, datetime): updated_at_dt = ( updated_at if updated_at.tzinfo is not None else updated_at.replace(tzinfo=timezone.utc) ) # If cancelling shortly after a submit attempt, keep waiting. if (now - updated_at_dt) < timedelta(minutes=2): print( f"Task {task['id']} is CANCELLING; workflow not found yet, " f"waiting for potential in-flight submission." ) continue db.update_task_status_if_current( task["id"], "CANCELLED", expected_current_statuses=["CANCELLING"], ) print( f"Task {task['id']} CANCELLED (workflow not found after grace)." ) continue else: db.update_task_status_if_current( task["id"], "CANCELLED", expected_current_statuses=["CANCELLING"], ) print(f"Task {task['id']} CANCELLED (no workflow name).") continue execution_status = tcms_client.get_execution_status( task["workflow_execution_id"] ) if execution_status: state = execution_status.get("status") print(f"Workflow State: {state}") if state == "CANCELLED": db.update_task_status_if_current( task["id"], "CANCELLED", expected_current_statuses=["CANCELLING"], ) print(f"Task {task['id']} CANCELLED.") continue if state == "SUCCESS": raw_execution_response = execution_status db.save_result( task["id"], "SUCCESS", result_data=None, raw_execution_response=raw_execution_response, parsed=False, ) db.update_task_status_if_current( task["id"], "SUCCESS", expected_current_statuses=["CANCELLING"], ) print(f"Task {task['id']} SUCCESS (while cancelling).") _save_cluster_metrics_for_task( task["id"], raw_execution_response ) continue if state == "FAILURE": error_msg = execution_status.get("output", "Unknown error") raw_execution_response = execution_status db.save_result( task["id"], "FAILURE", error_msg=error_msg, raw_execution_response=raw_execution_response, parsed=False, ) db.update_task_status_if_current( task["id"], "FAILURE", expected_current_statuses=["CANCELLING"], ) print(f"Task {task['id']} FAILURE (while cancelling).") _save_cluster_metrics_for_task( task["id"], raw_execution_response ) continue if state == "ERROR": error_msg = execution_status.get("output", "Unknown error") db.save_result( task["id"], "ERROR", error_msg=error_msg, raw_execution_response=execution_status, parsed=False, ) db.update_task_status_if_current( task["id"], "ERROR", expected_current_statuses=["CANCELLING"], ) print(f"Task {task['id']} ERROR (while cancelling).") continue if state == "CANCELLING": print( f"Workflow {task['workflow_execution_id']} is already cancelling; " f"skip sending another cancel request." ) continue cancelled = tcms_client.cancel_workflow_execution( str(task["workflow_execution_id"]) ) if cancelled: print( f"Cancel request sent for workflow {task['workflow_execution_id']}." ) else: print( f"Failed to send cancel request for workflow {task['workflow_execution_id']}." ) continue if task["status"] == "SUBMITTING": _process_submitting_task( task, allow_immediate_create=False, workflow_name_was_preexisting=True, ) continue if task["status"] == "REQUESTED": print(f"Processing REQUESTED task {task['id']}...") workflow_name_was_preexisting = bool(task.get("workflow_execution_name")) workload_template = canonicalize_workload_template_or_default( task.get("workload_template") ) workflow_name = ( task.get("workflow_execution_name") or f"SPU_auto_{workload_template.replace('-','_')}_{task['id']}_{int(time.time())}" ) claimed = db.update_task_status_if_current( task["id"], "SUBMITTING", expected_current_statuses=["REQUESTED"], workflow_name=workflow_name, ) if not claimed: continue task = db.get_task_by_id(task["id"]) or task task["status"] = "SUBMITTING" task["workflow_execution_name"] = workflow_name _process_submitting_task( task, allow_immediate_create=True, workflow_name_was_preexisting=workflow_name_was_preexisting, ) continue elif task["status"] == "PENDING": print( f"Checking PENDING task {task['id']} (Workflow: {task['workflow_execution_id']})..." ) if not task["workflow_execution_id"]: print( f"Task {task['id']} has no workflow ID but is PENDING. Marking FAILURE." ) db.update_task_status_if_current( task["id"], "FAILURE", expected_current_statuses=["PENDING"], ) continue execution_status = tcms_client.get_execution_status( task["workflow_execution_id"] ) if execution_status: state = execution_status.get( "status" ) # e.g. SUCCESS, FAILURE, RUNNING print(f"Workflow State: {state}") if state == "RUNNING": db.update_task_status_if_current( task["id"], state, expected_current_statuses=["PENDING"], ) elif state == "SUCCESS": # Save initial result, will be parsed later outputs = execution_status raw_execution_response = outputs db.save_result( task["id"], "SUCCESS", result_data=None, raw_execution_response=raw_execution_response, parsed=False, ) db.update_task_status_if_current( task["id"], "SUCCESS", expected_current_statuses=["PENDING"], ) print(f"Task {task['id']} SUCCESS.") # Save cluster metrics to database _save_cluster_metrics_for_task( task["id"], raw_execution_response ) elif state == "FAILURE": error_msg = execution_status.get("output", "Unknown error") raw_execution_response = execution_status db.save_result( task["id"], "FAILURE", error_msg=error_msg, raw_execution_response=raw_execution_response, parsed=False, ) db.update_task_status_if_current( task["id"], "FAILURE", expected_current_statuses=["PENDING"], ) print(f"Task {task['id']} FAILURE.") # Save cluster metrics to database (even on failure) _save_cluster_metrics_for_task( task["id"], raw_execution_response ) elif state == "ERROR": error_msg = execution_status.get("output", "Unknown error") db.save_result( task["id"], "ERROR", error_msg=error_msg, raw_execution_response=execution_status, parsed=False, ) db.update_task_status_if_current( task["id"], "ERROR", expected_current_statuses=["PENDING"], ) print(f"Task {task['id']} ERROR.") elif state == "CANCELLED": db.update_task_status_if_current( task["id"], "CANCELLED", expected_current_statuses=["PENDING"], ) print(f"Task {task['id']} CANCELLED.") elif task["status"] == "RUNNING": print( f"Checking RUNNING task {task['id']} (Workflow: {task['workflow_execution_id']})..." ) if not task["workflow_execution_id"]: print( f"Task {task['id']} has no workflow ID but is RUNNING. Marking FAILURE." ) db.update_task_status_if_current( task["id"], "FAILURE", expected_current_statuses=["RUNNING"], ) continue execution_status = tcms_client.get_execution_status( task["workflow_execution_id"] ) if execution_status: state = execution_status.get( "status" ) # e.g. SUCCESS, FAILURE, RUNNING print(f"Workflow State: {state}") if state == "SUCCESS": # Save initial result, will be parsed later outputs = execution_status raw_execution_response = outputs db.save_result( task["id"], "SUCCESS", result_data=None, raw_execution_response=raw_execution_response, parsed=False, ) db.update_task_status_if_current( task["id"], "SUCCESS", expected_current_statuses=["RUNNING"], ) print(f"Task {task['id']} SUCCESS.") # Save cluster metrics to database _save_cluster_metrics_for_task( task["id"], raw_execution_response ) elif state == "FAILURE": error_msg = execution_status.get("output", "Unknown error") raw_execution_response = execution_status db.save_result( task["id"], "FAILURE", error_msg=error_msg, raw_execution_response=raw_execution_response, parsed=False, ) db.update_task_status_if_current( task["id"], "FAILURE", expected_current_statuses=["RUNNING"], ) print(f"Task {task['id']} FAILURE.") # Save cluster metrics to database (even on failure) _save_cluster_metrics_for_task( task["id"], raw_execution_response ) elif state == "ERROR": error_msg = execution_status.get("output", "Unknown error") db.save_result( task["id"], "ERROR", error_msg=error_msg, raw_execution_response=execution_status, parsed=False, ) db.update_task_status_if_current( task["id"], "ERROR", expected_current_statuses=["RUNNING"], ) print(f"Task {task['id']} ERROR.") elif state == "CANCELLED": db.update_task_status_if_current( task["id"], "CANCELLED", expected_current_statuses=["RUNNING"], ) print(f"Task {task['id']} CANCELLED.") else: print(f"Could not fetch status for task {task['id']}") elif ( task["status"] == "CANCELLED" or task["status"] == "SUCCESS" or task["status"] == "FAILURE" or task["status"] == "ERROR" ): print(f"Task {task['id']} {task['status']}.") except Exception as e: print(f"Error processing task {task['id']}: {e}") traceback.print_exc() def process_ended_tasks(): """Process ended tasks and parse their results.""" tasks = db.get_ended_tasks() if not tasks: return print(f"Found {len(tasks)} ended tasks to parse.") for task in tasks: try: print( f"Parsing results for task {task['id']} " f"(status: {task['status']})..." ) # Get raw execution response from the task # Note: We need to fetch it from the result table # For now, we'll fetch it via get_execution_status if not task["workflow_execution_id"]: print(f"Task {task['id']} has no workflow ID, skipping.") continue execution_status = tcms_client.get_execution_status( task["workflow_execution_id"] ) if not execution_status: print(f"Could not fetch execution status for " f"task {task['id']}") continue # Parse the result based on status and workload template workload_template = canonicalize_workload_template_or_default( task.get("workload_template") ) result = parse_execution_result( execution_status, task["status"], workload_template, ) # Update the result in database with parsed=True # Save the entire result as dict (contains nodes, failed_node_name, etc.) db.save_result( task["id"], task["status"], result_data=result.to_dict(), error_msg=result.error, raw_execution_response=execution_status, parsed=True, ) print(f"Task {task['id']} results parsed successfully.") # Save cluster metrics (fallback mechanism with cache check) _save_cluster_metrics_for_task(task["id"], execution_status) except Exception as e: print(f"Error parsing results for task {task['id']}: {e}") traceback.print_exc() def run_scheduler(): """Run scheduler once - process all pending/running tasks and exit. Suitable for cron jobs or GitHub Actions workflows.""" print("Starting Scheduler (single run mode)...") # db.init_db() # Tables should be managed by Alembic process_pending_tasks() process_ended_tasks() # TODO: Add notification logic (notify users when tasks complete, # only if not already notified) print("Scheduler completed.") # TODO: Add cancel_task function to mark tasks as deleted and handle cleanup if __name__ == "__main__": # POLL_INTERVAL = 6000 # seconds # while True: # run_scheduler() # time.sleep(POLL_INTERVAL) run_scheduler()