Fix: Worker on master flag ignored and standardize boolean environment (#514)

Author: timotheeguerin

aztk/node_scripts/install/install.py

@@ -18,13 +18,13 @@ def setup_host(docker_repo: str):
     client = config.batch_client
 
     create_user.create_user(batch_client=client)
-    if os.environ['AZ_BATCH_NODE_IS_DEDICATED'] == "true" or os.environ['AZTK_MIXED_MODE'] == "False":
+    if os.environ['AZ_BATCH_NODE_IS_DEDICATED'] == "true" or os.environ['AZTK_MIXED_MODE'] == "false":
         is_master = pick_master.find_master(client)
     else:
         is_master = False
         wait_until_master_selected.main()
 
-    is_worker = not is_master or os.environ["AZTK_WORKER_ON_MASTER"]
+    is_worker = not is_master or os.environ.get("AZTK_WORKER_ON_MASTER") == "true"
     master_node_id = pick_master.get_master_node_id(config.batch_client.pool.get(config.pool_id))
     master_node = config.batch_client.compute_node.get(config.pool_id, master_node_id)
 
@@ -61,9 +61,6 @@ def setup_spark_container():
     spark.setup_conf()
     print("Done copying spark setup config")
 
-    master_node_id = pick_master.get_master_node_id(config.batch_client.pool.get(config.pool_id))
-    master_node = config.batch_client.compute_node.get(config.pool_id, master_node_id)
-
     spark.setup_connection()
 
     if is_master:

aztk/node_scripts/setup_host.sh

@@ -38,7 +38,7 @@ sudo curl -L https://github.com/docker/compose/releases/download/1.19.0/docker-c
 sudo chmod +x /usr/local/bin/docker-compose
 echo "Done installing Docker-Compose"
 
-if [ $AZTK_GPU_ENABLED == "True" ]; then
+if [ $AZTK_GPU_ENABLED == "true" ]; then
     echo "running nvidia install"
     sudo apt-get -y install nvidia-384
     sudo apt-get -y install nvidia-modprobe

aztk/spark/helpers/create_cluster.py

@@ -13,13 +13,9 @@
 
 def _get_aztk_environment(worker_on_master, mixed_mode):
     envs = []
-    envs.append(batch_models.EnvironmentSetting(name="AZTK_MIXED_MODE", value=mixed_mode))
-    if worker_on_master is not None:
-        envs.append(batch_models.EnvironmentSetting(
-                name="AZTK_WORKER_ON_MASTER", value=worker_on_master))
-    else:
-        envs.append(batch_models.EnvironmentSetting(
-                name="AZTK_WORKER_ON_MASTER", value=False))
+    envs.append(batch_models.EnvironmentSetting(name="AZTK_MIXED_MODE", value=helpers.bool_env(mixed_mode)))
+    envs.append(batch_models.EnvironmentSetting(
+            name="AZTK_WORKER_ON_MASTER", value=helpers.bool_env(worker_on_master)))
     return envs
 
 def __get_docker_credentials(spark_client):
@@ -152,7 +148,7 @@ def generate_cluster_start_task(
         batch_models.EnvironmentSetting(
             name="SPARK_SUBMIT_LOGS_FILE", value=spark_submit_logs_file),
         batch_models.EnvironmentSetting(
-            name="AZTK_GPU_ENABLED", value=gpu_enabled),
+            name="AZTK_GPU_ENABLED", value=helpers.bool_env(gpu_enabled)),
     ] + __get_docker_credentials(spark_client) + _get_aztk_environment(worker_on_master, mixed_mode)
 
     # start task command

aztk/utils/helpers.py

@@ -400,3 +400,24 @@ def read_cluster_config(cluster_id: str, blob_client: blob.BlockBlobService):
         logging.warn(
             "Cluster %s contains invalid cluster configuration in blob",
             cluster_id)
+
+
+def bool_env(value: bool):
+    """
+    Takes a boolean value(or None) and return the serialized version to be used as an environemnt variable
+
+    Examples:
+        >>> bool_env(True)
+        "true"
+
+        >>> bool_env(False)
+        "false"
+
+        >>> bool_env(None)
+        "false"
+    """
+
+    if value is True:
+        return "true"
+    else:
+        return "false"

tests/utils/test_helpers.py

@@ -0,0 +1,7 @@
+from aztk.utils import helpers
+
+def test_bool_env():
+    assert helpers.bool_env(True) == "true"
+    assert helpers.bool_env(False) == "false"
+    assert helpers.bool_env(None) == "false"
+    assert helpers.bool_env("some") == "false"