From 2fadb9662868a3325892660938f0fa2857cb5622 Mon Sep 17 00:00:00 2001 From: Philip Montgomery Date: Thu, 29 Aug 2024 22:30:54 -0400 Subject: [PATCH] Fix for --format provider-json --- dsub/commands/dstat.py | 8 ++++++-- dsub/providers/google_batch.py | 3 ++- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/dsub/commands/dstat.py b/dsub/commands/dstat.py index 642c0a6..eecdf2e 100755 --- a/dsub/commands/dstat.py +++ b/dsub/commands/dstat.py @@ -201,6 +201,7 @@ def main(): # to provide a username automatically. user_ids = set(args.users) if args.users else {dsub_util.get_os_user()} labels = param_util.parse_pair_args(args.label, job_model.LabelParam) + raw_format = bool(args.format == 'provider-json') job_producer = dstat_job_producer( provider=provider, @@ -216,12 +217,15 @@ def main(): full_output=args.full, summary_output=args.summary, poll_interval=poll_interval, - raw_format=bool(args.format == 'provider-json')) + raw_format=raw_format) # Track if any jobs are running in the event --wait was requested. for poll_event_tasks in job_producer: rows = poll_event_tasks - formatter.prepare_and_print_table(rows, args.summary) + if raw_format: + formatter.print_table(rows) + else: + formatter.prepare_and_print_table(rows, args.summary) def dstat_job_producer(provider, diff --git a/dsub/providers/google_batch.py b/dsub/providers/google_batch.py index 654539b..e24e477 100644 --- a/dsub/providers/google_batch.py +++ b/dsub/providers/google_batch.py @@ -32,6 +32,7 @@ from . import google_batch_operations from . import google_custom_machine from . import google_utils +from proto import message # pylint: disable=g-import-not-at-top try: @@ -207,7 +208,7 @@ def __init__(self, operation_data: batch_v1.types.Job): self._job_descriptor = self._try_op_to_job_descriptor() def raw_task_data(self): - return self._op + return message.Message.to_dict(self._op) def _try_op_to_job_descriptor(self): # The _META_YAML_REPR field in the 'prepare' action enables reconstructing