Refactor XpkTask using OOP subclasses

[ryanliuCienet]

Description

This PR refactors the XpkTask class structure to follow OOP principles, introducing clean inheritance subclasses for specific test sub-workflows (node interruption and name generation/quarantine).

This allows XpkTask to maintain a clean interface and establishes a uniform polymorphic run() entry point across all tasks.

Tests

DAG that uses XpkNameGenAndQuarantineTask class:
maxtext_convergence

DAG that uses XpkNodeInterruptionTask class:
maxtext_emc_orbax_res_gcs

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run one-shot tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[alfredyu-cienet]

1

have original class returns a derivative class does not make sense

the proper way should be like (dags/multipod/configs/gke_config.py)

def get_gke_config(...) -> task.XpkTask:
def get_gke_config_with_interrupt(...) -> task.XpkNodeInterruptionTask:
def get_gke_config_with_name_gen_and_quarantine(...) -> task.XpkNameGenAndQuarantineTask:

---

2

for XpkTask XpkNodeInterruptionTask XpkNameGenAndQuarantineTask, as mentioned before, want only one single entry point, which is run

but they all have multiple entry points

---

3

for XpkTask XpkNodeInterruptionTask XpkNameGenAndQuarantineTask
a key intention of OOP is to reuse code

sample figure for reference

graph LR
    subgraph Class_XpkNode ["Class: XpkNodeInterruptionTask"]
        subgraph TG_Benchmark_B ["TaskGroup: benchmark_id"]
            
            subgraph TG_RunModel_B ["TaskGroup: run_model"]
                subgraph TG_LaunchReach_B ["TaskGroup: launch_workload_with_node_reach_to_step"]
                    T_RunWorkload_B["Task: xpk.run_workload"] --> T_WaitStart_B["Task: xpk.wait_for_workload_start"]
                    T_WaitStart_B --> T_WaitStep_B["Task: xpk.wait_for_workload_reach_step"]
                    T_WaitStep_B --> B_Decider_B{"Branch: task_path_decider"}
                    
                    B_Decider_B -->|"True"| T_WaitFile_B["Task: gcs.wait_for_file_to_exist"]
                    B_Decider_B -->|"False"| T_DoNothing_B["Task: do_nothing"]
                end
                
                TG_LaunchReach_B --> T_DeleteNode_B["Task: xpk.delete_node"]
                T_DeleteNode_B --> T_WaitComplete_B["Task: xpk.wait_for_workload_completion"]
                T_WaitComplete_B --> T_Cleanup_B["Task: xpk.clean_up_workload"]
            end

            subgraph TG_PostProcess_B ["TaskGroup: post_process"]
                T_GenPID_B["Task: metric.generate_process_id"] --> T_Xplane_B["Task: metric.xplane_to_metrics"]
                T_Xplane_B --> T_ProcMetrics_B["Task: metric.process_metrics"]
            end

            TG_RunModel_B --> TG_PostProcess_B
        end
    end

    subgraph Class_XpkTask ["Class: XpkTask"]
        subgraph TG_Benchmark_A ["TaskGroup: benchmark_id"]
            
            subgraph TG_RunModel_A ["TaskGroup: run_model"]
                subgraph TG_Launch_A ["TaskGroup: launch_workload"]
                    T_RunWorkload_A["Task: xpk.run_workload"] --> T_WaitStart_A["Task: xpk.wait_for_workload_start"]
                end
                
                TG_Launch_A --> T_WaitComplete_A["Task: xpk.wait_for_workload_completion"]
                T_WaitComplete_A --> T_Cleanup_A["Task: xpk.clean_up_workload"]
            end

            subgraph TG_PostProcess_A ["TaskGroup: post_process"]
                T_GenPID_A["Task: metric.generate_process_id"] --> T_Xplane_A["Task: metric.xplane_to_metrics"]
                T_Xplane_A --> T_ProcMetrics_A["Task: metric.process_metrics"]
            end

            TG_RunModel_A --> TG_PostProcess_A
        end
    end

[ryanliuCienet]

I've tried to update it based on the comment, please have a look. Thank you.