Skip to content

Commit

Permalink
option for stored data
Browse files Browse the repository at this point in the history
  • Loading branch information
Fabian Peschel committed Dec 10, 2024
1 parent 41bf0b2 commit 77f9a89
Show file tree
Hide file tree
Showing 3 changed files with 45 additions and 3 deletions.
24 changes: 22 additions & 2 deletions src/jobflow_remote/cli/formatting.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,14 +26,20 @@
from jobflow_remote.jobs.upgrade import UpgradeAction


def get_job_info_table(jobs_info: list[JobInfo], verbosity: int) -> Table:
def get_job_info_table(
jobs_info: list[JobInfo],
verbosity: int,
stored_data_keys: list[str] | None = None,
skip_job_id: bool = False,
) -> Table:
time_zone_str = f" [{time.tzname[0]}]"

table = Table(title="Jobs info")
table.add_column("DB id")
table.add_column("Name")
table.add_column("State")
table.add_column("Job id (Index)")
if not skip_job_id:
table.add_column("Job id (Index)")

table.add_column("Worker")
table.add_column("Last updated" + time_zone_str)
Expand All @@ -50,6 +56,10 @@ def get_job_info_table(jobs_info: list[JobInfo], verbosity: int) -> Table:
table.add_column("Lock id")
table.add_column("Lock time" + time_zone_str)

if stored_data_keys:
for key in stored_data_keys:
table.add_column(f"{key}")

for ji in jobs_info:
state = ji.state.name

Expand All @@ -66,6 +76,8 @@ def get_job_info_table(jobs_info: list[JobInfo], verbosity: int) -> Table:
ji.worker,
convert_utc_time(ji.updated_on).strftime(fmt_datetime),
]
if skip_job_id:
row.pop(3)

if verbosity >= 1:
row.append(ji.remote.process_id)
Expand Down Expand Up @@ -98,6 +110,14 @@ def get_job_info_table(jobs_info: list[JobInfo], verbosity: int) -> Table:
else None
)

if stored_data_keys:
default_value = Text.from_markup("[italic grey58]none[/]")
stored_data = ji.stored_data
row += [
stored_data.get(key, default_value) if stored_data else default_value
for key in stored_data_keys
]

table.add_row(*row)

return table
Expand Down
23 changes: 22 additions & 1 deletion src/jobflow_remote/cli/job.py
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,22 @@ def jobs_list(
help="Select the jobs in FAILED and REMOTE_ERROR state. Incompatible with the --state option",
),
] = False,
stored_data_keys: Annotated[
Optional[list[str]],
typer.Option(
"--stored-data-key",
"-sdk",
help="Key to be shown from the stored_data field.",
),
] = None,
skip_job_id: Annotated[
bool,
typer.Option(
"--skip-job-id",
"-sji",
help="Skip the UUID field in the output table.",
),
] = False,
):
"""
Get the list of Jobs in the database
Expand Down Expand Up @@ -163,7 +179,12 @@ def jobs_list(
sort=db_sort,
)

table = get_job_info_table(jobs_info, verbosity=verbosity)
table = get_job_info_table(
jobs_info,
verbosity=verbosity,
stored_data_keys=stored_data_keys,
skip_job_id=skip_job_id,
)

out_console.print(table)
if SETTINGS.cli_suggestions:
Expand Down
1 change: 1 addition & 0 deletions src/jobflow_remote/jobs/data.py
Original file line number Diff line number Diff line change
Expand Up @@ -147,6 +147,7 @@ class JobInfo(BaseModel):
end_time: Optional[datetime] = None
priority: int = 0
metadata: Optional[dict] = None
stored_data: Optional[dict] = None

@property
def is_locked(self) -> bool:
Expand Down

0 comments on commit 77f9a89

Please sign in to comment.