diff --git a/src/jobflow_remote/cli/formatting.py b/src/jobflow_remote/cli/formatting.py index b8121654..34219a20 100644 --- a/src/jobflow_remote/cli/formatting.py +++ b/src/jobflow_remote/cli/formatting.py @@ -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) @@ -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 @@ -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) @@ -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 diff --git a/src/jobflow_remote/cli/job.py b/src/jobflow_remote/cli/job.py index 6a9cec18..9e13d486 100644 --- a/src/jobflow_remote/cli/job.py +++ b/src/jobflow_remote/cli/job.py @@ -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 @@ -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: diff --git a/src/jobflow_remote/jobs/data.py b/src/jobflow_remote/jobs/data.py index a2ef47a0..2fb54d4c 100644 --- a/src/jobflow_remote/jobs/data.py +++ b/src/jobflow_remote/jobs/data.py @@ -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: