Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix for --format provider-json when run with google-batch provider #297

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions dsub/commands/dstat.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand Down
3 changes: 2 additions & 1 deletion dsub/providers/google_batch.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down