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.

Note

The majority of the melissa.server module is statically typed, which can be helpful for identifying attributes and methods from super classes. This is especially useful if users are using IDEs with support for Pylance or similar extensions.

Melissa server class hierarchy¶

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.

Important

Due to the deeper inheritance structure, understanding the available attributes and methods can be challenging. We recommend reviewing the server documentation to familiarize yourself with all the attributes exposed in the user server class.

Sensitivity Analysis Server¶

In most SA use cases, users typically don't need to make many modifications to the newly created class. However, the parameter sampling strategy might need adjustments depending on how the client or solver expects the input arguments.

import logging

from melissa.server.sensitivity_analysis import SensitivityAnalysisServer
from melissa.server.parameters import ParameterSamplerType
from typing import Dict, Any

logger = logging.getLogger("melissa")


class HeatPDEServerSA(SensitivityAnalysisServer):
    """
    Use-case specific server
    """

    def __init__(self, config_dict: Dict[str, Any]):
        super().__init__(config_dict)
        Tmin, Tmax = config_dict["study_options"]["parameter_range"]

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

Looking at the given code snippet, all users need to do is:

Import the SensitivityAnalysisServer class.
Import the default parameter sampling strategies from the ParameterSamplerType Enum.
Use the config_dict options to access all the configuration settings defined in the JSON file. For example, in this code, we load the temperature min and max values for the Heat-PDE use case.
Finally, call self.set_parameter_sampler, which accepts either the pre-defined Enum values or a custom sampler class type.

Important

In the given example, ParameterSamplerType.RANDOM_UNIFORM is used. However, users have the flexibility to choose their own parameter sampling strategy. For more information on setting up a custom parameter sampler in a server class, we recommend reading the Design of Experiments (DoE) guide.

If everything is set up correctly, users should be able to run the SA study now.

Deep-Learning Server¶

Melissa's DeepMelissaServer class handles several key aspects, including:

Initializing a default TensorBoard logger, accessible as self.tb_logger.
Creating an IterableDataset instance for the specified buffer and framework.
Generating a buffer instance from the iterable dataset.
Setting up a training dataloader.
Implementing a framework-agnostic training loop with optional periodic validation and checkpointing.

Following code snippet showcases an example provided in examples/heat-pde/heat-pde-dl/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


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]
        )

        # user-defined method for getting validation dataloader
        self.valid_dataloader = self.get_validation_dataloader()

    @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.tb_logger.log_scalar("Loss/train", loss.item(), batch_idx)


    @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.data[field].astype(np.float32))

        return x, y


    class MyModel(torch.nn.Module):
        def __init__(self, input_features, output_features, output_dim):
            super().__init__()
            self.output_dim = output_dim
            self.output_features = output_features
            self.hidden_features = 256
            self.net = torch.nn.Sequential(
                torch.nn.Linear(input_features, self.hidden_features),
                torch.nn.ReLU(),
                torch.nn.Linear(self.hidden_features, self.hidden_features),
                torch.nn.ReLU(),
                torch.nn.Linear(self.hidden_features, output_features * output_dim),
            )

        def forward(self, x):
            y = self.net(x)
            return y

Modifications for training¶

Define a parameter sampling strategy, as explained in the Sensitivity Analysis Server section.
Assuming your model architecture is already implemented, override prepare_training_attributes to return the model and optimizer as a tuple.

Note

When working on DataDistributedParallel with PyTorch, users must call self.wrap_model_ddp on the model instance to convert it into a DDP-compatible model.

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 processes the transformed data (batch) and takes the current batch index (batch_idx) as input.

Note

If neither TorchServer nor TensorflowServer is used for training, then GeneralDataLoader and MelissaIterableDataset instances will be used, by default.

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. In the provided code snippet, self.valid_dataloader is initialized inside the __init__ method of the server class.
Override the validation_step method, which takes the validation data (batch) from self.valid_dataloader, the validation batch index (valid_batch_idx), and the training batch index (batch_idx) as inputs.

Note

Validation loop executes only when the condition batch_idx > 0 and (batch_idx + 1) % self.nb_batches_update == 0 is satisfied.

More Control (Optional)¶

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

Method	Description
`on_train_start()`	Called at the start of training.
`on_train_end()`	Called at the end of training.
`on_batch_start(batch_idx)`	Called at the start of a batch iteration.
`on_batch_end(batch_idx)`	Called at the end of a batch iteration.
`on_validation_start(batch_idx)`	Called at the start of validation.
`on_validation_end(batch_idx)`	Called at the end of validation.

Warning

Advanced users familiar with the DeepMelissaServer class can override the train and validation methods of the super class instead of using the hook or step methods described above.

Exposed properties¶

Users can access (with self) the following server properties while defining their server classes:

Property	Description
`tb_logger`	Provides access to the TensorBoard logger instance.
`buffer`	Returns the buffer instance.
`optimizer`	Gets or sets the optimizer. Must be set using `prepare_training_attributes`.
`model`	Gets or sets the model. Must be set using `prepare_training_attributes`.
`dataset`	Gets or sets the dataset created from the buffer.
`valid_dataloader`	Gets or sets the validation dataloader. Must be set by the user.

Finally, users can execute their study by following the approach outlined in the Running Your First DL Study section.