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 ServicePrincipalConfiguration(Model):
"""
Container class for AAD authentication
"""
tenant_id = fields.String()
client_id = fields.String()
credential = fields.String()
batch_account_resource_id = fields.String()
storage_account_resource_id = fields.String()
class DockerConfiguration(Model):
"""
Configuration for connecting to private docker
Args:
endpoint (str): Which docker endpoint to use. Default to docker hub.
username (str): Docker endpoint username
password (str): Docker endpoint password
"""
endpoint = fields.String(default=None)
username = fields.String(default=None)
password = fields.String(default=None)
class PluginFile(Model):
"""
Reference to a file for a plugin.
"""
target = fields.String()
local_path = fields.String()
def __init__(self, target: str = None, local_path: str = None):
super().__init__(target=target, local_path=local_path)
def content(self):
with open(self.local_path, "r", encoding='UTF-8') as f:
return f.read()
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 SparkConfiguration(Model):
spark_defaults_conf = fields.String(default=None)
spark_env_sh = fields.String(default=None)
core_site_xml = fields.String(default=None)
jars = fields.List()
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.ssh_key_pair = self.__generate_ssh_key_pair()
def __generate_ssh_key_pair(self):
key = RSA.generate(2048)
priv_key = key.exportKey("PEM")
pub_key = key.publickey().exportKey("OpenSSH")
return {"pub_key": pub_key, "priv_key": priv_key}
class SecretsConfiguration(Model):
service_principal = fields.Model(ServicePrincipalConfiguration, default=None)
shared_key = fields.Model(SharedKeyConfiguration, default=None)
docker = fields.Model(DockerConfiguration, default=None)
ssh_pub_key = fields.String(default=None)
ssh_priv_key = fields.String(default=None)
def __validate__(self):
if self.service_principal and self.shared_key:
raise InvalidModelError(
"Both service_principal and shared_key auth are configured, must use only one"
)
if not self.service_principal and not self.shared_key:
raise InvalidModelError(
"Neither service_principal and shared_key auth are configured, must use only one"
)
def is_aad(self):
return self.service_principal is not None
class SharedKeyConfiguration(Model):
"""
Container class for shared key authentication
"""
batch_account_name = fields.String()
batch_account_key = fields.String()
batch_service_url = fields.String()
storage_account_name = fields.String()
storage_account_key = fields.String()
storage_account_suffix = fields.String()
class Task(Model):
id = fields.String()
node_id = fields.String(default=None)
state = fields.String(default=None)
state_transition_time = fields.String(default=None)
command_line = fields.String(default=None)
exit_code = fields.Integer(default=None)
start_time = fields.Datetime(datetime, default=None)
end_time = fields.Datetime(datetime, default=None)
failure_info = fields.String(default=None)
class TextPluginFile(Model):
"""
Reference to a file for a plugin.
Args:
target (str): Where should the file be uploaded relative to the plugin working dir
content (str|io.StringIO): Content of the file. Can either be a string or a StringIO
"""
target = fields.String()
def __init__(self, target: str, content: Union[str, io.StringIO]):
super().__init__(target=target)
if isinstance(content, str):
self._content = content
else:
self._content = content.getValue()
def content(self):
return self._content
class SimpleMergeModel(Model):
name = fields.String()
enabled = fields.Boolean(default=True)
class CustomScript(Model):
name = fields.String()
script = fields.String()
run_on = fields.String()
class Toolkit(Model):
"""
Toolkit for a cluster.
This will help pick the docker image needed
Args:
software (str): Name of the toolkit(spark)
version (str): Version of the toolkit
environment (str): Which environment to use for this toolkit
environment_version (str): If there is multiple version for an environment you can specify which one
docker_repo (str): Optional docker repo
"""
software = fields.String()
version = fields.String()
environment = fields.String(default=None)
environment_version = fields.String(default=None)
docker_repo = fields.String(default=None)
def __validate__(self):
if self.software not in TOOLKIT_MAP:
raise InvalidModelError(
"Toolkit '{0}' is not in the list of allowed toolkits {1}".
format(self.software, list(TOOLKIT_MAP.keys())))
toolkit_def = TOOLKIT_MAP[self.software]
if self.version not in toolkit_def.versions:
raise InvalidModelError(
"Toolkit '{0}' with version '{1}' is not available. Use one of: {2}"
.format(self.software, self.version, toolkit_def.versions))
if self.version == "1.6":
deprecate("0.9.0",
"Spark version 1.6 is being deprecated for Aztk.",
"Please use 2.1 and above.")
if self.environment:
if self.environment not in toolkit_def.environments:
raise InvalidModelError(
"Environment '{0}' for toolkit '{1}' is not available. Use one of: {2}"
.format(self.environment, self.software,
list(toolkit_def.environments.keys())))
env_def = toolkit_def.environments[self.environment]
if self.environment_version and self.environment_version not in env_def.versions:
raise InvalidModelError(
"Environment '{0}' version '{1}' for toolkit '{2}' is not available. Use one of: {3}"
.format(self.environment, self.environment_version,
self.software, env_def.versions))
def get_docker_repo(self, gpu: bool):
if self.docker_repo:
return self.docker_repo
repo = "aztk/{0}".format(self.software)
return "{repo}:{tag}".format(
repo=repo,
tag=self._get_docker_tag(gpu),
)
def _get_docker_tag(self, gpu: bool):
environment = self.environment or "base"
environment_def = self._get_environment_definition()
environment_version = self.environment_version or (
environment_def and environment_def.default)
array = [
"v{docker_image_version}".format(
docker_image_version=constants.DOCKER_IMAGE_VERSION),
"{toolkit}{version}".format(toolkit=self.software,
version=self.version),
]
if self.environment:
array.append("{0}{1}".format(environment, environment_version))
array.append("gpu" if gpu else "base")
return '-'.join(array)
def _get_environment_definition(self) -> ToolkitEnvironmentDefinition:
toolkit = TOOLKIT_MAP.get(self.software)
if toolkit:
return toolkit.environments.get(self.environment or "base")
return None
class SimpleNameModel(Model):
name = fields.String()
class ServiceUser(User):
service = fields.String()
class ComplexModel(Model):
model_id = fields.String()
info = fields.Model(UserInfo, merge_strategy=ModelMergeStrategy.Override)
class Toolkit(Model):
"""
Toolkit for a cluster.
This will help pick the docker image needed
Args:
software (str): Name of the toolkit(spark)
version (str): Version of the toolkit
environment (str): Which environment to use for this toolkit
environment_version (str): If there is multiple version for an environment you can specify which one
docker_repo (str): Optional docker repo
docker_run_options (str): Optional command-line options for `docker run`
"""
software = fields.String()
version = fields.String()
environment = fields.String(default=None)
environment_version = fields.String(default=None)
docker_repo = fields.String(default=None)
docker_run_options = fields.String(default=None)
def __validate__(self):
if self.software not in TOOLKIT_MAP:
raise InvalidModelError("Toolkit '{0}' is not in the list of allowed toolkits {1}".format(
self.software, list(TOOLKIT_MAP.keys())))
toolkit_def = TOOLKIT_MAP[self.software]
if self.version not in toolkit_def.versions:
raise InvalidModelError("Toolkit '{0}' with version '{1}' is not available. Use one of: {2}".format(
self.software, self.version, toolkit_def.versions))
if self.environment:
if self.environment not in toolkit_def.environments:
raise InvalidModelError("Environment '{0}' for toolkit '{1}' is not available. Use one of: {2}".format(
self.environment, self.software, list(toolkit_def.environments.keys())))
env_def = toolkit_def.environments[self.environment]
if self.environment_version and self.environment_version not in env_def.versions:
raise InvalidModelError(
"Environment '{0}' version '{1}' for toolkit '{2}' is not available. Use one of: {3}".format(
self.environment, self.environment_version, self.software, env_def.versions))
if self.docker_run_options:
invalid_character = re.search(r'[^A-Za-z0-9 _./:=\-"]', self.docker_run_options)
if invalid_character:
raise InvalidModelError(
"Docker run options contains invalid character '{0}'. Only A-Z, a-z, 0-9, space, hyphen (-), "
"underscore (_), period (.), forward slash (/), colon (:), equals(=), comma (,), and "
'double quote (") are allowed.'.format(invalid_character.group(0)))
def get_docker_repo(self, gpu: bool):
if self.docker_repo:
return self.docker_repo
repo = "aztk/{0}".format(self.software)
return "{repo}:{tag}".format(repo=repo, tag=self._get_docker_tag(gpu))
def get_docker_run_options(self):
return self.docker_run_options
def _get_docker_tag(self, gpu: bool):
environment = self.environment or "base"
environment_def = self._get_environment_definition()
environment_version = self.environment_version or (environment_def and environment_def.default)
array = [
"v{docker_image_version}".format(docker_image_version=constants.DOCKER_IMAGE_VERSION),
"{toolkit}{version}".format(toolkit=self.software, version=self.version),
]
if self.environment:
array.append("{0}{1}".format(environment, environment_version))
array.append("gpu" if gpu else "base")
return "-".join(array)
def _get_environment_definition(self) -> ToolkitEnvironmentDefinition:
toolkit = TOOLKIT_MAP.get(self.software)
if toolkit:
return toolkit.environments.get(self.environment or "base")
return None
class FileShare(Model):
storage_account_name = fields.String()
storage_account_key = fields.String()
file_share_path = fields.String()
mount_path = fields.String()
class SimpleRequiredModel(Model):
name = fields.String()
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")
class UserConfiguration(Model):
username = fields.String()
ssh_key = fields.String(default=None)
password = fields.String(default=None)
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 UserInfo(Model):
name = fields.String()
age = fields.Integer()