Building a New Server¶

The creation of a Melissa server is largely based on inheritance, which establishes the structure for study execution. Depending on the type of study, users must inherit a specific server class defined in Melissa and extend it with their own functionalities.

Melissa server class hierarchy¶

Based on the hierarchy shown above, users are expected to inherit one of the child classes.

To begin, create a new Python script that defines the custom server class. Both the script and the class name are specified in the server_filename and server_class options in the configuration file.

Deep-Learning Server¶

Melissa's DeepMelissaServer class key aspects:

Following code snippet showcases an example provided in examples/heat-pde/heatpde_dl_server.py script. Users must refer to this and then build their own use-case specific servers.

import torch
from melissa.server.deep_learning.torch_server import TorchServer
from melissa.server.parameters import ParameterSamplerType

from local_module import MyModel

logger = logging.getLogger("melissa")


class HeatPDEServerDL(TorchServer):
    """Use-case specific server"""

    def __init__(self, config_dict: dict[str, Any]):
        super().__init__(config_dict)
        self.param_list = ["ic", "b1", "b2", "b3", "b4", "t"]
        study_options = self.config_dict["study_options"]

        # custom options
        self.mesh_size = study_options["mesh_size"]
        Tmin, Tmax = study_options["parameter_range"]

        # example of random uniform sampling
        self.set_parameter_sampler(
            sampler_t=ParameterSamplerType.RANDOM_UNIFORM, l_bounds=[Tmin], u_bounds=[Tmax]
        )

        # setting the attribute valid_dataloader allows validation
        # to be run alongside training
        self.valid_dataloader = self.get_validation_dataloader()  # user-defined

    @override
    def prepare_training_attributes(self):
        """Abstract method that must return model and optimizer."""

        model = self.wrap_model_ddp(
            self.MyModel(self.nb_parameters + 1, self.mesh_size * self.mesh_size, 1).to(self.device)
        )

        optimizer = torch.optim.Adam(
            model.parameters(), lr=self.dl_config.get("lr", 1e-3), weight_decay=1e-4
        )

        return model, optimizer

    @override
    def training_step(self, batch, batch_idx, **kwargs):

        # backprogation
        self.optimizer.zero_grad()
        x, y_target = batch
        x = x.to(self.device)
        y_target = y_target.to(self.device)
        y_pred = self.model(x)
        loss = self.criterion(y_pred, y_target)
        loss.backward()
        self.optimizer.step()
        self.learning_rate_scheduler.step()
        self.metric_logger.log_scalar("Loss/train", loss.item(), batch_idx)
        logger.info(f"[TRAINING] batch-id={batch_id} loss={loss.item():.2e}")

    @override
    def process_simulation_data(self, msg: SimulationData, config_dict: dict):
        """Abstract method for transformation while batch creation."""

        field = "temperature"
        # cast msg.data to float32
        x = torch.from_numpy(
            np.array(msg.parameters[-self.nb_parameters :] + [msg.time_step], dtype=np.float32)
        )
        y = torch.from_numpy(msg[field].astype(np.float32))

        return x, y

@override
def process_simulation_data(self, msg: SimulationData, config_dict: dict):
    """Abstract method for transformation while batch creation."""

    field = "temperature"

    # cast msg.data to float32
    x = torch.from_numpy(
        np.array(msg.parameters[-self.nb_parameters :] + [msg.time_step], dtype=np.float32)
    )
    temp_node: conduit.Node = msg[field]
    data: np.ndarray = temp_node["mesh/data"]
    current_sim_time: float = temp_node["metadata/sim_time"]
    dt: float = temp_node["metadata/dt"]

    y = torch.from_numpy(data.astype(np.float32))

    return x, y

melissa::send("temperature", temp_node);
melissa::send("pressure", pressure_node);
melissa::send("velocity", velocity_vector);

@override
def process_simulation_data(self, msg: SimulationData, config_dict: dict):
    """Abstract method for transformation while batch creation."""
    assert isinstance(msg, SimulationData)
    temp_node: conduit.Node = msg["temperature"]
    pressure_node: conduit.Node = msg["pressure"]
    velocity_arr: np.array = msg["velocity"]

Modifications for training¶

• Users can access JSON configuration options through config_dict dictionary. Although, some are already exposed as attributes in Melissa server super classes such as nb_simulations, nb_time_steps, etc. (See exposed properties section below)

• Call self.set_parameter_sampler, which accepts either the pre-defined Enum values (ParameterSamplerType) or a custom sampler class type that inherits base classes.

• Assuming your model architecture is already implemented, override prepare_training_attributes to return the (model, optimizer) tuple.

• Override process_simulation_data, a transformation method applied to data retrieved from the buffer when creating a batch. This method takes an instance of SimulationData and the config_dict containing all configuration settings from the JSON file.

• Override training_step, which accepts the transformed data (batch) and the current batch index (batch_idx) as inputs.

Modifications for Validation (Optional)¶

To enable validation, users need to follow these steps:

• The DeepMelissaServer training loop expects self.valid_dataloader to be set, but its definition is left to the user. The provided code snippet initializes self.valid_dataloader through a custom method.

• Override the validation_step method, which takes 3 inputs:

• the validation data (batch) from self.valid_dataloader
• the validation batch index (valid_batch_idx)
• the training batch index (batch_idx)

More Control (Optional)¶

For users who need greater flexibility over the training loop, DeepMelissaServer provides several hook methods (pytorch-lightning like) that are triggered at specific points during training:

Exposed properties¶

Following are some of the important server attributes users may requires: