feat: Implement method to get number of pending jobs.

PiperOrigin-RevId: 716684266

Author: vertex-sdk-bot

tests/unit/vertexai/test_batch_prediction.py

@@ -723,3 +723,9 @@ def tes_list_batch_prediction_jobs(self, list_batch_prediction_jobs_mock):
         list_batch_prediction_jobs_mock.assert_called_once_with(
             request={"parent": _TEST_PARENT}
         )
+
+    def test_num_pending_jobs(self, list_batch_prediction_jobs_mock):
+        num_pending_jobs = batch_prediction.BatchPredictionJob.num_pending_jobs()
+
+        assert num_pending_jobs == 1
+        list_batch_prediction_jobs_mock.assert_called_once()

vertexai/batch_prediction/_batch_prediction.py

@@ -328,6 +328,28 @@ def _is_genai_model(cls, model_name: str) -> bool:
 
         return False
 
+    @classmethod
+    def num_pending_jobs(cls) -> int:
+        """Returns the number of pending batch prediction jobs.
+
+        The pending jobs are those defined in _JOB_PENDING_STATES from
+        google/cloud/aiplatform/jobs.py
+        e.g. JOB_STATE_QUEUED, JOB_STATE_PENDING, JOB_STATE_RUNNING,
+        JOB_STATE_CANCELLING, JOB_STATE_UPDATING.
+        It will be used to manage the number of concurrent batch that is limited
+        according to
+        https://cloud.google.com/vertex-ai/generative-ai/docs/quotas#concurrent-batch-requests
+        """
+        return len(
+            cls._list(
+                cls_filter=lambda gca_resource: cls._is_genai_model(gca_resource.model),
+                filter=" OR ".join(
+                    f'state="{pending_state.name}"'
+                    for pending_state in jobs._JOB_PENDING_STATES
+                ),
+            )
+        )
+
     @classmethod
     def _complete_bq_uri(cls, uri: Optional[str] = None):
         """Completes a BigQuery uri to a BigQuery table uri."""