Allow submitting jobs into a VNET

aztk/client.py

@@ -179,6 +179,12 @@ def __submit_job(self,
             helpers.select_latest_verified_vm_image_with_node_agent_sku(
                 vm_image_model.publisher, vm_image_model.offer, vm_image_model.sku, self.batch_client)
 
+        # set up subnet if necessary
+        network_conf = None
+        if job_configuration.subnet_id:
+            network_conf = batch_models.NetworkConfiguration(
+                subnet_id=job_configuration.subnet_id)
+
         # set up a schedule for a recurring job
         auto_pool_specification = batch_models.AutoPoolSpecification(
             pool_lifetime_option=batch_models.PoolLifetimeOption.job_schedule,
@@ -195,6 +201,7 @@ def __submit_job(self,
                 auto_scale_evaluation_interval=timedelta(minutes=5),
                 start_task=start_task,
                 enable_inter_node_communication=True,
+                network_configuration=network_conf,
                 max_tasks_per_node=1,
                 metadata=[
                     batch_models.MetadataItem(

aztk/spark/models.py

@@ -178,7 +178,8 @@ def __init__(
             spark_configuration=None,
             docker_repo=None,
             max_dedicated_nodes=None,
-            max_low_pri_nodes=None):
+            max_low_pri_nodes=None,
+            subnet_id=None):
         self.id = id
         self.applications = applications
         self.custom_scripts = custom_scripts
@@ -188,6 +189,7 @@ def __init__(
         self.docker_repo = docker_repo
         self.max_dedicated_nodes = max_dedicated_nodes
         self.max_low_pri_nodes = max_low_pri_nodes
+        self.subnet_id = subnet_id
 
 
 class JobState():

cli/config.py

@@ -333,6 +333,7 @@ def __init__(self):
         self.spark_defaults_conf = None
         self.spark_env_sh = None
         self.core_site_xml = None
+        self.subnet_id = None
 
     def _merge_dict(self, config):
         config = config.get('job')
@@ -347,6 +348,7 @@ def _merge_dict(self, config):
             self.max_dedicated_nodes = cluster_configuration.get('size')
             self.max_low_pri_nodes = cluster_configuration.get('size_low_pri')
             self.custom_scripts = cluster_configuration.get('custom_scripts')
+            self.subnet_id = cluster_configuration.get('subnet_id')
 
         self.applications = config.get('applications')
 

cli/spark/endpoints/job/submit.py

@@ -56,7 +56,8 @@ def execute(args: typing.NamedTuple):
         vm_size=job_conf.vm_size,
         docker_repo=job_conf.docker_repo,
         max_dedicated_nodes=job_conf.max_dedicated_nodes,
-        max_low_pri_nodes=job_conf.max_low_pri_nodes
+        max_low_pri_nodes=job_conf.max_low_pri_nodes,
+        subnet_id=job_conf.subnet_id
     )
 
     #TODO: utils.print_job_conf(job_configuration)

config/job.yaml

@@ -9,6 +9,7 @@ job:
         vm_size: standard_d2_v2
         size: 2
         size_low_pri: 0
+        subnet_id: 
         docker_repo: # defaults to aztk/base:spark2.2.0
         # custom_scripts:
         #     - script: </path/to/script.sh or /path/to/script/directory/>