File size: 2,067 Bytes
b84549f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
from dataclasses import dataclass
from typing import List, Optional
from .base import ConfigBase
from .common import TrainingServiceConfig
from . import util
__all__ = [
'FrameworkControllerConfig',
'FrameworkControllerRoleConfig',
'FrameworkControllerNfsConfig',
'FrameworkControllerAzureStorageConfig'
]
@dataclass(init=False)
class _FrameworkControllerStorageConfig(ConfigBase):
storage: str
server: Optional[str] = None
path: Optional[str] = None
azure_account: Optional[str] = None
azure_share: Optional[str] = None
key_vault: Optional[str] = None
key_vault_secret: Optional[str] = None
@dataclass(init=False)
class FrameworkControllerNfsConfig(ConfigBase):
storage: str = 'nfs'
server: str
path: str
@dataclass(init=False)
class FrameworkControllerAzureStorageConfig(ConfigBase):
storage: str = 'azureStorage'
azure_account: str
azure_share: str
key_vault: str
key_vault_secret: str
@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: str
attempt_completion_min_failed_tasks: int
attempt_completion_min_succeeded_tasks: int
@dataclass(init=False)
class FrameworkControllerConfig(TrainingServiceConfig):
platform: str = 'frameworkcontroller'
service_account_name: str
storage: _FrameworkControllerStorageConfig
task_roles: List[FrameworkControllerRoleConfig]
def __init__(self, **kwargs):
kwargs = util.case_insensitive(kwargs)
kwargs['storage'] = util.load_config(_FrameworkControllerStorageConfig, kwargs.get('storage'))
kwargs['taskroles'] = util.load_config(FrameworkControllerRoleConfig, kwargs.get('taskroles'))
super().__init__(**kwargs)
_validation_rules = {
'platform': lambda value: (value == 'frameworkcontroller', 'cannot be modified')
}
|