vae_parallel_train.py

import os
from pathlib import Path
import concurrent.futures

def run_command(month, day, year, antenna):
    """Function to execute the os.system command."""
    command = f"python vae_pipeline.py --month {month} --day {day} --year {year} --antenna {antenna}"
    print(f"Running: {command}")
    os.system(command)

def process_folders(base_folder, max_workers=4):
    """
    Processes each A2 file in the folders under base_folder in parallel.
    
    Args:
        base_folder (str): Path to the base folder containing subfolders with A2 files.
        max_workers (int): Maximum number of parallel commands.
    """
    # Collect all A2 files and their respective month, day, year information
    tasks = []
    for folder in Path(base_folder).iterdir():
        if folder.is_dir():
            folder_name = folder.name
            try:
                # Parse month, day, and year from the folder name
                month, day, year = folder_name.split('-')
                day = day.zfill(2)  # Ensure day is two digits
            except ValueError:
                print(f"Skipping invalid folder name: {folder_name}")
                continue
            
            # Check if A2 file exists
            file_path = folder / "A2.npy"
            if file_path.exists():
                tasks.append((month, day, year, "A2"))
            else:
                print(f"Skipping folder {folder_name}: No A2 file found.")

    # Run commands in parallel
    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = [
            executor.submit(run_command, month, day, year, file_path)
            for month, day, year, file_path in tasks
        ]
        for future in concurrent.futures.as_completed(futures):
            try:
                future.result()  # Wait for each future to complete
            except Exception as e:
                print(f"Error occurred: {e}")

if __name__ == "__main__":
    base_folder = "data/" 
    max_workers = 6  # Define the maximum number of parallel commands
    process_folders(base_folder, max_workers=max_workers)