Make job ID optional for polling and improve usability (#203)

* Make job ID optional for polling and improve usability

* Modified list_jobs and updated polling logic

* modified list_jobs and poll_job

* minor modifications on cli.py and run.py

* modify test cli file

* remove mocked JobManager

---------

Co-authored-by: Alessandro Cosentino <[email protected]>

metriq_gym/cli.py

@@ -4,31 +4,31 @@
 
 from tabulate import tabulate
 
-from metriq_gym.job_manager import JobManager, MetriqGymJob
+from metriq_gym.job_manager import MetriqGymJob
 from metriq_gym.provider import ProviderType
 
 LIST_JOBS_HEADERS = ["ID", "Provider", "Device", "Type", "Dispatch time (UTC)"]
 
-
-def list_jobs(job_manager: JobManager) -> None:
+def list_jobs(jobs: list[MetriqGymJob], show_index: bool = True) -> None:
     """List jobs recorded in the job manager.
 
     Args:
-        job_manager: Job manager instance.
+        jobs: List of MetriqGymJob instances.
+        show_index: Whether to show the job index in the output table.
     """
-    # Retrieve all jobs from JobManager.
-    jobs: list[MetriqGymJob] = job_manager.get_jobs()
-
-    # Display jobs in a tabular format.
     if not jobs:
         print("No jobs found.")
         return
-
+    # Display jobs in a tabular format.
     print(
-        tabulate([job.to_table_row() for job in jobs], headers=LIST_JOBS_HEADERS, tablefmt="grid")
+        tabulate(
+            [job.to_table_row() for job in jobs],
+            headers=LIST_JOBS_HEADERS,
+            tablefmt="grid",
+            showindex=show_index,
+        )
     )
 
-
 def parse_arguments() -> argparse.Namespace:
     """
     Parse command-line arguments for the quantum volume benchmark.
@@ -64,7 +64,7 @@ def parse_arguments() -> argparse.Namespace:
 
     # Subparser for poll.
     poll_parser = subparsers.add_parser("poll", help="Poll jobs")
-    poll_parser.add_argument("--job_id", type=str, required=True, help="Job ID to poll")
+    poll_parser.add_argument("--job_id", type=str, required=False, help="Job ID to poll (optional)")
 
     # Subparser for list-jobs.
     subparsers.add_parser("list-jobs", help="List dispatched jobs")

metriq_gym/run.py

@@ -58,6 +58,24 @@ def dispatch_job(args: argparse.Namespace, job_manager: JobManager) -> None:
 
 
 def poll_job(args: argparse.Namespace, job_manager: JobManager) -> None:
+    if not args.job_id:
+        jobs = job_manager.get_jobs()
+        if not jobs:
+            print("No jobs available for polling.")
+            return
+        print("Available jobs:")
+        list_jobs(jobs, show_index=True)
+        while True:
+            try:
+                selected_index = int(input("Select a job index: "))
+                if 0 <= selected_index < len(jobs):
+                    break
+                else:
+                    print(f"Invalid index. Please enter a number between 0 and {len(jobs) - 1}")
+            except ValueError:
+                print("Invalid input. Please enter a valid number.")
+        args.job_id = jobs[selected_index].id
+
     logger.info("Polling job...")
     metriq_job: MetriqGymJob = job_manager.get_job(args.job_id)
     job_type: JobType = JobType(metriq_job.job_type)
@@ -85,7 +103,8 @@ def main() -> int:
     elif args.action == "poll":
         poll_job(args, job_manager)
     elif args.action == "list-jobs":
-        list_jobs(job_manager)
+        jobs: list[MetriqGymJob] = job_manager.get_jobs()
+        list_jobs(jobs)
 
     else:
         logging.error("Invalid action specified. Run with --help for usage information.")

tests/test_cli.py

@@ -12,11 +12,10 @@ def mock_job_manager():
     """Fixture to provide a mocked JobManager."""
     return MagicMock(spec=JobManager)
 
-
-def test_list_jobs_all(mock_job_manager, capsys):
+def test_list_jobs_all(capsys):
     """Test listing all jobs without filters."""
     # Mock jobs
-    mock_job_manager.get_jobs.return_value = [
+    mock_jobs = [
         MetriqGymJob(
             id="1234",
             device_name="ibmq_qasm_simulator",
@@ -36,8 +35,8 @@ def test_list_jobs_all(mock_job_manager, capsys):
             data={},
         ),
     ]
-
-    list_jobs(mock_job_manager)
+    
+    list_jobs(mock_jobs, show_index=False)
 
     # Capture the output
     captured = capsys.readouterr()
@@ -52,12 +51,12 @@ def test_list_jobs_all(mock_job_manager, capsys):
     assert captured.out == expected_output
 
 
-def test_list_jobs_no_jobs(mock_job_manager, capsys):
+def test_list_jobs_no_jobs(capsys):
     """Test listing jobs when no jobs are recorded."""
     # Mock no jobs
-    mock_job_manager.get_jobs.return_value = []
-
-    list_jobs(mock_job_manager)
+    mock_jobs = []
+    
+    list_jobs(mock_jobs, show_index=False)
 
     # Capture the output
     captured = capsys.readouterr()