ClI update n' bug fix (#42)

* Added support for updating project external IDs
* Fixed cluster discovery bug
* Version bump
* Ignore tasks without a cluster

Author: rmoneys

sync/__init__.py

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

sync/_databricks.py

@@ -29,8 +29,8 @@ def create_prediction(
     compute_type: str,
     cluster: dict,
     cluster_events: dict,
-    instances: dict,
     eventlog: bytes,
+    instances: dict = None,
     volumes: dict = None,
     tasks: List[dict] = None,
     project_id: str = None,
@@ -49,19 +49,20 @@ def create_prediction(
         If the cluster is a long-running cluster, this should only include events relevant to the time window that a
         run occurred in.
     :type cluster_events: dict
+    :param eventlog: encoded event log zip
+    :type eventlog: bytes
     :param instances: All EC2 Instances that were a part of the cluster. Expects a data format as is returned by
         `boto3's EC2.describe_instances API <https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ec2/client/describe_instances.html>`_
         Instances should be narrowed to just those instances relevant to the Databricks Run. This can be done by passing
         a `tag:ClusterId` filter to the describe_instances call like -
         ``Filters=[{"Name": "tag:ClusterId", "Values": ["my-dbx-clusterid"]}]``
         If there are multiple pages of instances, all pages should be accumulated into 1 dictionary and passed to this
         function
-    :param eventlog: encoded event log zip
-    :type eventlog: bytes
-    :param tasks: The Databricks Tasks associated with the cluster
-    :type tasks: List[dict]
+    :type instances: dict, optional
     :param volumes: The EBS volumes that were attached to this cluster
     :type volumes: dict, optional
+    :param tasks: The Databricks Tasks associated with the cluster
+    :type tasks: List[dict]
     :param project_id: Sync project ID, defaults to None
     :type project_id: str, optional
     :return: prediction ID
@@ -1032,7 +1033,9 @@ def _get_cluster_id_and_tasks_from_run_tasks(
     project_cluster_ids = defaultdict(list)
     all_cluster_tasks = defaultdict(list)
     for task in run["tasks"]:
-        if not exclude_tasks or task["task_key"] not in exclude_tasks:
+        if "cluster_instance" in task and (
+            not exclude_tasks or task["task_key"] not in exclude_tasks
+        ):
             cluster_id = task["cluster_instance"]["cluster_id"]
             all_cluster_tasks[cluster_id].append(task)
 
@@ -1051,7 +1054,7 @@ def _get_cluster_id_and_tasks_from_run_tasks(
         if cluster_ids:
             project_cluster_tasks = {
                 cluster_id: tasks
-                for cluster_id, tasks in all_cluster_tasks
+                for cluster_id, tasks in all_cluster_tasks.items()
                 if cluster_id in cluster_ids
             }
         else:
@@ -1062,15 +1065,17 @@ def _get_cluster_id_and_tasks_from_run_tasks(
     cluster_tasks = project_cluster_tasks or all_cluster_tasks
     num_clusters = len(cluster_tasks)
     if num_clusters == 0:
-        raise Exception("No cluster found for tasks")
+        raise RuntimeError("No cluster found for tasks")
     elif num_clusters > 1:
-        raise Exception("More than 1 cluster found for tasks")
+        raise RuntimeError("More than 1 cluster found for tasks")
 
     return cluster_tasks.popitem()
 
 
 def _get_run_spark_context_id(tasks: List[dict]) -> Response[str]:
-    context_ids = {task["cluster_instance"]["spark_context_id"] for task in tasks}
+    context_ids = {
+        task["cluster_instance"]["spark_context_id"] for task in tasks if "cluster_instance" in task
+    }
     num_ids = len(context_ids)
 
     if num_ids == 1:

sync/api/projects.py

@@ -102,6 +102,7 @@ def update_project(
     project_id: str,
     description: str = None,
     s3_url: str = None,
+    app_id: str = None,
     prediction_preference: Preference = None,
     prediction_params: dict = None,
 ) -> Response[dict]:
@@ -113,6 +114,8 @@ def update_project(
     :type description: str, optional
     :param s3_url: location of project event logs and configurations, defaults to None
     :type s3_url: str, optional
+    :param app_id: external identifier, defaults to None
+    :type app_id: str, optional
     :param prediction_preference: default preference for predictions, defaults to None
     :type prediction_preference: Preference, optional
     :param prediction_params: dictionary of prediction parameters, defaults to None. Valid options are documented `here <https://developers.synccomputing.com/reference/update_project_v1_projects__project_id__put>`__
@@ -125,6 +128,8 @@ def update_project(
         project_update["description"] = description
     if s3_url:
         project_update["s3_url"] = s3_url
+    if app_id:
+        project_update["app_id"] = app_id
     if prediction_preference:
         project_update["prediction_preference"] = prediction_preference
     if prediction_params:

sync/cli/projects.py

@@ -97,17 +97,24 @@ def create(
 @click.argument("project-id")
 @click.option("-d", "--description")
 @click.option("-l", "--location", help="S3 URL under which to store event logs and configuration")
+@click.option(
+    "-i", "--app-id", help="External identifier often based on the project's target application"
+)
 @click.option(
     "-p",
     "--preference",
     type=click.Choice(Preference),
     default=CONFIG.default_prediction_preference,
 )
 def update(
-    project_id: str, description: str = None, location: str = None, preference: Preference = None
+    project_id: str,
+    description: str = None,
+    location: str = None,
+    app_id: str = None,
+    preference: Preference = None,
 ):
     """Update a project"""
-    response = update_project(project_id, description, location, preference)
+    response = update_project(project_id, description, location, app_id, preference)
     if response.result:
         click.echo("Project updated")
     else:

sync/clients/sync.py

@@ -50,7 +50,7 @@ def update_access_token(self, response: httpx.Response):
 
 
 class SyncClient(RetryableHTTPClient):
-    def __init__(self, api_url, api_key):
+    def __init__(self, api_url: str, api_key: APIKey):
         super().__init__(
             client=httpx.Client(
                 base_url=api_url,
@@ -138,7 +138,7 @@ def _send(self, request: httpx.Request) -> dict:
 
 
 class ASyncClient(RetryableHTTPClient):
-    def __init__(self, api_url, api_key):
+    def __init__(self, api_url: str, api_key: APIKey):
         super().__init__(
             client=httpx.AsyncClient(
                 base_url=api_url,