Add get-all command to get all job info in one go

Author: emlyn

cli/spark/endpoints/job/get_all.py

@@ -0,0 +1,21 @@
+import argparse
+import typing
+import time
+import aztk.spark
+from cli import config
+from cli import utils
+
+def setup_parser(parser: argparse.ArgumentParser):
+    parser.add_argument('--id',
+                        dest='job_id',
+                        required=True,
+                        help='The unique id of your AZTK job')
+    parser.add_argument('--logs',
+                        action='store_true',
+                        help='Also print logs for completed applications')
+
+
+def execute(args: typing.NamedTuple):
+    spark_client = aztk.spark.Client(config.load_aztk_screts())
+
+    utils.print_job_verbose(spark_client, spark_client.get_job(args.job_id), args.logs)

cli/spark/endpoints/job/job.py

@@ -4,6 +4,7 @@
 from . import get_app_logs
 from . import get_app
 from . import get
+from . import get_all
 from . import list
 from . import list_apps
 from . import stop_app
@@ -16,6 +17,7 @@ class ClusterAction:
     get_app = "get-app"
     delete = "delete"
     get = "get"
+    get_all = "get-all"
     list = "list"
     list_apps = "list-apps"
     stop_app = "stop-app"
@@ -36,6 +38,8 @@ def setup_parser(parser: argparse.ArgumentParser):
         ClusterAction.delete, help="Delete a Job")
     get_parser = subparsers.add_parser(
         ClusterAction.get, help="Get information about a Job")
+    get_all_parser = subparsers.add_parser(
+        ClusterAction.get_all, help="Get all information about a Job including applications")
     list_parser = subparsers.add_parser(
         ClusterAction.list, help="List Jobs in your account")
     list_apps_parser = subparsers.add_parser(
@@ -52,6 +56,7 @@ def setup_parser(parser: argparse.ArgumentParser):
     get_app.setup_parser(get_app_parser)
     delete.setup_parser(delete_parser)
     get.setup_parser(get_parser)
+    get_all.setup_parser(get_all_parser)
     list.setup_parser(list_parser)
     list_apps.setup_parser(list_apps_parser)
     stop_app.setup_parser(stop_app_parser)
@@ -66,6 +71,7 @@ def execute(args: typing.NamedTuple):
     actions[ClusterAction.get_app] = get_app.execute
     actions[ClusterAction.delete] = delete.execute
     actions[ClusterAction.get] = get.execute
+    actions[ClusterAction.get_all] = get_all.execute
     actions[ClusterAction.list] = list.execute
     actions[ClusterAction.list_apps] = list_apps.execute
     actions[ClusterAction.stop_app] = stop_app.execute

cli/utils.py

@@ -4,7 +4,7 @@
 import threading
 import time
 from subprocess import call
-from typing import List
+from typing import Dict, List
 import azure.batch.models as batch_models
 import aztk.spark
 from aztk import error
@@ -247,6 +247,13 @@ def print_jobs(jobs: List[aztk.spark.models.Job]):
         )
 
 
+def get_applications(client, job: aztk.spark.models.Job):
+    if job.applications:
+        return {j.name: j for j in job.applications}
+    else:
+        return client.list_applications(job.id)
+
+
 def print_job(client, job: aztk.spark.models.Job):
     print_format = '{:<36}| {:<15}'
 
@@ -267,11 +274,41 @@ def print_job(client, job: aztk.spark.models.Job):
             log.info(print_format.format("Cluster", "Provisioning"))
             log.info("")
 
-    if job.applications:
-        application_summary(job.applications)
+    application_summary(get_applications(client, job))
+    log.info("")
+
+
+def print_job_verbose(client, job: aztk.spark.models.Job, print_logs: bool):
+    print_format = '{:<36}| {:<15}'
+
+    log.info("")
+    log.info("Job               %s", job.id)
+    log.info("------------------------------------------")
+    log.info("State:            %s", job.state)
+    log.info("Transition Time:  %s", utc_to_local(job.state_transition_time))
+    log.info("")
+
+    if job.cluster:
+        print_cluster_summary(job.cluster)
     else:
-        application_summary(client.list_applications(job.id))
+        if job.state == 'completed':
+            log.info("Cluster           %s", "Job completed, cluster deallocated.")
+            log.info("")
+        else:
+            log.info(print_format.format("Cluster", "Provisioning"))
+            log.info("")
+
+    applications = get_applications(client, job)
+    application_summary(applications)
     log.info("")
+    print_applications(applications)
+    log.info("")
+
+    if print_logs:
+        for name, app in applications.items():
+            if app and app.exit_code is not None:
+                log.info("Logs for %s:", name)
+                print(client.get_job_application_log(job.id, name).log)
 
 
 def node_state_count(cluster: aztk.spark.models.Cluster):
@@ -296,22 +333,22 @@ def print_cluster_summary(cluster: aztk.spark.models.Cluster):
     log.info("")
 
 
-def application_summary(applications):
+def application_summary(applications: Dict[str, aztk.spark.models.Application]):
     states = {"scheduling": 0}
     for state in batch_models.TaskState:
         states[state.name] = 0
 
     warn_scheduling = False
 
-    for application in applications:
-        if type(application) == str:
+    for name, application in applications.items():
+        if application is None:
             states["scheduling"] += 1
             warn_scheduling = True
         else:
             states[application.state] += 1
 
     print_format = '{:<17} {:<14}'
-    log.info("Applications")
+    log.info("Application States")
     log.info("-"*42)
     for state in states:
         if states[state] > 0:
@@ -320,15 +357,15 @@ def application_summary(applications):
     if warn_scheduling:
         log.warning("\nNo Spark applications will be scheduled until the master is selected.")
 
-def print_applications(applications):
+def print_applications(applications: Dict[str, aztk.spark.models.Application]):
     print_format = '{:<36}| {:<15}| {:<16} | {:^9} |'
     print_format_underline = '{:-<36}|{:-<16}|{:-<18}|{:-<11}|'
     log.info(print_format.format("Applications", "State", "Transition Time", "Exit Code"))
     log.info(print_format_underline.format('', '', '', ''))
 
     warn_scheduling = False
-    for name in applications:
-        if applications[name] is None:
+    for name, application in applications.items():
+        if application is None:
             log.info(
                 print_format.format(
                     name,
@@ -339,10 +376,9 @@ def print_applications(applications):
             )
             warn_scheduling = True
         else:
-            application = applications[name]
             log.info(
                 print_format.format(
-                    application.name,
+                    name,
                     application.state,
                     utc_to_local(application.state_transition_time),
                     application.exit_code if application.exit_code is not None else "-"