Skip to main content

MAIA Workspace Setup

This module provides utilities for setting up MAIA workspaces with proper file structure for MPMD coupling. It handles downloading from HF Hub and creating necessary symbolic links with correct file extensions.

MaiaWorkspace Objects​

class MaiaWorkspace()

Manages MAIA workspace setup including file downloads and symbolic links.

File preparation is driven by solver profiles in :data:~hydrogym.data_manager.SOLVER_PROFILES. The correct profile ('MAIA_LB' or 'MAIA_STRCTRD') is auto-detected from sentinel files.

__init__​

def __init__(environment_name: str,
             work_dir: Optional[str] = None,
             hf_repo_id: str = "dynamicslab/HydroGym-environments",
             local_fallback_dir: Optional[str] = None,
             use_clean_cache: bool = True,
             solver_type: Optional[str] = None)

Initialize the MAIA workspace.

Arguments:

  • environment_name - Name of the environment (e.g., 'Cylinder_2D_Re200').
  • work_dir - Working directory path. If None, uses './run_{environment_name}'.
  • hf_repo_id - Hugging Face repository ID.
  • local_fallback_dir - Optional local fallback directory.
  • use_clean_cache - Whether to use clean cache for HF downloads.
  • solver_type - Solver profile key ('MAIA_LB' or 'MAIA_STRCTRD'). Auto-detected from sentinel files if None (recommended). Defaults to 'MAIA_LB' as fallback for legacy environments.

setup​

def setup(force_download: bool = False,
          verbose: bool = True) -> Dict[str, str]

Set up the workspace with all required files.

This method:

  1. Downloads/locates environment data from HF Hub
  2. Creates work directory structure
  3. Creates symbolic links with proper names/extensions

Arguments:

  • force_download - Force re-download from HF Hub.
  • verbose - Print setup progress.

Returns:

Dictionary containing paths to key files:

  • 'work_dir': Path to work directory
  • 'properties_file': Path to properties file for MAIA
  • 'config_file': Path to environment config file
  • 'env_data_path': Path to source environment data

cleanup​

def cleanup(remove_work_dir: bool = False, verbose: bool = True) -> None

Clean up the workspace.

Arguments:

  • remove_work_dir - If True, removes the entire work directory. If False, only removes symbolic links.
  • verbose - Print cleanup messages.

prepare_maia_workspace​

def prepare_maia_workspace(
        environment_name: str,
        work_dir: Optional[str] = None,
        hf_repo_id: str = "dynamicslab/HydroGym-environments",
        force_download: bool = False,
        **kwargs) -> Tuple[str, str]

Convenience function to prepare a MAIA workspace for MPMD coupling.

This function handles all the file preparation needed before launching mpirun with Python and MAIA. It downloads data from HF Hub (if needed) and creates symbolic links with proper file extensions.

Arguments:

  • environment_name - Name of the environment (e.g., 'Cylinder_2D_Re200').
  • work_dir - Working directory path. If None, auto-generates name.
  • hf_repo_id - Hugging Face repository ID.
  • force_download - Force re-download from HF Hub.
  • **kwargs - Additional arguments passed to MaiaWorkspace.

Returns:

Tuple of (work_dir_path, properties_file_path) for use with mpirun.

Example:

>>> import hydrogym.maia as maia >>> >>> # Prepare workspace (before submitting HPC job) >>> work_dir, props_file = maia.prepare_maia_workspace('Cylinder_2D_Re200') >>> >>> # Then use work_dir and props_file in your job script: >>> # sbatch job.slurm # where job.slurm references work_dir