class PluginConfiguration(Model):
"""
Plugin manifest that should be returned in the main.py of your plugin
Args
name: Name of the plugin. Used to reference the plugin
runOn: Where the plugin should run
execute: Path to the file to execute(This must match the target of one of the files)
files: List of files to upload
args: List of arguments to pass to the executing script
env: Dict of environment variables to pass to the script
"""
name = fields.String()
files = fields.List(PluginFile)
execute = fields.String()
args = fields.List(default=[])
env = fields.List(default=[])
target = fields.Enum(PluginTarget, default=PluginTarget.SparkContainer)
target_role = fields.Enum(PluginTargetRole,
default=PluginTargetRole.Master)
ports = fields.List(PluginPort, default=[])
def has_arg(self, name: str):
for x in self.args:
if x.name == name:
return True
return False
class PluginReference(Model):
"""
Contains the configuration to use a plugin
Args:
name (str): Name of the plugin(Must be the name of one of the provided plugins if no script provided)
script (str): Path to a custom script to run as the plugin
target_role (PluginTarget): Target for the plugin. Default to SparkContainer.
This can only be used if providing a script
target_role (PluginTargetRole): Target role default to All nodes. This can only be used if providing a script
args: (dict): If using name this is the arguments to pass to the plugin
"""
name = fields.String(default=None)
script = fields.String(default=None)
target = fields.Enum(PluginTarget, default=None)
target_role = fields.Enum(PluginTargetRole, default=None)
args = fields.Field(default=None)
def get_plugin(self) -> PluginConfiguration:
self.validate()
if self.script:
return self._plugin_from_script()
return plugin_manager.get_plugin(self.name, self.args)
def __validate__(self):
if not self.name and not self.script:
raise InvalidModelError(
"Plugin must either specify a name of an existing plugin or the path to a script."
)
if self.script and not os.path.isfile(self.script):
raise InvalidModelError(
"Plugin script file doesn't exists: '{0}'".format(self.script))
def _plugin_from_script(self):
script_filename = os.path.basename(self.script)
name = self.name or os.path.splitext(script_filename)[0]
return PluginConfiguration(
name=name,
execute=script_filename,
target=self.target,
target_role=self.target_role or PluginConfiguration,
files=[
PluginFile(script_filename, self.script),
],
)
class ClusterConfiguration(Model):
"""
Cluster configuration model
Args:
cluster_id (str): Id of the Aztk cluster
toolkit (aztk.models.Toolkit): Toolkit to be used for this cluster
size (int): Number of dedicated nodes for this cluster
size_low_priority (int): Number of low priority nodes for this cluster
vm_size (int): Azure Vm size to be used for each node
subnet_id (str): Full resource id of the subnet to be used(Required for mixed mode clusters)
plugins (List[aztk.models.plugins.PluginConfiguration]): List of plugins to be used
file_shares (List[aztk.models.FileShare]): List of File shares to be used
user_configuration (aztk.models.UserConfiguration): Configuration of the user to
be created on the master node to ssh into.
"""
cluster_id = fields.String()
toolkit = fields.Model(Toolkit)
size = fields.Integer(default=0)
size_low_priority = fields.Integer(default=0)
vm_size = fields.String()
subnet_id = fields.String(default=None)
plugins = fields.List(PluginConfiguration)
file_shares = fields.List(FileShare)
user_configuration = fields.Model(UserConfiguration, default=None)
scheduling_target = fields.Enum(SchedulingTarget, default=None)
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def mixed_mode(self) -> bool:
"""
Return:
if the pool is using mixed mode(Both dedicated and low priority nodes)
"""
return self.size > 0 and self.size_low_priority > 0
def gpu_enabled(self):
return helpers.is_gpu_enabled(self.vm_size)
def get_docker_repo(self):
return self.toolkit.get_docker_repo(self.gpu_enabled())
def get_docker_run_options(self) -> str:
return self.toolkit.get_docker_run_options()
def __validate__(self) -> bool:
if self.size == 0 and self.size_low_priority == 0:
raise error.InvalidModelError(
"Please supply a valid (greater than 0) size or size_low_priority value either "
"in the cluster.yaml configuration file or with a parameter (--size or --size-low-priority)"
)
if self.vm_size is None:
raise error.InvalidModelError(
"Please supply a vm_size in either the cluster.yaml configuration file or with a parameter (--vm-size)"
)
if self.mixed_mode() and not self.subnet_id:
raise error.InvalidModelError(
"You must configure a VNET to use AZTK in mixed mode (dedicated and low priority nodes). "
"Set the VNET's subnet_id in your cluster.yaml or with a parameter (--subnet-id)."
)
class SimpleStateModel(Model):
state = fields.Enum(UserState)
class User(Model):
info = fields.Model(UserInfo)
enabled = fields.Boolean(default=True)
state = fields.Enum(UserState, default=UserState.Ready)
class ClusterConfiguration(Model):
"""
Cluster configuration model
Args:
cluster_id (str): Id of the Aztk cluster
toolkit (aztk.models.Toolkit): Toolkit to be used for this cluster
size (int): Number of dedicated nodes for this cluster
size_low_priority (int): Number of low priority nodes for this cluster
vm_size (int): Azure Vm size to be used for each node
subnet_id (str): Full resource id of the subnet to be used(Required for mixed mode clusters)
plugins (List[aztk.models.plugins.PluginConfiguration]): List of plugins to be used
file_shares (List[aztk.models.FileShare]): List of File shares to be used
user_configuration (aztk.models.UserConfiguration): Configuration of the user to
be created on the master node to ssh into.
"""
cluster_id = fields.String()
toolkit = fields.Model(Toolkit)
size = fields.Integer(default=0)
size_low_priority = fields.Integer(default=0)
vm_size = fields.String()
subnet_id = fields.String(default=None)
plugins = fields.List(PluginConfiguration)
custom_scripts = fields.List(CustomScript)
file_shares = fields.List(FileShare)
user_configuration = fields.Model(UserConfiguration, default=None)
scheduling_target = fields.Enum(SchedulingTarget, default=None)
def __init__(self, *args, **kwargs):
if 'vm_count' in kwargs:
deprecate("0.9.0", "vm_count is deprecated for ClusterConfiguration.", "Please use size instead.")
kwargs['size'] = kwargs.pop('vm_count')
if 'vm_low_pri_count' in kwargs:
deprecate("vm_low_pri_count is deprecated for ClusterConfiguration.", "Please use size_low_priority instead.")
kwargs['size_low_priority'] = kwargs.pop('vm_low_pri_count')
super().__init__(*args, **kwargs)
@property
@deprecated("0.9.0")
def vm_count(self):
return self.size
@vm_count.setter
@deprecated("0.9.0")
def vm_count(self, value):
self.size = value
@property
@deprecated("0.9.0")
def vm_low_pri_count(self):
return self.size_low_priority
@vm_low_pri_count.setter
@deprecated("0.9.0")
def vm_low_pri_count(self, value):
self.size_low_priority = value
def mixed_mode(self) -> bool:
"""
Return:
if the pool is using mixed mode(Both dedicated and low priority nodes)
"""
return self.size > 0 and self.size_low_priority > 0
def gpu_enabled(self):
return helpers.is_gpu_enabled(self.vm_size)
def get_docker_repo(self):
return self.toolkit.get_docker_repo(self.gpu_enabled())
def __validate__(self) -> bool:
if self.size == 0 and self.size_low_priority == 0:
raise error.InvalidModelError(
"Please supply a valid (greater than 0) size or size_low_priority value either in the cluster.yaml configuration file or with a parameter (--size or --size-low-pri)"
)
if self.vm_size is None:
raise error.InvalidModelError(
"Please supply a vm_size in either the cluster.yaml configuration file or with a parameter (--vm-size)"
)
if self.mixed_mode() and not self.subnet_id:
raise error.InvalidModelError(
"You must configure a VNET to use AZTK in mixed mode (dedicated and low priority nodes). Set the VNET's subnet_id in your cluster.yaml."
)
if self.custom_scripts:
deprecate("0.9.0", "Custom scripts are DEPRECATED.", "Use plugins instead. See https://aztk.readthedocs.io/en/v0.7.0/15-plugins.html.")
if self.scheduling_target == SchedulingTarget.Dedicated and self.size == 0:
raise error.InvalidModelError("Scheduling target cannot be Dedicated if dedicated vm size is 0")
