| #!/usr/bin/env python3 |
| |
| import os |
| import sys |
| import time |
| import logging |
| from concurrent.futures import ThreadPoolExecutor |
| from typing import List, Dict, Any, Union |
| import pydantic |
| from rich.progress import Progress, BarColumn, TextColumn, SpinnerColumn, TimeElapsedColumn |
| from rich.console import Console |
| from rich.live import Live |
| from rich.panel import Panel |
| from rich.layout import Layout |
| from collections import deque |
| |
| os.environ['AUTH_SERVER_URL'] = "https://auth.dev.cybershuttle.org" |
| os.environ['API_SERVER_HOSTNAME'] = "api.dev.cybershuttle.org" |
| os.environ['GATEWAY_URL'] = "https://gateway.dev.cybershuttle.org" |
| os.environ['STORAGE_RESOURCE_HOST'] = "gateway.dev.cybershuttle.org" |
| |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| |
| from create_launch_experiment_with_storage import create_and_launch_experiment |
| from airavata_experiments.airavata import AiravataOperator |
| from airavata.model.status.ttypes import ExperimentState |
| |
| |
| class ExperimentLaunchResult(pydantic.BaseModel): |
| """Result from creating and launching an experiment.""" |
| experiment_id: str |
| process_id: str |
| experiment_dir: str |
| storage_host: str |
| mount_point: str |
| |
| |
| class JobConfig(pydantic.BaseModel): |
| """Configuration for a batch job submission.""" |
| experiment_name: str |
| project_name: str |
| application_name: str |
| computation_resource_name: str |
| queue_name: str |
| node_count: int |
| cpu_count: int |
| walltime: int |
| group_name: str = "Default" |
| input_storage_host: str | None = None |
| output_storage_host: str | None = None |
| input_files: Dict[str, Union[str, List[str]]] | None = None |
| data_inputs: Dict[str, Union[str, int, float]] | None = None |
| gateway_id: str | None = None |
| auto_schedule: bool = False |
| |
| |
| class JobResult(pydantic.BaseModel): |
| """Result from submitting and monitoring a single job.""" |
| job_index: int |
| experiment_id: str | None |
| status: str |
| result: ExperimentLaunchResult | None = None |
| success: bool |
| error: str | None = None |
| |
| |
| def get_experiment_state_value(status) -> tuple[int, str, ExperimentState]: |
| """Extract state value, name, and enum from status. Returns (value, name, enum).""" |
| if isinstance(status, ExperimentState): |
| return status.value, status.name, status |
| |
| # Handle ExperimentStatus object |
| if hasattr(status, 'state'): |
| state = status.state |
| if isinstance(state, ExperimentState): |
| return state.value, state.name, state |
| elif hasattr(state, 'value'): |
| return state.value, state.name if hasattr(state, 'name') else str(state), state |
| |
| # Handle direct value/name access |
| status_value = status.value if hasattr(status, 'value') else (status if isinstance(status, int) else None) |
| status_name = status.name if hasattr(status, 'name') else str(status) |
| |
| # Convert to ExperimentState enum |
| if status_value is not None: |
| try: |
| enum_state = ExperimentState(status_value) |
| return status_value, status_name, enum_state |
| except (ValueError, TypeError): |
| pass |
| |
| # Fallback |
| return None, status_name, ExperimentState.FAILED |
| |
| |
| def monitor_experiment_silent(operator: AiravataOperator, experiment_id: str, check_interval: int = 30) -> ExperimentState: |
| """Monitor experiment silently until completion. Returns final status.""" |
| logger = logging.getLogger(__name__) |
| max_checks = 3600 # Maximum number of checks (about 5 hours at 5s interval) |
| check_count = 0 |
| |
| # Use shorter interval initially, then increase |
| initial_interval = min(check_interval, 5) # Check every 5 seconds initially |
| |
| while check_count < max_checks: |
| try: |
| status = operator.get_experiment_status(experiment_id) |
| |
| # Extract state information |
| status_value, status_name, status_enum = get_experiment_state_value(status) |
| |
| # Log status periodically for debugging |
| if check_count % 12 == 0: # Log every minute (12 * 5s) |
| logger.debug(f"Experiment {experiment_id} status check {check_count}: value={status_value}, name={status_name}") |
| |
| # Check terminal states: COMPLETED (7), CANCELED (6), FAILED (8) |
| if status_value is not None: |
| is_terminal = status_value in [ |
| ExperimentState.COMPLETED.value, # 7 |
| ExperimentState.CANCELED.value, # 6 |
| ExperimentState.FAILED.value # 8 |
| ] |
| else: |
| is_terminal = status_name in ['COMPLETED', 'CANCELED', 'FAILED'] |
| |
| if is_terminal: |
| logger.info(f"Experiment {experiment_id} reached terminal state: {status_name} (value: {status_value})") |
| return status_enum |
| |
| except Exception as e: |
| # If we can't get status, log but continue monitoring |
| logger.warning(f"Error checking experiment {experiment_id} status (check {check_count}): {e}") |
| import traceback |
| logger.debug(traceback.format_exc()) |
| if check_count > 10: # After several failed checks, assume failed |
| logger.error(f"Multiple status check failures for {experiment_id}, assuming FAILED") |
| return ExperimentState.FAILED |
| |
| # Sleep before next check |
| sleep_time = initial_interval if check_count < 6 else check_interval |
| time.sleep(sleep_time) |
| check_count += 1 |
| |
| # If we've exceeded max checks, assume failed |
| logger.error(f"Experiment {experiment_id} monitoring timeout after {check_count} checks, assuming FAILED") |
| return ExperimentState.FAILED |
| |
| |
| def submit_and_monitor_job( |
| job_index: int, |
| job_config: JobConfig | Dict[str, Any], |
| access_token: str, |
| ) -> JobResult: |
| """Submit and monitor a single job. Returns job result with status.""" |
| # Convert dict to JobConfig if needed |
| if isinstance(job_config, dict): |
| job_config = JobConfig(**job_config) |
| |
| try: |
| # Make experiment name unique for each job to avoid directory conflicts |
| # Using job_index ensures uniqueness and makes it easy to track |
| unique_experiment_name = f"{job_config.experiment_name}-job{job_index}" |
| |
| # Handle input_files and data_inputs same way as working version |
| input_files = job_config.input_files if job_config.input_files else None |
| data_inputs = job_config.data_inputs if job_config.data_inputs else None |
| |
| result_dict = create_and_launch_experiment( |
| access_token=access_token, |
| experiment_name=unique_experiment_name, |
| project_name=job_config.project_name, |
| application_name=job_config.application_name, |
| computation_resource_name=job_config.computation_resource_name, |
| queue_name=job_config.queue_name, |
| node_count=job_config.node_count, |
| cpu_count=job_config.cpu_count, |
| walltime=job_config.walltime, |
| group_name=job_config.group_name, |
| input_storage_host=job_config.input_storage_host, |
| output_storage_host=job_config.output_storage_host, |
| input_files=input_files, |
| data_inputs=data_inputs, |
| gateway_id=job_config.gateway_id, |
| auto_schedule=job_config.auto_schedule, |
| monitor=False, |
| ) |
| |
| operator = AiravataOperator(access_token=access_token) |
| experiment_id = result_dict['experiment_id'] |
| |
| # Check status immediately after submission to catch early failures |
| try: |
| initial_status = operator.get_experiment_status(experiment_id) |
| status_value, status_name, status_enum = get_experiment_state_value(initial_status) |
| |
| # Check if already in terminal state |
| if status_value is not None and status_value in [ |
| ExperimentState.COMPLETED.value, |
| ExperimentState.CANCELED.value, |
| ExperimentState.FAILED.value |
| ]: |
| # Already in terminal state |
| final_status = status_enum |
| else: |
| # Monitor until completion |
| final_status = monitor_experiment_silent(operator, experiment_id) |
| except Exception as e: |
| # If we can't check status, log and assume failed |
| logger = logging.getLogger(__name__) |
| logger.error(f"Error monitoring experiment {experiment_id}: {e}") |
| import traceback |
| logger.debug(traceback.format_exc()) |
| final_status = ExperimentState.FAILED |
| |
| result = ExperimentLaunchResult(**result_dict) |
| |
| return JobResult( |
| job_index=job_index, |
| experiment_id=result.experiment_id, |
| status=final_status.name, |
| result=result, |
| success=final_status == ExperimentState.COMPLETED, |
| ) |
| except Exception as e: |
| # Log the error for debugging |
| import traceback |
| error_msg = f"{str(e)}\n{traceback.format_exc()}" |
| logger = logging.getLogger(__name__) |
| logger.error(f"Job {job_index} failed: {error_msg}") |
| |
| return JobResult( |
| job_index=job_index, |
| experiment_id=None, |
| status='ERROR', |
| result=None, |
| success=False, |
| error=str(e), |
| ) |
| |
| |
| def batch_submit_jobs( |
| job_config: JobConfig | Dict[str, Any], |
| num_copies: int = 10, |
| max_concurrent: int = 5, |
| access_token: str | None = None, |
| ) -> List[JobResult]: |
| """Submit multiple job copies in batches with progress bar.""" |
| if access_token is None: |
| from airavata_auth.device_auth import AuthContext |
| access_token = AuthContext.get_access_token() |
| |
| console = Console() |
| results = [] |
| log_buffer = deque(maxlen=50) # Keep last 50 log lines for display |
| |
| # Custom handler to capture logs to buffer |
| class ListHandler(logging.Handler): |
| def __init__(self, buffer): |
| super().__init__() |
| self.buffer = buffer |
| |
| def emit(self, record): |
| msg = self.format(record) |
| self.buffer.append(msg) |
| |
| log_handler = ListHandler(log_buffer) |
| log_handler.setLevel(logging.INFO) |
| log_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')) |
| |
| # Add to root logger and module logger |
| logging.root.addHandler(log_handler) |
| logger = logging.getLogger('create_launch_experiment_with_storage') |
| logger.addHandler(log_handler) |
| |
| # Configure progress bar |
| progress = Progress( |
| SpinnerColumn(), |
| TextColumn("[progress.description]{task.description}"), |
| BarColumn(), |
| TextColumn("[progress.percentage]{task.percentage:>3.0f}%"), |
| TextColumn("•"), |
| TextColumn("{task.completed}/{task.total}"), |
| TimeElapsedColumn(), |
| console=console, |
| ) |
| |
| task = progress.add_task( |
| f"{num_copies} total, 0 running, 0 completed, 0 failed", |
| total=num_copies |
| ) |
| |
| # Create layout with logs above and progress below |
| layout = Layout() |
| layout.split_column( |
| Layout(name="logs", size=None), |
| Layout(progress, name="progress", size=3) |
| ) |
| |
| def make_display(): |
| # Get logs from buffer - always show the latest logs (they're added to end of deque) |
| log_lines = list(log_buffer) if log_buffer else ["No logs yet..."] |
| # Show last 20 lines to keep display manageable and scrolled to bottom |
| display_lines = log_lines[-20:] if len(log_lines) > 20 else log_lines |
| log_text = '\n'.join(display_lines) |
| log_panel = Panel( |
| log_text, |
| title="Logs (latest)", |
| border_style="blue", |
| height=None, |
| expand=False |
| ) |
| layout["logs"].update(log_panel) |
| return layout |
| |
| try: |
| # Use Live to keep layout fixed, progress at bottom |
| with Live(make_display(), console=console, refresh_per_second=4, screen=True) as live: |
| with ThreadPoolExecutor(max_workers=max_concurrent) as executor: |
| active_futures = {} |
| next_job_index = 0 |
| |
| # Submit initial batch |
| while next_job_index < min(max_concurrent, num_copies): |
| future = executor.submit(submit_and_monitor_job, next_job_index, job_config, access_token) |
| active_futures[future] = next_job_index |
| next_job_index += 1 |
| |
| # Process completed jobs and submit new ones |
| # Continue until all jobs are submitted AND all active futures are done |
| while active_futures or next_job_index < num_copies: |
| completed_futures = [f for f in active_futures if f.done()] |
| |
| for future in completed_futures: |
| job_idx = active_futures.pop(future) |
| |
| try: |
| result = future.result() |
| results.append(result) |
| except Exception as e: |
| # Handle unexpected exceptions |
| results.append(JobResult( |
| job_index=job_idx, |
| experiment_id=None, |
| status='ERROR', |
| result=None, |
| success=False, |
| error=str(e), |
| )) |
| |
| # Submit next jobs if available and we have capacity |
| while next_job_index < num_copies and len(active_futures) < max_concurrent: |
| try: |
| new_future = executor.submit(submit_and_monitor_job, next_job_index, job_config, access_token) |
| active_futures[new_future] = next_job_index |
| next_job_index += 1 |
| except Exception as e: |
| # If submission itself fails, mark as error and continue |
| results.append(JobResult( |
| job_index=next_job_index, |
| experiment_id=None, |
| status='ERROR', |
| result=None, |
| success=False, |
| error=f"Submission failed: {str(e)}", |
| )) |
| next_job_index += 1 |
| |
| # Update progress bar with counts |
| completed_count = len(results) |
| running_count = len(active_futures) |
| submitted_count = next_job_index |
| successful_count = sum(1 for r in results if r.success) |
| failed_count = completed_count - successful_count |
| |
| # Show submitted count if not all jobs submitted yet |
| if submitted_count < num_copies: |
| status_desc = f"{num_copies} total, {submitted_count} submitted, {running_count} running, {completed_count} completed, {failed_count} failed" |
| else: |
| status_desc = f"{num_copies} total, {running_count} running, {completed_count} completed, {failed_count} failed" |
| |
| progress.update( |
| task, |
| completed=completed_count, |
| description=status_desc |
| ) |
| live.update(make_display()) |
| |
| if not completed_futures and next_job_index >= num_copies: |
| # Only sleep if nothing changed |
| time.sleep(1) |
| |
| # Sort results by job_index |
| results.sort(key=lambda x: x.job_index) |
| return results |
| finally: |
| # Clean up log handlers |
| logging.root.removeHandler(log_handler) |
| if log_handler in logger.handlers: |
| logger.removeHandler(log_handler) |
| |
| |
| def main(): |
| """Main function that sets up job configuration and runs batch submission.""" |
| from airavata_auth.device_auth import AuthContext |
| |
| access_token = AuthContext.get_access_token() |
| |
| # Job configuration - matching create_launch_experiment_with_storage.py exactly |
| EXPERIMENT_NAME = "Test" |
| PROJECT_NAME = "Default Project" |
| APPLICATION_NAME = "NAMD-test" |
| GATEWAY_ID = None |
| COMPUTATION_RESOURCE_NAME = "NeuroData25VC2" |
| QUEUE_NAME = "cloud" |
| NODE_COUNT = 1 |
| CPU_COUNT = 1 |
| WALLTIME = 5 |
| GROUP_NAME = "Default" |
| INPUT_STORAGE_HOST = "gateway.dev.cybershuttle.org" |
| OUTPUT_STORAGE_HOST = "149.165.169.12" |
| INPUT_FILES = {} |
| DATA_INPUTS = {} |
| AUTO_SCHEDULE = False |
| |
| job_config = JobConfig( |
| experiment_name=EXPERIMENT_NAME, |
| project_name=PROJECT_NAME, |
| application_name=APPLICATION_NAME, |
| computation_resource_name=COMPUTATION_RESOURCE_NAME, |
| queue_name=QUEUE_NAME, |
| node_count=NODE_COUNT, |
| cpu_count=CPU_COUNT, |
| walltime=WALLTIME, |
| group_name=GROUP_NAME, |
| input_storage_host=INPUT_STORAGE_HOST, |
| output_storage_host=OUTPUT_STORAGE_HOST, |
| input_files=INPUT_FILES if INPUT_FILES else None, |
| data_inputs=DATA_INPUTS if DATA_INPUTS else None, |
| gateway_id=GATEWAY_ID, |
| auto_schedule=AUTO_SCHEDULE, |
| ) |
| |
| num_copies = 10 |
| |
| try: |
| results = batch_submit_jobs( |
| job_config=job_config, |
| num_copies=num_copies, |
| max_concurrent=5, |
| access_token=access_token, |
| ) |
| |
| # Print summary |
| print("\n" + "="*60) |
| print(f"Batch submission complete: {num_copies} jobs") |
| print("="*60) |
| successful = sum(1 for r in results if r.success) |
| print(f"Successful: {successful}/{num_copies}") |
| print(f"Failed: {num_copies - successful}/{num_copies}") |
| print("\nJob Results:") |
| for result in results: |
| status_symbol = "✓" if result.success else "✗" |
| exp_id = result.experiment_id or 'N/A' |
| print(f" {status_symbol} Job {result.job_index}: {result.status} " |
| f"(ID: {exp_id})") |
| print("="*60) |
| |
| return results |
| |
| except Exception as e: |
| print(f"Failed to run batch submission: {repr(e)}", file=sys.stderr) |
| import traceback |
| traceback.print_exc() |
| sys.exit(1) |
| |
| |
| if __name__ == "__main__": |
| main() |
| |
