Prod-1746-Hosted-Monitor-Idempotency (#115)

* add lambda option to monitor

* Tests =

* bump version

* lint

* reduce complexity

Author: Brandon Kaplan

sync/__init__.py

@@ -1,5 +1,5 @@
 """Library for leveraging the power of Sync"""
 
-__version__ = "1.6.0"
+__version__ = "1.7.0"
 
 TIME_FORMAT = "%Y-%m-%dT%H:%M:%SZ"

sync/awsdatabricks.py

@@ -12,7 +12,6 @@
 import sync._databricks
 from sync._databricks import (
     _cluster_log_destination,
-    get_all_cluster_events,
     _get_cluster_instances_from_dbfs,
     _update_monitored_timelines,
     _wait_for_cluster_termination,
@@ -23,6 +22,7 @@
     create_run,
     create_submission_for_run,
     create_submission_with_cluster_info,
+    get_all_cluster_events,
     get_cluster,
     get_cluster_report,
     get_project_cluster,
@@ -48,9 +48,9 @@
     AccessReportLine,
     AccessStatusCode,
     AWSDatabricksClusterReport,
+    DatabricksComputeType,
     DatabricksError,
     DatabricksPlanType,
-    DatabricksComputeType,
     Response,
 )
 from sync.utils.dbfs import format_dbfs_filepath, write_dbfs_file
@@ -367,9 +367,11 @@ def monitor_cluster(
         spark_context_id = cluster.get("spark_context_id")
 
     (log_url, filesystem, bucket, base_prefix) = _cluster_log_destination(cluster)
+    write_function = None
     if cluster_report_destination_override:
         filesystem = cluster_report_destination_override.get("filesystem", filesystem)
         base_prefix = cluster_report_destination_override.get("base_prefix", base_prefix)
+        write_function = cluster_report_destination_override.get("write_function")
 
     if log_url or cluster_report_destination_override:
         _monitor_cluster(
@@ -378,6 +380,7 @@ def monitor_cluster(
             spark_context_id,
             polling_period,
             kill_on_termination,
+            write_function,
         )
     else:
         logger.warning("Unable to monitor cluster due to missing cluster log destination - exiting")
@@ -389,6 +392,7 @@ def _monitor_cluster(
     spark_context_id: int,
     polling_period: int,
     kill_on_termination: bool = False,
+    write_function=None,
 ) -> None:
 
     (log_url, filesystem, bucket, base_prefix) = cluster_log_destination
@@ -400,7 +404,7 @@ def _monitor_cluster(
     aws_region_name = DB_CONFIG.aws_region_name
     ec2 = boto.client("ec2", region_name=aws_region_name)
 
-    write_file = _define_write_file(file_key, filesystem, bucket)
+    write_file = _define_write_file(file_key, filesystem, bucket, write_function)
 
     all_inst_by_id = {}
     active_timelines_by_id = {}
@@ -454,8 +458,14 @@ def _monitor_cluster(
         sleep(polling_period)
 
 
-def _define_write_file(file_key, filesystem, bucket):
-    if filesystem == "file":
+def _define_write_file(file_key, filesystem, bucket, write_function):
+    if filesystem == "lambda":
+
+        def write_file(body: bytes):
+            logger.info("Using custom lambda function to write data")
+            write_function(body)
+
+    elif filesystem == "file":
         file_path = Path(file_key)
 
         def ensure_path_exists(report_path: Path):

sync/azuredatabricks.py

@@ -9,14 +9,13 @@
 
 from azure.common.credentials import get_cli_profile
 from azure.core.exceptions import ClientAuthenticationError
-from azure.identity import DefaultAzureCredential, ClientSecretCredential
+from azure.identity import ClientSecretCredential, DefaultAzureCredential
 from azure.mgmt.compute import ComputeManagementClient
 from azure.mgmt.resource import ResourceManagementClient
 
 import sync._databricks
 from sync._databricks import (
     _cluster_log_destination,
-    get_all_cluster_events,
     _get_cluster_instances_from_dbfs,
     _update_monitored_timelines,
     _wait_for_cluster_termination,
@@ -27,6 +26,7 @@
     create_run,
     create_submission_for_run,
     create_submission_with_cluster_info,
+    get_all_cluster_events,
     get_cluster,
     get_cluster_report,
     get_project_cluster,
@@ -51,9 +51,9 @@
     AccessReportLine,
     AccessStatusCode,
     AzureDatabricksClusterReport,
+    DatabricksComputeType,
     DatabricksError,
     DatabricksPlanType,
-    DatabricksComputeType,
     Response,
 )
 from sync.utils.dbfs import format_dbfs_filepath, write_dbfs_file
@@ -338,9 +338,11 @@ def monitor_cluster(
         spark_context_id = cluster.get("spark_context_id")
 
     (log_url, filesystem, bucket, base_prefix) = _cluster_log_destination(cluster)
+    write_function = None
     if cluster_report_destination_override:
         filesystem = cluster_report_destination_override.get("filesystem", filesystem)
         base_prefix = cluster_report_destination_override.get("base_prefix", base_prefix)
+        write_function = cluster_report_destination_override.get("write_function")
 
     if log_url:
         _monitor_cluster(
@@ -349,6 +351,7 @@ def monitor_cluster(
             spark_context_id,
             polling_period,
             kill_on_termination,
+            write_function,
         )
     else:
         logger.warning("Unable to monitor cluster due to missing cluster log destination - exiting")
@@ -360,6 +363,7 @@ def _monitor_cluster(
     spark_context_id: int,
     polling_period: int,
     kill_on_termination: bool = False,
+    write_function=None,
 ) -> None:
     (log_url, filesystem, bucket, base_prefix) = cluster_log_destination
     # If the event log destination is just a *bucket* without any sub-path, then we don't want to include
@@ -370,7 +374,7 @@ def _monitor_cluster(
     azure_logger = logging.getLogger("azure.core.pipeline.policies.http_logging_policy")
     azure_logger.setLevel(logging.WARNING)
 
-    write_file = _define_write_file(file_key, filesystem)
+    write_file = _define_write_file(file_key, filesystem, write_function)
 
     resource_group_name = _get_databricks_resource_group_name()
     if not resource_group_name:
@@ -418,8 +422,14 @@ def _monitor_cluster(
         sleep(polling_period)
 
 
-def _define_write_file(file_key, filesystem):
-    if filesystem == "file":
+def _define_write_file(file_key, filesystem, write_function):
+    if filesystem == "lambda":
+
+        def write_file(body: bytes):
+            logger.info("Using custom lambda function to write data")
+            write_function(body)
+
+    elif filesystem == "file":
         file_path = Path(file_key)
 
         def ensure_path_exists(report_path: Path):

tests/test_awsdatabricks.py

@@ -39,14 +39,24 @@ def test_monitor_cluster_with_override(
 
         expected_log_destination_override = ("s3://bucket/path", "file", "bucket", "test_file_path")
         mock_monitor_cluster.assert_called_with(
-            expected_log_destination_override, "0101-214342-tpi6qdp2", 1443449481634833945, 1, True
+            expected_log_destination_override,
+            "0101-214342-tpi6qdp2",
+            1443449481634833945,
+            1,
+            True,
+            None,
         )
 
         mock_cluster_log_destination.return_value = (None, "s3", None, "path")
         monitor_cluster("0101-214342-tpi6qdp2", 1, cluster_report_destination_override, True)
         expected_log_destination_override = (None, "file", None, "test_file_path")
         mock_monitor_cluster.assert_called_with(
-            expected_log_destination_override, "0101-214342-tpi6qdp2", 1443449481634833945, 1, True
+            expected_log_destination_override,
+            "0101-214342-tpi6qdp2",
+            1443449481634833945,
+            1,
+            True,
+            None,
         )
 
     def test_monitor_cluster_without_override(
@@ -70,4 +80,5 @@ def test_monitor_cluster_without_override(
             1443449481634833945,
             1,
             False,
+            None,
         )