Untitled

import argparse
import ast
import sys
from datetime import datetime, date  # CHANGE: added `date` for daily retry tracking
import time
import re
import json
from azureml.core import Workspace, Experiment, Run
from azureml.pipeline.core import PipelineRun, StepRun
from azure.ai.ml import MLClient, dsl, load_component
from azure.ai.ml.entities import Job
from azure.identity import DefaultAzureCredential, InteractiveBrowserCredential, ManagedIdentityCredential

# CHANGE: Introduce global retry tracking variables
RETRY_TRACKER = {"count": 0, "day": date.today()}
MAX_RETRIES_PER_DAY = 10  # CHANGE: configurable retry limit


def create_aml_client(config_path: str = "config.json") -> MLClient:
    """
    initialize AML client based on credentials
    """
    with open(config_path) as json_file:
        data = json.load(json_file)
    print(data)
    credential = DefaultAzureCredential()
    ml_client = MLClient.from_config(credential=credential, path=config_path)
    return ml_client


def parse_args() -> argparse.Namespace:
    """
    command line parser
    """
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--experiment-name",
        type=str,
        default="One sentence summary billing",
        help="Name of the experiment",
    )
    parser.add_argument(
        "--schedule-name",
        type=str,
        help="Name of the schedule",
    )
    parser.add_argument(
        "--components-monitor-list",
        type=str,
        help="list of components to be monitored",
    )
    parser.add_argument(
        "--max-retries-per-day",  # CHANGE: allow overriding retry limit via CLI
        type=int,
        default=10,
        help="Maximum number of retries allowed per day",
    )
    return parser.parse_args()


# CHANGE: Helper function to check retry quota
def can_retry(max_retries: int) -> bool:
    global RETRY_TRACKER
    today = date.today()

    # Reset counter if a new day starts
    if RETRY_TRACKER["day"] != today:
        RETRY_TRACKER["count"] = 0
        RETRY_TRACKER["day"] = today

    if RETRY_TRACKER["count"] < max_retries:
        RETRY_TRACKER["count"] += 1
        print(f"[Retry Tracker] Retry attempt {RETRY_TRACKER['count']} of {max_retries} for {today}")
        return True
    else:
        print(f"[Retry Tracker] Daily retry limit ({max_retries}) reached for {today}. No more retries.")
        return False


def mointor_components_sdk_v2(experiment_name, components_monitor_list, schedule_name, max_retries):
    ml_client = create_aml_client()
    components_monitor_list = ast.literal_eval(components_monitor_list.replace("\\", ""))
    print(components_monitor_list)
    time.sleep(5)
    ws = Run.get_context().experiment.workspace

    # Get the experiment
    experiment = Experiment(workspace=ws, name=experiment_name)
    runs = experiment.get_runs()
    for index, run in enumerate(runs):
        if index == 0:
            print(f"Run ID: {run.id}, Status: {run.status} ,{run.display_name}")
            job_id = run.id
        break

    start_time = time.time()
    completed_count = 0
    while len(components_monitor_list) > completed_count:
        step_status_list = []
        print(job_id)
        job_status = ml_client.jobs.get(job_id).status
        print(f"jobs_name:{run.id}")

        components = ml_client.jobs.list(parent_job_name=run.id)

        for component in components:
            print(
                f"component Name & ID: {component.display_name},{component.name}, "
                f"Status: {component.status}: Created at :{component.creation_context.created_at}"
            )
            dt = datetime.fromisoformat(str(component.creation_context.created_at))
            epoch_time = dt.timestamp()
            tup = (component.display_name, component.status, epoch_time)
            step_status_list.append(tup)
            print(step_status_list)

        completed_count = 0
        for step_name, step_status, step_start_time in step_status_list:
            if step_status == "Completed":
                completed_count += 1
                print(f"completed_count:{completed_count}")
            elapsed_time = time.time() - step_start_time
            for component_name, threshold_duration in components_monitor_list:
                if component_name == step_name and step_status == "Running" and elapsed_time > threshold_duration:
                    print(f"{component_name} running for more than {threshold_duration} secs. Cancelling...")

                    # CHANGE: check retry quota before triggering
                    if can_retry(max_retries):
                        print(f"Triggering {schedule_name} schedule")
                        ml_client.schedules.trigger(schedule_name)
                    else:
                        print("Retry skipped due to daily cap.")

                    print("Cancelling Job....")
                    cancel = ml_client.jobs.begin_cancel(run.id)
                    print(cancel.status())
                    sys.exit()

                elif component_name == step_name and step_status == "Completed":
                    print(f"{component_name} completed successfully.")

                elif component_name == step_name and step_status == "Failed":
                    print(f"{component_name} failed, Retriggering the schedule in 15 mins...")

                    # CHANGE: check retry quota before triggering
                    if can_retry(max_retries):
                        time.sleep(900)
                        print(f"Triggering {schedule_name} schedule")
                        ml_client.schedules.trigger(schedule_name)
                    else:
                        print("Retry skipped due to daily cap.")

                    print("Cancelling Job....")
                    cancel = ml_client.jobs.begin_cancel(run.id)
                    print(cancel.status())
                    sys.exit()
        time.sleep(10)
    return True


if __name__ == "__main__":
    args = parse_args()

    # CHANGE: pass retry limit from CLI args
    mointor_components_sdk_v2(
        args.experiment_name,
        args.components_monitor_list,
        args.schedule_name,
        args.max_retries_per_day,
    )