Feature: Disable scheduling on group of nodes (#540)

Author: timotheeguerin

aztk/models/__init__.py

@@ -15,4 +15,5 @@
 from .vm_image import VmImage
 from .software import Software
 from .cluster import Cluster
+from .scheduling_target import SchedulingTarget
 from .plugins import *

aztk/models/cluster_configuration.py

@@ -1,13 +1,13 @@
 import aztk.error as error
 from aztk.core.models import Model, fields
-from aztk.utils import deprecate, helpers
+from aztk.utils import deprecated,deprecate, helpers
 
 from .custom_script import CustomScript
 from .file_share import FileShare
 from .plugins import PluginConfiguration
 from .toolkit import Toolkit
 from .user_configuration import UserConfiguration
-
+from .scheduling_target import SchedulingTarget
 
 class ClusterConfiguration(Model):
     """
@@ -37,6 +37,7 @@ class ClusterConfiguration(Model):
     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:
@@ -49,6 +50,26 @@ def __init__(self, *args, **kwargs):
 
         super().__init__(*args, **kwargs)
 
+    @property
+    @deprecated()
+    def vm_count(self):
+        return self.size
+
+    @vm_count.setter
+    @deprecated()
+    def vm_count(self, value):
+        self.size = value
+
+    @property
+    @deprecated()
+    def vm_low_pri_count(self):
+        return self.size_low_priority
+
+    @vm_low_pri_count.setter
+    @deprecated()
+    def vm_low_pri_count(self, value):
+        self.size_low_priority = value
+
     def mixed_mode(self) -> bool:
         """
         Return:
@@ -81,3 +102,6 @@ def __validate__(self) -> bool:
 
         if self.custom_scripts:
             deprecate("Custom scripts are DEPRECATED and will be removed in 0.8.0. Use plugins instead See https://aztk.readthedocs.io/en/v0.7.0/15-plugins.html")
+
+        if self.scheduling_target == SchedulingTarget.Dedicated and self.vm_count == 0:
+            raise error.InvalidModelError("Scheduling target cannot be Dedicated if dedicated vm size is 0")

aztk/models/scheduling_target.py

@@ -0,0 +1,22 @@
+from enum import Enum
+
+class SchedulingTarget(Enum):
+    """
+    Target where task will get scheduled.
+    For spark this is where the driver will live.
+    """
+
+    Master = "master"
+    """
+    Only master is allowed to run task
+    """
+
+    Dedicated = "dedicated"
+    """
+    Any dedicated node is allowed to run task(Default)
+    """
+
+    Any = "any"
+    """
+    Any node(Not reconmmended if using low pri)
+    """

aztk/node_scripts/core/__init__.py

@@ -0,0 +1 @@
+from .logger import log

aztk/node_scripts/core/config.py

@@ -4,6 +4,7 @@
 import azure.batch.batch_service_client as batch
 import azure.storage.blob as blob
 import azure.batch.batch_auth as batchauth
+from core import log
 from azure.common.credentials import ServicePrincipalCredentials
 from azure.mgmt.batch import BatchManagementClient
 from azure.mgmt.storage import StorageManagementClient
@@ -28,7 +29,7 @@
 cluster_id = os.environ.get("AZTK_CLUSTER_ID")
 pool_id = os.environ["AZ_BATCH_POOL_ID"]
 node_id = os.environ["AZ_BATCH_NODE_ID"]
-is_dedicated = os.environ["AZ_BATCH_NODE_IS_DEDICATED"]
+is_dedicated = os.environ["AZ_BATCH_NODE_IS_DEDICATED"] == "true"
 
 spark_web_ui_port = os.environ["SPARK_WEB_UI_PORT"]
 spark_worker_ui_port = os.environ["SPARK_WORKER_UI_PORT"]
@@ -87,7 +88,7 @@ def get_batch_client() -> batch.BatchServiceClient:
 batch_client = get_batch_client()
 blob_client = get_blob_client()
 
-logging.info("Pool id is %s", pool_id)
-logging.info("Node id is %s", node_id)
-logging.info("Batch account name %s", batch_account_name)
-logging.info("Is dedicated %s", is_dedicated)
+log.info("Pool id is %s", pool_id)
+log.info("Node id is %s", node_id)
+log.info("Batch account name %s", batch_account_name)
+log.info("Is dedicated %s", is_dedicated)

aztk/node_scripts/core/logger.py

@@ -0,0 +1,13 @@
+import sys
+import logging
+
+log = logging.getLogger("aztk.node-agent")
+
+DEFAULT_FORMAT = '%(message)s'
+
+def setup_logging():
+    for handler in logging.root.handlers[:]:
+        logging.root.removeHandler(handler)
+
+    log.setLevel(logging.INFO)
+    logging.basicConfig(stream=sys.stdout, format=DEFAULT_FORMAT)

aztk/node_scripts/install/install.py

@@ -4,6 +4,8 @@
 import wait_until_master_selected
 from aztk.models.plugins import PluginTarget
 from aztk.internal import cluster_data
+from .node_scheduling import setup_node_scheduling
+
 
 def read_cluster_config():
     data = cluster_data.ClusterData(config.blob_client, config.cluster_id)
@@ -41,6 +43,8 @@ def setup_host(docker_repo: str):
 
     cluster_conf = read_cluster_config()
 
+    setup_node_scheduling(client, cluster_conf, is_master)
+
     #TODO pass azure file shares
     spark_container.start_spark_container(
         docker_repo=docker_repo,

aztk/node_scripts/install/node_scheduling.py

@@ -0,0 +1,58 @@
+from azure import batch
+from aztk.models import ClusterConfiguration, SchedulingTarget
+from core import config, log
+
+def disable_scheduling(batch_client: batch.BatchServiceClient):
+    """
+    Disable scheduling for the current node
+    """
+    pool_id = config.pool_id
+    node_id = config.node_id
+
+    node = batch_client.compute_node.get(pool_id, node_id)
+    if node.scheduling_state == batch.models.SchedulingState.enabled:
+        batch_client.compute_node.disable_scheduling(pool_id, node_id)
+        log.info("Disabled task scheduling for this node")
+    else:
+        log.info("Task scheduling is already disabled for this node")
+
+def enable_scheduling(batch_client: batch.BatchServiceClient):
+    """
+    Disable scheduling for the current node
+    """
+    pool_id = config.pool_id
+    node_id = config.node_id
+
+    node = batch_client.compute_node.get(pool_id, node_id)
+    if node.scheduling_state == batch.models.SchedulingState.disabled:
+        batch_client.compute_node.disable_scheduling(pool_id, node_id)
+        log.info("Enabled task scheduling for this node")
+    else:
+        log.info("Task scheduling is already enabled for this node")
+
+
+def setup_node_scheduling(
+        batch_client: batch.BatchServiceClient,
+        cluster_config: ClusterConfiguration,
+        is_master: bool):
+
+    is_dedicated = config.is_dedicated
+    enable = False
+    log.info("Resolving scheduling for this node")
+    log.info("  Scheduling target: %s", cluster_config.scheduling_target and cluster_config.scheduling_target.value)
+    log.info("  Is node dedicated: %s", str(is_dedicated))
+    log.info("  Is node master:    %s", str(is_master))
+
+    if cluster_config.scheduling_target == SchedulingTarget.Any or cluster_config.scheduling_target is None:
+        enable = True
+    elif cluster_config.scheduling_target == SchedulingTarget.Dedicated and is_dedicated:
+        enable = True
+    elif cluster_config.scheduling_target == SchedulingTarget.Master and is_master:
+        enable = True
+
+    if enable:
+        log.info("Scheduling will be enabled on this node as it satifies the right conditions")
+        enable_scheduling(batch_client)
+    else:
+        log.info("Scheduling will be disabled on this node as it does NOT satifies the right conditions")
+        disable_scheduling(batch_client)

aztk/node_scripts/main.py

@@ -1,6 +1,8 @@
 import sys
-import aztk.spark
 from install import install
+from core import logger
+
+
 
 def run():
     if len(sys.argv) < 2:
@@ -18,4 +20,5 @@ def run():
 
 
 if __name__ == "__main__":
+    logger.setup_logging()
     run()

aztk/spark/client.py

@@ -21,8 +21,6 @@ class Client(BaseClient):
     Args:
         secrets_config(aztk.spark.models.models.SecretsConfiguration): Configuration with all the needed credentials
     """
-    def __init__(self, secrets_config):
-        super().__init__(secrets_config)
 
     def create_cluster(self, cluster_conf: models.ClusterConfiguration, wait: bool = False):
         """
@@ -35,6 +33,7 @@ def create_cluster(self, cluster_conf: models.ClusterConfiguration, wait: bool =
         Returns:
             aztk.spark.models.Cluster
         """
+        cluster_conf = _apply_default_for_cluster_config(cluster_conf)
         cluster_conf.validate()
 
         cluster_data = self._get_cluster_data(cluster_conf.cluster_id)
@@ -200,8 +199,9 @@ def cluster_download(self, cluster_id: str, source_path: str, destination_path:
     '''
         job submission
     '''
-    def submit_job(self, job_configuration):
+    def submit_job(self, job_configuration: models.JobConfiguration):
         try:
+            job_configuration = _apply_default_for_job_config(job_configuration)
             job_configuration.validate()
             cluster_data = self._get_cluster_data(job_configuration.id)
             node_data =  NodeData(job_configuration.to_cluster_config()).add_core().done()
@@ -319,3 +319,23 @@ def run_cluster_diagnostics(self, cluster_id, output_directory):
             return output
         except batch_error.BatchErrorException as e:
             raise error.AztkError(helpers.format_batch_exception(e))
+
+
+def _default_scheduling_target(vm_count: int):
+    if vm_count == 0:
+        return models.SchedulingTarget.Any
+    else:
+        return models.SchedulingTarget.Dedicated
+
+def _apply_default_for_cluster_config(configuration: models.ClusterConfiguration):
+    cluster_conf = models.ClusterConfiguration()
+    cluster_conf.merge(configuration)
+    if cluster_conf.scheduling_target is None:
+        cluster_conf.scheduling_target = _default_scheduling_target(cluster_conf.size)
+    return cluster_conf
+
+def _apply_default_for_job_config(job_conf: models.JobConfiguration):
+    if job_conf.scheduling_target is None:
+        job_conf.scheduling_target = _default_scheduling_target(job_conf.max_dedicated_nodes)
+
+    return job_conf