Improve docs on Batch tasks

src/bloqade/submission/ir/task_results.py

@@ -16,6 +16,23 @@ class QuEraShotStatusCode(str, Enum):
 
 
 class QuEraTaskStatusCode(str, Enum):
+    """
+    An Enum representing the various states a task can be in within the QuEra system.
+
+    Attributes:
+    - Created: The task has been created but not yet started.
+    - Running: The task is currently running.
+    - Completed: The task has completed successfully.
+    - Failed: The task has failed.
+    - Cancelled: The task has been cancelled.
+    - Executing: The task is currently being executed.
+    - Enqueued: The task is in the queue waiting to be executed.
+    - Accepted: The task has been accepted for execution.
+    - Unaccepted: The task has not been accepted for execution.
+    - Partial: The task has partially completed.
+    - Unsubmitted: The task has not been submitted for execution.
+    """
+
     Created = "Created"
     Running = "Running"
     Completed = "Completed"

src/bloqade/task/batch.py

@@ -1,5 +1,6 @@
 from decimal import Decimal
 from numbers import Real
+from typing import Literal
 from bloqade.builder.typing import LiteralType
 from bloqade.serialize import Serializer
 from bloqade.task.base import Report
@@ -216,6 +217,22 @@ def rerun(
     def _run(
         self, multiprocessing: bool = False, num_workers: Optional[int] = None, **kwargs
     ):
+        """
+        Private method to run tasks in the batch.
+
+        Args:
+            multiprocessing (bool, optional): If True, tasks are run in parallel using multiple processes.
+                If False, tasks are run sequentially in a single process. Defaults to False.
+            num_workers (Optional[int], optional): The maximum number of processes that can be used to
+                execute the given calls if multiprocessing is True. If None, the number of workers will be the number of processors on the machine.
+            **kwargs: Arbitrary keyword arguments passed to the task's run method.
+
+        Raises:
+            ValueError: If num_workers is not None and multiprocessing is False.
+
+        Returns:
+            self: The instance of the batch with tasks run.
+        """
         if multiprocessing:
             from concurrent.futures import ProcessPoolExecutor as Pool
 
@@ -468,6 +485,19 @@ def resubmit(self, shuffle_submit_order: bool = True) -> "RemoteBatch":
     def _submit(
         self, shuffle_submit_order: bool = True, ignore_submission_error=False, **kwargs
     ) -> "RemoteBatch":
+        """
+        Private method to submit tasks in the RemoteBatch.
+
+        Args:
+            shuffle_submit_order (bool, optional): If True, tasks are submitted in a random order.
+                If False, tasks are submitted in the order they were added to the batch. Defaults to True.
+            ignore_submission_error (bool, optional): If True, submission errors are ignored and the method continues to submit the remaining tasks.
+                If False, the method stops at the first submission error. Defaults to False.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            RemoteBatch: The RemoteBatch instance with tasks submitted.
+        """
         from bloqade import save
 
         # online, non-blocking
@@ -565,7 +595,22 @@ def get_tasks(self, *status_codes: str) -> "RemoteBatch":
         return RemoteBatch(self.source, new_task_results, name=self.name)
 
     @beartype
-    def remove_tasks(self, *status_codes: str) -> "RemoteBatch":
+    def remove_tasks(
+        self,
+        *status_codes: Literal[
+            "Created",
+            "Running",
+            "Completed",
+            "Failed",
+            "Cancelled",
+            "Executing",
+            "Enqueued",
+            "Accepted",
+            "Unaccepted",
+            "Partial",
+            "Unsubmitted",
+        ],
+    ) -> "RemoteBatch":
         """
         Remove Tasks with specify status_codes.