Feature: Plugins (#387)

.editorconfig

@@ -4,10 +4,7 @@ indent_size = 4
 insert_final_newline = true
 trim_trailing_whitespace = true
 
-[*.json]
-indent_size = 2
-
-[*.yml]
+[*.{json,yml,yaml}]
 indent_size = 2
 
 [*.xml]

.vscode/launch.json

@@ -9,10 +9,10 @@
           "type": "python",
           "request": "launch",
           "stopOnEntry": false,
-          "internalConsoleOptions": "openOnSessionStart",
           "pythonPath": "${config:python.pythonPath}",
           "program": "${workspaceFolder}/cli/entrypoint.py",
           "cwd": "${workspaceFolder}",
+          "console": "integratedTerminal",
           "args": [
             "spark", "cluster", "list"
           ],
@@ -27,9 +27,9 @@
           "type": "python",
           "request": "launch",
           "stopOnEntry": false,
-          "internalConsoleOptions": "openOnSessionStart",
           "pythonPath": "${config:python.pythonPath}",
           "program": "${workspaceFolder}/cli/entrypoint.py",
+          "console": "integratedTerminal",
           "cwd": "${workspaceFolder}",
           "args": [
             "spark", "cluster", "create", "--id", "spark-debug"

aztk/client.py

@@ -1,6 +1,7 @@
 import asyncio
 import concurrent.futures
 import sys
+import yaml
 from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime, timedelta, timezone
 
@@ -25,6 +26,7 @@ def __init__(self, secrets_config: models.SecretsConfiguration):
         self.batch_client = azure_api.make_batch_client(secrets_config)
         self.blob_client = azure_api.make_blob_client(secrets_config)
 
+
     '''
     General Batch Operations
     '''
@@ -54,7 +56,7 @@ def __delete_pool_and_job(self, pool_id: str):
 
         return job_exists or pool_exists
 
-    def __create_pool_and_job(self, cluster_conf, software_metadata_key: str, start_task, VmImageModel):
+    def __create_pool_and_job(self, cluster_conf: models.ClusterConfiguration, software_metadata_key: str, start_task, VmImageModel):
         """
             Create a pool and job
             :param cluster_conf: the configuration object used to create the cluster
@@ -64,6 +66,7 @@ def __create_pool_and_job(self, cluster_conf, software_metadata_key: str, start_
             :param VmImageModel: the type of image to provision for the cluster
             :param wait: wait until the cluster is ready
         """
+        helpers.save_cluster_config(cluster_conf, self.blob_client)
         # reuse pool_id as job_id
         pool_id = cluster_conf.cluster_id
         job_id = cluster_conf.cluster_id

aztk/error.py

@@ -1,15 +1,31 @@
+"""
+Contains all errors used in Aztk.
+All error should inherit from `AztkError`
+"""
 
 
 
 class AztkError(Exception):
     def __init__(self, message: str = None):
-        super().__init__()
-        self.message = message
+        super().__init__(message)
 
 class ClusterNotReadyError(AztkError):
     pass
 
 class AzureApiInitError(AztkError):
-    def __init__(self, message: str = None):
-        super().__init__()
-        self.message = message
+    pass
+
+class InvalidPluginConfigurationError(AztkError):
+    pass
+
+class InvalidModelError(AztkError):
+    pass
+
+class MissingRequiredAttributeError(InvalidModelError):
+    pass
+
+class InvalidCustomScriptError(InvalidModelError):
+    pass
+
+class InvalidPluginReferenceError(InvalidModelError):
+    pass

aztk/internal/__init__.py

@@ -0,0 +1,5 @@
+"""
+Module containing classes used in the library but without any use for SDK user
+"""
+
+from .configuration_base import *

aztk/internal/configuration_base.py

@@ -0,0 +1,42 @@
+import yaml
+from aztk.error import AztkError
+
+class ConfigurationBase:
+    """
+    Base class for any configuration.
+    Include methods to help with validation
+    """
+
+    @classmethod
+    def from_dict(cls, args: dict):
+        """
+        Create a new model from a dict values
+        The dict is cleaned from null values and passed expanded to the constructor
+        """
+        try:
+            clean = dict((k, v) for k, v in args.items() if v)
+            return cls(**clean)
+        except TypeError as e:
+            pretty_args = yaml.dump(args, default_flow_style=False)
+            raise AztkError("{0} {1}\n{2}".format(cls.__name__, str(e), pretty_args))
+
+    def validate(self):
+        raise NotImplementedError("Validate not implemented")
+
+    def valid(self):
+        try:
+            self.validate()
+            return True
+        except AztkError:
+            return False
+
+    def _validate_required(self, attrs):
+        for attr in attrs:
+            if not getattr(self, attr):
+                raise AztkError("{0} missing {1}.".format(self.__class__.__name__, attr))
+
+    def _merge_attributes(self, other, attrs):
+        for attr in attrs:
+            val = getattr(other, attr)
+            if val is not None:
+                setattr(self, attr, val)

aztk/models/__init__.py

@@ -0,0 +1 @@
+from .models import *

aztk/models.py → aztk/models/models.py

@@ -1,37 +1,14 @@
 import io
 from typing import List
-import azure.batch.models as batch_models
-import aztk.utils.constants as constants
 from aztk import error
+from aztk.utils import constants
+import azure.batch.models as batch_models
+from aztk.models.plugins import PluginConfiguration
+from aztk.internal import ConfigurationBase
+import yaml
+import logging
 
 
-class ConfigurationBase:
-    """
-    Base class for any configuration.
-    Include methods to help with validation
-    """
-    def validate(self):
-        raise NotImplementedError("Validate not implemented")
-
-    def valid(self):
-        try:
-            self.validate()
-            return True
-        except error.AztkError:
-            return False
-
-    def _validate_required(self, attrs):
-        for attr in attrs:
-            if not getattr(self, attr):
-                raise error.AztkError(
-                    "{0} missing {1}.".format(self.__class__.__name__, attr))
-
-    def _merge_attributes(self, other, attrs):
-        for attr in attrs:
-            val = getattr(other, attr)
-            if val is not None:
-                setattr(self, attr, val)
-
 class FileShare:
     def __init__(self,
                  storage_account_name: str = None,
@@ -57,7 +34,10 @@ def __init__(self, name: str = None, script = None, run_on=None):
 
 
 class UserConfiguration(ConfigurationBase):
-    def __init__(self, username: str, ssh_key: str = None, password: str = None):
+    def __init__(self,
+                 username: str,
+                 ssh_key: str = None,
+                 password: str = None):
         self.username = username
         self.ssh_key = ssh_key
         self.password = password
@@ -69,21 +49,24 @@ def merge(self, other):
             "password",
         ])
 
+
+
 class ClusterConfiguration(ConfigurationBase):
     """
     Cluster configuration model
     """
-    def __init__(
-            self,
-            custom_scripts: List[CustomScript] = None,
-            file_shares: List[FileShare] = None,
-            cluster_id: str = None,
-            vm_count=0,
-            vm_low_pri_count=0,
-            vm_size=None,
-            subnet_id=None,
-            docker_repo: str=None,
-            user_configuration: UserConfiguration=None):
+
+    def __init__(self,
+                 custom_scripts: List[CustomScript] = None,
+                 file_shares: List[FileShare] = None,
+                 cluster_id: str = None,
+                 vm_count=0,
+                 vm_low_pri_count=0,
+                 vm_size=None,
+                 subnet_id=None,
+                 docker_repo: str = None,
+                 plugins: List[PluginConfiguration] = None,
+                 user_configuration: UserConfiguration = None):
         super().__init__()
         self.custom_scripts = custom_scripts
         self.file_shares = file_shares
@@ -94,6 +77,7 @@ def __init__(
         self.subnet_id = subnet_id
         self.docker_repo = docker_repo
         self.user_configuration = user_configuration
+        self.plugins = plugins
 
     def merge(self, other):
         """
@@ -110,6 +94,7 @@ def merge(self, other):
             "docker_repo",
             "vm_count",
             "vm_low_pri_count",
+            "plugins",
         ])
 
         if other.user_configuration:
@@ -118,6 +103,10 @@ def merge(self, other):
             else:
                 self.user_configuration = other.user_configuration
 
+        if self.plugins:
+            for plugin in self.plugins:
+                plugin.validate()
+
     def mixed_mode(self) -> bool:
         return self.vm_count > 0 and self.vm_low_pri_count > 0
 
@@ -133,15 +122,18 @@ def validate(self) -> bool:
 
         if self.vm_count == 0 and self.vm_low_pri_count == 0:
             raise error.AztkError(
-                "Please supply a valid (greater than 0) size or size_low_pri value either in the cluster.yaml configuration file or with a parameter (--size or --size-low-pri)")
+                "Please supply a valid (greater than 0) size or size_low_pri 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.AztkError(
-                "Please supply a vm_size in either the cluster.yaml configuration file or with a parameter (--vm-size)")
+                "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.AztkError(
-                "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.")
+                "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."
+            )
 
 
 class RemoteLogin:
@@ -217,11 +209,7 @@ def validate(self) -> bool:
 
 
 class DockerConfiguration(ConfigurationBase):
-    def __init__(
-            self,
-            endpoint=None,
-            username=None,
-            password=None):
+    def __init__(self, endpoint=None, username=None, password=None):
 
         self.endpoint = endpoint
         self.username = username
@@ -232,13 +220,12 @@ def validate(self):
 
 
 class SecretsConfiguration(ConfigurationBase):
-    def __init__(
-            self,
-            service_principal=None,
-            shared_key=None,
-            docker=None,
-            ssh_pub_key=None,
-            ssh_priv_key=None):
+    def __init__(self,
+                 service_principal=None,
+                 shared_key=None,
+                 docker=None,
+                 ssh_pub_key=None,
+                 ssh_priv_key=None):
         self.service_principal = service_principal
         self.shared_key = shared_key
         self.docker = docker
@@ -249,14 +236,16 @@ def __init__(
     def validate(self):
         if self.service_principal and self.shared_key:
             raise error.AztkError(
-                "Both service_principal and shared_key auth are configured, must use only one")
+                "Both service_principal and shared_key auth are configured, must use only one"
+            )
         elif self.service_principal:
             self.service_principal.validate()
         elif self.shared_key:
             self.shared_key.validate()
         else:
             raise error.AztkError(
-                "Neither service_principal and shared_key auth are configured, must use only one")
+                "Neither service_principal and shared_key auth are configured, must use only one"
+            )
 
     def is_aad(self):
         return self.service_principal is not None
@@ -270,7 +259,9 @@ def __init__(self, publisher, offer, sku):
 
 
 class Cluster:
-    def __init__(self, pool: batch_models.CloudPool, nodes: batch_models.ComputeNodePaged = None):
+    def __init__(self,
+                 pool: batch_models.CloudPool,
+                 nodes: batch_models.ComputeNodePaged = None):
         self.id = pool.id
         self.pool = pool
         self.nodes = nodes
@@ -288,11 +279,13 @@ def __init__(self, pool: batch_models.CloudPool, nodes: batch_models.ComputeNode
         self.target_dedicated_nodes = pool.target_dedicated_nodes
         self.target_low_pri_nodes = pool.target_low_priority_nodes
 
+
 class SSHLog():
     def __init__(self, output, node_id):
         self.output = output
         self.node_id = node_id
 
+
 class Software:
     """
         Enum with list of available softwares

aztk/models/plugins/__init__.py

@@ -0,0 +1,2 @@
+from .plugin_file import *
+from .plugin_configuration import *

aztk/models/plugins/internal/__init__.py

@@ -0,0 +1,2 @@
+from .plugin_manager import *
+from .plugin_reference import *