fl_server_core.models ¶

class GlobalModel(Model):
    """
    Model class for global models.
    """

    name: CharField = CharField(max_length=256)
    """Name of the model."""
    description: TextField = TextField()
    """Description of the model."""
    # alternative to be postgres independent: create a new model for each nullable integer field
    # and map the corresponding list of integers to the model (but pay attention to the order)
    input_shape: ArrayField = ArrayField(IntegerField(null=True), null=True)
    """Input shape of the model."""
    preprocessing: BinaryField = BinaryField(null=True)
    """Preprocessing of the model."""

    def get_preprocessing_torch_model(self) -> torch.nn.Module:
        """
        Converts the preprocessing to a PyTorch model.

        Returns:
            torch.nn.Module: The PyTorch model.
        """
        return to_torch_module(self.preprocessing)

    def set_preprocessing_torch_model(self, value: torch.nn.Module):
        """
        Sets the preprocessing from a PyTorch model.

        Args:
            value (torch.nn.Module): The PyTorch model.
        """
        self.preprocessing = from_torch_module(value)

class LocalModel(Model):
    """
    Model class for local models.
    """

    base_model: ForeignKey = ForeignKey(GlobalModel, on_delete=CASCADE)
    """Base model of the local model."""
    sample_size: IntegerField = IntegerField()
    """Sample size of the local model."""

    def get_training(self) -> Optional["models.Training"]:
        """
        Gets the training associated with the base model.

        Returns:
            models.Training: The training associated with the base model.
        """
        return models.Training.objects.filter(model=self.base_model).first()

class MeanModel(GlobalModel):
    """
    Model class for mean models.
    """

    models: ManyToManyField = ManyToManyField(GlobalModel, related_name="mean_models")
    """Models of the mean model."""

    def get_torch_model(self) -> torch.nn.Module:
        """
        Converts the models to a PyTorch model.

        Returns:
            torch.nn.Module: The PyTorch model.
        """
        torch_models: List[torch.nn.Module] = [model.get_torch_model() for model in self.models.all()]
        model = MeanModule(torch_models)
        if all(is_torchscript_instance(m) for m in torch_models):
            return torch.jit.script(model)
        return model

    def set_torch_model(self, value: torch.nn.Module):
        """
        Sets the models from a PyTorch model.

        Args:
            value (torch.nn.Module): The PyTorch model.
        """
        raise NotImplementedError()

class Metric(models.Model):
    """
    Metric model class.
    """

    model: ForeignKey = ForeignKey(Model, on_delete=CASCADE)
    """Model associated with the metric."""
    identifier: CharField = CharField(max_length=64, null=True, blank=True)
    """Identifier of the metric."""
    key: CharField = CharField(max_length=32)
    """Key of the metric."""
    value_float: FloatField = FloatField(null=True, blank=True)
    """Float value of the metric."""
    value_binary: BinaryField = BinaryField(null=True, blank=True)
    """Binary value of the metric."""
    step: IntegerField = IntegerField(null=True, blank=True)
    """Step of the metric."""
    reporter: ForeignKey = ForeignKey(User, null=True, blank=True, on_delete=CASCADE)
    """User who reported the metric."""

    @property
    def value(self) -> float | bytes:
        """
        Value of the metric.

        Returns:
            float | bytes: The value of the metric.
        """
        if self.is_float():
            return self.value_float
        return self.value_binary

    @value.setter
    def value(self, value: float | int | bytes | Module | Tensor):
        """
        Setter for the value of the metric.

        Args:
            value (float | int | bytes | Module | Tensor): The value to set.
        """
        if isinstance(value, float):
            self.value_float = value
        elif isinstance(value, int):
            self.value_float = float(value)
        elif isinstance(value, (Module, Tensor)):
            self.value_binary = from_torch_module_or_tensor(value)
        else:
            self.value_binary = value

    @value.deleter
    def value(self):
        """
        Deleter for the value of the metric.
        """
        self.value_float = None
        self.value_binary = None

    def is_float(self) -> bool:
        """
        Check if the value of the metric is a float.

        Returns:
            bool: `True` if the value of the metric is a float, otherwise `False`.
        """
        return self.value_float is not None

    def is_binary(self) -> bool:
        """
        Check if the value of the metric is binary.

        Returns:
            bool: `True` if the value of the metric is binary, otherwise `False`.
        """
        return self.value_binary is not None

    def to_torch(self) -> Module | Tensor:
        """
        Convert the binary value of the metric to a torch module or tensor.

        Returns:
            Module | Tensor: The converted torch module or tensor.
        """
        return to_torch_module_or_tensor(self.value_binary)

class Model(PolymorphicModel):
    """
    Base model class for all types of model models.
    """

    id: UUIDField = UUIDField(primary_key=True, editable=False, default=uuid4)
    """Unique identifier for the model."""
    owner: ForeignKey = ForeignKey(User, on_delete=CASCADE)
    """User who owns the model."""
    round: IntegerField = IntegerField()
    """Round number of the model."""
    weights: BinaryField = BinaryField()
    """Weights of the model."""

    def is_global_model(self):
        """
        Checks if the model is a global model.

        Returns:
            bool: True if the model is a global model, False otherwise.
        """
        return isinstance(self, GlobalModel)

    def is_local_model(self):
        """
        Checks if the model is a local model.

        Returns:
            bool: True if the model is a local model, False otherwise.
        """
        return isinstance(self, LocalModel)

    def get_torch_model(self) -> torch.nn.Module:
        """
        Converts the model weights to a PyTorch model.

        Returns:
            torch.nn.Module: The PyTorch model.
        """
        return to_torch_module(self.weights)

    def set_torch_model(self, value: torch.nn.Module):
        """
        Sets the model weights from a PyTorch model.

        Args:
            value (torch.nn.Module): The PyTorch model.
        """
        self.weights = from_torch_module(value)

    def get_training(self) -> Optional["models.Training"]:
        """
        Gets the training associated with the model.

        Returns:
            models.Training: The training associated with the model.
        """
        return models.Training.objects.filter(model=self).first()

class SWAGModel(GlobalModel):
    """
    Model class for SWAG models.
    """

    swag_first_moment: BinaryField = BinaryField()
    """First moment of the SWAG model."""
    swag_second_moment: BinaryField = BinaryField()
    """Second moment of the SWAG model."""

    @property
    def first_moment(self) -> Tensor:
        """
        Gets the first moment of the SWAG model.

        Returns:
            Tensor: The first moment of the SWAG model.
        """
        return to_torch_tensor(self.swag_first_moment)

    @first_moment.setter
    def first_moment(self, value: Tensor):
        """
        Sets the first moment of the SWAG model.

        Args:
            value (Tensor): The first moment of the SWAG model.
        """
        self.swag_first_moment = from_torch_tensor(value)

    @property
    def second_moment(self) -> Tensor:
        """
        Gets the second moment of the SWAG model.

        Returns:
            Tensor: The second moment of the SWAG model.
        """
        return to_torch_tensor(self.swag_second_moment)

    @second_moment.setter
    def second_moment(self, value: Tensor):
        """
        Sets the second moment of the SWAG model.

        Args:
            value (Tensor): The second moment of the SWAG model.
        """
        self.swag_second_moment = from_torch_tensor(value)

class Training(models.Model):
    """
    Training model class.
    """

    id: UUIDField = UUIDField(primary_key=True, editable=False, default=uuid4)
    """Unique identifier for the training."""
    model: OneToOneField = OneToOneField(Model, on_delete=CASCADE)
    """Model used in the training."""
    actor: ForeignKey = ForeignKey(User, on_delete=CASCADE, related_name="actors")
    """User who is the actor of the training."""
    participants: ManyToManyField = ManyToManyField(User)
    """Users who are the participants of the training."""
    state: CharField = CharField(max_length=1, choices=TrainingState.choices)
    """State of the training."""
    target_num_updates: IntegerField = IntegerField()
    """Target number of updates for the training."""
    last_update = models.DateTimeField(auto_now=True)
    """Time of the last update."""
    uncertainty_method: CharField = CharField(
        max_length=32, choices=UncertaintyMethod.choices, default=UncertaintyMethod.NONE
    )
    """Uncertainty method used in the training."""
    aggregation_method: CharField = CharField(
        max_length=32, choices=AggregationMethod.choices, default=AggregationMethod.FED_AVG
    )
    """Aggregation method used in the training."""
    # HINT: https://docs.djangoproject.com/en/4.2/topics/db/queries/#querying-jsonfield
    options: JSONField = JSONField(default=dict, encoder=DjangoJSONEncoder)
    """Options for the training."""
    locked: BooleanField = BooleanField(default=False)
    """Flag indicating whether the training is locked."""

class User(AbstractUser, NotificationReceiver):
    """
    User class.

    Inherits from Django's AbstractUser and NotificationReceiver.
    """

    id: UUIDField = UUIDField(primary_key=True, editable=False, default=uuid4)
    """Unique identifier for the user."""
    actor: BooleanField = BooleanField(default=False)
    """Flag indicating whether the user is an actor."""
    client: BooleanField = BooleanField(default=False)
    """Flag indicating whether the user is a client."""
    message_endpoint: URLField = URLField()
    """Endpoint to send the message to."""