Interactive Terminate

from kubernetes import client, config

def terminate_vcjob(vcjob: dict, namespace):
    vcjob_name = vcjob["metadata"]["name"]
    api_instance = client.CustomObjectsApi()

    try:
        vcjob["status"]["state"]["phase"] = "Aborted"
        api_instance.replace_namespaced_custom_object_status(
            group="batch.volcano.sh",
            version="v1alpha1",
            namespace=namespace,
            plural="jobs",
            name=vcjob_name,
            body=vcjob,
        )
        print(f"Successfully changed job {vcjob_name} status.")
    except Exception as e:
        print(f"Failed to change status for job {vcjob_name}: {e}")

def is_vcjob_running(vcjob: dict) -> bool:
    return vcjob["status"]["state"]["phase"] == "Running"

def get_vcjob_by_name(vcjob_name, namespace):
    api_instance = client.CustomObjectsApi()
    try:
        vcjob = api_instance.get_namespaced_custom_object(
            group="batch.volcano.sh",
            version="v1alpha1",
            namespace=namespace,
            plural="jobs",
            name=vcjob_name,
        )
        return vcjob
    except client.exceptions.ApiException as e:
        print(f"Failed to get job {vcjob_name}: {e}.")
        return None

def main():
    # Prompt and validate namespace
    valid_namespaces = {"application", "application-nonprod"}
    while True:
        namespace = input("Enter the namespace ('application' or 'application-nonprod'): ").strip()
        if namespace in valid_namespaces:
            break
        print("Invalid namespace. Please enter either 'application' or 'application-nonprod'.")

    print("Loading Kubernetes config...")
    config.load_kube_config()
    api_instance = client.CustomObjectsApi()

    # List all jobs and filter for running ones
    try:
        jobs = api_instance.list_namespaced_custom_object(
            group="batch.volcano.sh",
            version="v1alpha1",
            namespace=namespace,
            plural="jobs",
        )["items"]
    except Exception as e:
        print(f"Error listing jobs: {e}")
        return

    running_jobs = [job for job in jobs if job.get("status", {}).get("state", {}).get("phase") == "Running"]

    if not running_jobs:
        print(f"No running vcjobs found in namespace '{namespace}'.")
        return

    # Display and prompt user to select a job
    print("\nRunning Jobs:")
    for idx, job in enumerate(running_jobs):
        print(f"{idx + 1}: {job['metadata']['name']}")

    while True:
        try:
            selection = int(input("Select a job to terminate by number: "))
            if 1 <= selection <= len(running_jobs):
                selected_job = running_jobs[selection - 1]
                break
            else:
                print("Invalid selection. Enter a number from the list.")
        except ValueError:
            print("Please enter a valid number.")

    job_name = selected_job["metadata"]["name"]

    # Prompt for dryrun
    dryrun_input = input("Dry run mode? (yes/no): ").strip().lower()
    dryrun = dryrun_input in ['yes', 'y']

    if dryrun:
        print(f"[DryRun] Would terminate job {job_name} in namespace {namespace}")
        return

    # Confirm termination
    confirm = input(f"Are you sure you want to terminate job '{job_name}'? (yes/no): ").strip().lower()
    if confirm not in ['yes', 'y']:
        print("Termination aborted.")
        return

    # Additional warning for production jobs
    if namespace == "application":
        print("\n⚠️ WARNING: You are about to terminate a PRODUCTION job.")
        while True:
            username = input("Enter your username to confirm: ").strip()
            if username:
                print(f"Username '{username}' confirmed. Proceeding with termination.")
                break
            else:
                print("Username cannot be empty. Please enter your username to proceed.")

    print(f"Terminating job {job_name} in namespace {namespace}...")

    if is_vcjob_running(selected_job):
        terminate_vcjob(selected_job, namespace)
        print(f"Job {job_name} is terminated successfully. Please use this command to check the job status:\n kubectl get vcjob -n {namespace} -o wide | grep '{job_name}'\n")
    else:
        print(f"Job {job_name} is not in Running state. Skip termination.")

if __name__ == "__main__":
    main()