# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

"""
Configuration for FrameworkController training service.

Check the reference_ for explaination of each field.

You may also want to check `FrameworkController training service doc`_.

.. _reference: https://nni.readthedocs.io/en/stable/reference/experiment_config.html

.. _FrameworkController training service doc: https://nni.readthedocs.io/en/stable/TrainingService/FrameworkControllerMode.html

"""

__all__ = ['FrameworkControllerConfig', 'FrameworkControllerRoleConfig', 'FrameworkAttemptCompletionPolicy']

from dataclasses import dataclass
from typing import List, Optional, Union

from ..base import ConfigBase
from ..training_service import TrainingServiceConfig
from .k8s_storage import K8sStorageConfig

@dataclass(init=False)
class FrameworkAttemptCompletionPolicy(ConfigBase):
    min_failed_task_count: int
    min_succeed_task_count: int

@dataclass(init=False)
class FrameworkControllerRoleConfig(ConfigBase):
    name: str
    docker_image: str = 'msranni/nni:latest'
    task_number: int
    command: str
    gpu_number: int
    cpu_number: int
    memory_size: Union[str, int]
    framework_attempt_completion_policy: FrameworkAttemptCompletionPolicy

@dataclass(init=False)
class FrameworkControllerConfig(TrainingServiceConfig):
    platform: str = 'frameworkcontroller'
    storage: K8sStorageConfig
    service_account_name: Optional[str]
    task_roles: List[FrameworkControllerRoleConfig]
    reuse_mode: Optional[bool] = True

    def _canonicalize(self, parents):
        super()._canonicalize(parents)
        # framework controller does not need these fields, set empty string for type check
        if self.trial_command is None:
            self.trial_command = ''