pydantic base for parsl configurations.

[braceal]

Allow full control of parsl configuration from input YAML file. For the workstation case:

compute_settings:
  name: workstation
  available_accelerators: 4

Cons: The BaseComputeSettings currently requires an abstract function to return a Parsl Config object which couples the runtime to Parsl. This could be generalized down the road with minimal changes.

What do you think @WardLT ?

[WardLT]

I'm iffy about it. It is nice to allow people to say "create compute setting X" with local options "Y" but don't want you to have to write a lot of boilerplate wrappers over Parsl classes.

Perhaps we have just a few of those for really common systems (e.g., local workstations, Polaris) but make sure to have the ability to add a new one w/o modifying the deepdrive python module

[WardLT]

I was thinking about this more, and am curious if this is also your goal?

A function like

def make_config(name: str, run_dir: Path, **kwargs) -> Config:
    # Lookup the right function (
     build_fn = {
         'local': build_local,
         'polaris': build_polaris
     }

     return build_fn(run_dir, **kwargs)

[braceal]

Yeah, that is the functionality. Pydantic just lets us specify the config kwargs from yaml so all the experiment settings can be contained together. I pushed a small change that isolates all the Parsl config stuff to parsl.py (the ComputeSettings are no longer required). So if users want to use the pre-packaged configs (with minimal settings updates, e.g., specify number of GPUs in a workstation) they need to do two things:

from deepdrivemd.api import DeepDriveMDSettings
from deepdrivemd.parsl import ComputeSettingsTypes

class ExperimentSettings(DeepDriveMDSettings): 
    # This class is where users put specific application settings
    simulation_settings: MDSimulationSettings
    train_settings: CVAETrainSettings
    inference_settings: CVAEInferenceSettings

    # Add in optional compute settings
    compute_settings: ComputeSettingsTypes

# Define the parsl configuration (this can be done using the config_factory
# for common use cases or by defining your own configuration.)
parsl_config = cfg.compute_settings.config_factory(cfg.run_dir / "run-info")

doer = ParslTaskServer(
        [run_simulation, run_train, run_inference], queues, parsl_config
)

This is all user-level, so users can choose to use their own parsl configs and remove the compute_settings: ComputeSettingsTypes line in the data class.

As we make examples for Polaris, etc, we can define a generic settings mixin to reduce code duplication. Something like this in parsl.py:

class BaseComputeSettings(BaseSettings, ABC):
    """Compute settings (HPC platform, number of GPUs, etc)."""

    name: Literal[""] = ""
    """Name of the platform to use."""

    @abstractmethod
    def config_factory(self, run_dir: PathLike) -> Config:
        """Create a new Parsl configuration.

        Parameters
        ----------
        run_dir : PathLike
            Path to store monitoring DB and parsl logs.

        Returns
        -------
        Config
            Parsl configuration.
        """
        ...

# All the settings in these classes are directly configurable from yaml, e.g.,
# compute_settings:
#   name: polaris
#   num_nodes: 2
class HPCSettings(BaseSettings):
    num_nodes: int = 1
    worker_init: str
    scheduler_options: str = ""
    account: str
    queue: str
    walltime: str

class PolarisSettings(BaseComputeSettings, HPCSettings):
    name: Literal["polaris"] = "polaris"
    # specific polaris options ...

class PerlmutterSettings(BaseComputeSettings, HPCSettings):
    name: Literal["perlmutter"] = "perlmutter"
    # specific perlmutter options ...

ComputeSettingsTypes = Union[LocalSettings, WorkstationSettings, PolarisSettings, PerlmutterSettings]

How's this?

[WardLT]

I like the sound of that. Let's go for this route and then cross the bridge of "do we expect users to add a new class to the python module?" when we actually have that problem I invented

[WardLT]

Neat use of Pydantics auto-matching type based on name 👍🏼