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()