8000 feat: Export global quota configs in preview sdk · googleapis/python-aiplatform@7f964d5 · GitHub
[go: up one dir, main page]
Skip to content

Commit 7f964d5

Browse files
vertex-sdk-botcopybara-github
vertex-sdk-bot
authored and
copybara-github
committed
feat: Export global quota configs in preview sdk
PiperOrigin-RevId: 770276336
1 parent da09275 commit 7f964d5

File tree

5 files changed

+81
-1
lines changed

5 files changed

+81
-1
lines changed

‎tests/unit/vertex_rag/test_rag_constants_preview.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -485,10 +485,12 @@
485485
TEST_IMPORT_FILES_CONFIG_GCS = ImportRagFilesConfig(
486486
rag_file_transformation_config=TEST_RAG_FILE_TRANSFORMATION_CONFIG,
487487
rebuild_ann_index=False,
488+
max_embedding_requests_per_min=1000,
488489
)
489490
TEST_IMPORT_FILES_CONFIG_GCS_REBUILD_ANN_INDEX = ImportRagFilesConfig(
490491
rag_file_transformation_config=TEST_RAG_FILE_TRANSFORMATION_CONFIG,
491492
rebuild_ann_index=True,
493+
max_embedding_requests_per_min=1000,
492494
)
493495
TEST_IMPORT_FILES_CONFIG_GCS_REBUILD_ANN_INDEX.gcs_source.uris = [TEST_GCS_PATH]
494496
TEST_IMPORT_FILES_CONFIG_GCS_REBUILD_ANN_INDEX.rag_file_parsing_config.advanced_parser.use_advanced_pdf_parsing = (
@@ -517,6 +519,7 @@
517519
TEST_IMPORT_FILES_CONFIG_DRIVE_FOLDER = ImportRagFilesConfig(
518520
rag_file_transformation_config=TEST_RAG_FILE_TRANSFORMATION_CONFIG,
519521
rebuild_ann_index=False,
522+
max_embedding_requests_per_min=1000,
520523
)
521524
TEST_IMPORT_FILES_CONFIG_DRIVE_FOLDER.google_drive_source.resource_ids = [
522525
GoogleDriveSource.ResourceId(
@@ -530,6 +533,7 @@
530533
TEST_IMPORT_FILES_CONFIG_DRIVE_FOLDER_PARSING = ImportRagFilesConfig(
531534
rag_file_transformation_config=TEST_RAG_FILE_TRANSFORMATION_CONFIG,
532535
rebuild_ann_index=False,
536+
max_embedding_requests_per_min=1000,
533537
)
534538
TEST_IMPORT_FIL 10000 ES_CONFIG_DRIVE_FOLDER_PARSING.google_drive_source.resource_ids = [
535539
GoogleDriveSource.ResourceId(
@@ -589,6 +593,7 @@
589593
)
590594
),
591595
rebuild_ann_index=False,
596+
max_embedding_requests_per_min=1000,
592597
)
593598
TEST_IMPORT_FILES_CONFIG_DRIVE_FILE.max_embedding_requests_per_min = 800
594599

@@ -603,6 +608,14 @@
603608
import_rag_files_config=TEST_IMPORT_FILES_CONFIG_DRIVE_FILE,
604609
)
605610

611+
TEST_IMPORT_REQUEST_DRIVE_FILE_GLOBAL_QUOTA_CONTROL = ImportRagFilesRequest(
612+
parent=TEST_RAG_CORPUS_RESOURCE_NAME,
613+
import_rag_files_config=TEST_IMPORT_FILES_CONFIG_DRIVE_FILE,
614+
)
615+
TEST_IMPORT_REQUEST_DRIVE_FILE_GLOBAL_QUOTA_CONTROL.import_rag_files_config.global_max_embedding_requests_per_min = (
616+
8000
617+
)
618+
606619
TEST_IMPORT_RESPONSE = ImportRagFilesResponse(imported_rag_files_count=2)
607620

608621
TEST_GAPIC_RAG_FILE = GapicRagFile(
@@ -649,6 +662,7 @@
649662
rag_file_parsing_config=TEST_RAG_FILE_PARSING_CONFIG,
650663
rag_file_transformation_config=TEST_RAG_FILE_TRANSFORMATION_CONFIG,
651664
rebuild_ann_index=False,
665+
max_embedding_requests_per_min=1000,
652666
)
653667
TEST_IMPORT_FILES_CONFIG_SLACK_SOURCE.slack_source.channels = [
654668
GapicSlackSource.SlackChannels(
@@ -703,6 +717,7 @@
703717
rag_file_parsing_config=TEST_RAG_FILE_PARSING_CONFIG,
704718
rag_file_transformation_config=TEST_RAG_FILE_TRANSFORMATION_CONFIG,
705719
rebuild_ann_index=False,
720+
max_embedding_requests_per_min=1000,
706721
)
707722
TEST_IMPORT_FILES_CONFIG_JIRA_SOURCE.jira_source.jira_queries = [
708723
GapicJiraSource.JiraQueries(
@@ -736,6 +751,7 @@
736751
TEST_IMPORT_FILES_CONFIG_SHARE_POINT_SOURCE = ImportRagFilesConfig(
737752
rag_file_parsing_config=TEST_RAG_FILE_PARSING_CONFIG,
738753
rag_file_transformation_config=TEST_RAG_FILE_TRANSFORMATION_CONFIG,
754+
max_embedding_requests_per_min=1000,
739755
share_point_sources=GapicSharePointSources(
740756
share_point_sources=[
741757
GapicSharePointSources.SharePointSource(
@@ -813,6 +829,7 @@
813829
TEST_LAYOUT_PARSER_WITH_PROCESSOR_PATH_CONFIG = LayoutParserConfig(
814830
processor_name="projects/test-project/locations/us/processors/abc123",
815831
max_parsing_requests_per_min=100,
832+
global_max_parsing_requests_per_min=1000,
816833
)
817834

818835
TEST_LAYOUT_PARSER_WITH_PROCESSOR_VERSION_PATH_CONFIG = LayoutParserConfig(
@@ -885,6 +902,7 @@
885902

886903
TEST_IMPORT_FILES_CONFIG_SHARE_POINT_SOURCE_NO_FOLDERS = ImportRagFilesConfig(
887904
rag_file_transformation_config=TEST_RAG_FILE_TRANSFORMATION_CONFIG,
905+
max_embedding_requests_per_min=1000,
888906
share_point_sources=GapicSharePointSources(
889907
share_point_sources=[
890908
GapicSharePointSources.SharePointSource(
@@ -914,6 +932,7 @@< 72A1 /div>
914932
layout_parser=RagFileParsingConfig.LayoutParser(
915933
processor_name="projects/test-project/locations/us/processors/abc123",
916934
max_parsing_requests_per_min=100,
935+
global_max_parsing_requests_per_min=1000,
917936
)
918937
)
919938
)

‎tests/unit/vertex_rag/test_rag_data_preview.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -659,6 +659,14 @@ def import_files_request_eq(returned_request, expected_request):
659659
returned_request.import_rag_files_config.rebuild_ann_index
660660
== expected_request.import_rag_files_config.rebuild_ann_index
661661
)
662+
assert (
663+
returned_request.import_rag_files_config.max_embedding_requests_per_min
664+
== expected_request.import_rag_files_config.max_embedding_requests_per_min
665+
)
666+
assert (
667+
returned_request.import_rag_files_config.global_max_embedding_requests_per_min
668+
== expected_request.import_rag_files_config.global_max_embedding_requests_per_min
669+
)
662670

663671

664672
def rag_engine_config_eq(returned_config, expected_config):
@@ -1349,6 +1357,20 @@ def test_prepare_import_files_request_drive_files(self):
13491357
request, test_rag_constants_preview.TEST_IMPORT_REQUEST_DRIVE_FILE
13501358
)
13511359

1360+
def test_prepare_import_files_request_drive_files_with_global_quota_control(self):
1361+
paths = [test_rag_constants_preview.TEST_DRIVE_FILE]
1362+
request = prepare_import_files_request(
1363+
corpus_name=test_rag_constants_preview.TEST_RAG_CORPUS_RESOURCE_NAME,
1364+
paths=paths,
1365+
transformation_config=create_transformation_config(),
1366+
max_embedding_requests_per_min=800,
1367+
global_max_embedding_requests_per_min=8000,
1368+
)
1369+
import_files_request_eq(
1370+
request,
1371+
test_rag_constants_preview.TEST_IMPORT_REQUEST_DRIVE_FILE_GLOBAL_QUOTA_CONTROL,
1372+
)
1373+
13521374
def test_prepare_import_files_request_invalid_drive_path(self):
13531375
with pytest.raises(ValueError) as F438 e:
13541376
paths = ["https://drive.google.com/bslalsdfk/whichever_file/456"]

‎vertexai/preview/rag/rag_data.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -496,6 +496,7 @@ def import_files(
496496
transformation_config: Optional[TransformationConfig] = None,
497497
timeout: int = 600,
498498
max_embedding_requests_per_min: int = 1000,
499+
global_max_embedding_requests_per_min: Optional[int] = None,
499500
use_advanced_pdf_parsing: Optional[bool] = False,
500501
partial_failures_sink: Optional[str] = None,
501502
layout_parser: Optional[LayoutParserConfig] = None,
@@ -605,6 +606,13 @@ def import_files(
605606
page on the project to set an appropriate value
606607
here. If unspecified, a default value of 1,000
607608
QPM would be used.
609+
global_max_embedding_requests_per_min:
610+
Optional. The max number of queries per minute that the indexing
611+
pipeline job is allowed to make to the embedding model specified in
612+
the project. Please follow the quota usage guideline of the embedding
613+
model you use to set the value properly. If this value is not specified,
614+
max_embedding_requests_per_min will be used by indexing pipeline job
615+
as the global limit and this means parallel import jobs are not allowed.
608616
timeout: Default is 600 seconds.
609617
use_advanced_pdf_parsing: Whether to use advanced PDF
610618
parsing on uploaded files. This field is deprecated.
@@ -663,6 +671,7 @@ def import_files(
663671
chunk_overlap=chunk_overlap,
664672
transformation_config=transformation_config,
665673
max_embedding_requests_per_min=max_embedding_requests_per_min,
674+
global_max_embedding_requests_per_min=global_max_embedding_requests_per_min,
666675
use_advanced_pdf_parsing=use_advanced_pdf_parsing,
667676
partial_failures_sink=partial_failures_sink,
668677
layout_parser=layout_parser,
@@ -686,6 +695,7 @@ async def import_files_async(
686695
chunk_overlap: int = 200,
687696
transformation_config: Optional[TransformationConfig] = None,
688697
max_embedding_requests_per_min: int = 1000,
698+
global_max_embedding_requests_per_min: Optional[int] = None,
689699
use_advanced_pdf_parsing: Optional[bool] = False,
690700
partial_failures_sink: Optional[str] = None,
691701
layout_parser: Optional[LayoutParserConfig] = None,
@@ -796,6 +806,13 @@ async def import_files_async(
796806
page on the project to set an appropriate value
797807
here. If unspecified, a default value of 1,000
798808
QPM would be used.
809+
global_max_embedding_requests_per_min:
810+
Optional. The max number of queries per minute that the indexing
811+
pipeline job is allowed to make to the embedding model specified in
812+
the project. Please follow the quota usage guideline of the embedding
813+
model you use to set the value properly. If this value is not specified,
814+
max_embedding_requests_per_min will be used by indexing pipeline job
815+
as the global limit and this means parallel import jobs are not allowed.
799816
use_advanced_pdf_parsing: Whether to use advanced PDF
800817
parsing on uploaded files.
801818
partial_failures_sink: Either a GCS path to store partial failures or a
@@ -852,6 +869,7 @@ async def import_files_async(
852869
chunk_overlap=chunk_overlap,
853870
transformation_config=transformation_config,
854871
max_embedding_requests_per_min=max_embedding_requests_per_min,
872+
global_max_embedding_requests_per_min=global_max_embedding_requests_per_min,
855873
use_advanced_pdf_parsing=use_advanced_pdf_parsing,
856874
partial_failures_sink=partial_failures_sink,
857875
layout_parser=layout_parser,

‎vertexai/preview/rag/utils/_gapic_utils.py

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -537,6 +537,7 @@ def prepare_import_files_request(
537537
chunk_overlap: int = 200,
538538
transformation_config: Optional[TransformationConfig] = None,
539539
max_embedding_requests_per_min: int = 1000,
540+
global_max_embedding_requests_per_min: Optional[int] = None,
540541
use_advanced_pdf_parsing: bool = False,
541542
partial_failures_sink: Optional[str] = None,
542543
layout_parser: Optional[LayoutParserConfig] = None,
@@ -569,8 +570,15 @@ def prepare_import_files_request(
569570
)
570571
rag_file_parsing_config.layout_parser = RagFileParsingConfig.LayoutParser(
571572
processor_name=layout_parser.processor_name,
572-
max_parsing_requests_per_min=layout_parser.max_parsing_requests_per_min,
573573
)
574+
if layout_parser.max_parsing_requests_per_min is not None:
575+
rag_file_parsing_config.layout_parser.max_parsing_requests_per_min = (
576+
layout_parser.max_parsing_requests_per_min
577+
)
578+
if layout_parser.global_max_parsing_requests_per_min is not None:
579+
rag_file_parsing_config.layout_parser.global_max_parsing_requests_per_min = (
580+
layout_parser.global_max_parsing_requests_per_min
581+
)
574582
if llm_parser is not None:
575583
rag_file_parsing_config.llm_parser = RagFileParsingConfig.LlmParser(
576584
model_name=llm_parser.model_name
@@ -609,6 +617,10 @@ def prepare_import_files_request(
609617
rebuild_ann_index=rebuild_ann_index,
610618
)
611619

620+
if global_max_embedding_requests_per_min is not None:
621+
import_rag_files_config.global_max_embedding_requests_per_min = (
622+
global_max_embedding_requests_per_min
623+
)
612624
if source is not None:
613625
gapic_source = convert_source_for_rag_import(source)
614626
if isinstance(gapic_source, GapicSlackSource):

‎vertexai/preview/rag/utils/resources.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -515,10 +515,19 @@ class LayoutParserConfig:
515515
https://cloud.google.com/document-ai/quotas and the Quota page for
516516
your project to set an appropriate value here. If unspecified, a
517517
default value of 120 QPM will be used.
518+
global_max_parsing_requests_per_min (int):
519+
The maximum number of requests the job is allowed to make to
520+
the Document AI processor per minute in this project.
521+
Consult https://cloud.google.com/document-ai/quotas and the
522+
Quota page for your project to set an appropriate value
523+
here. If this value is not specified,
524+
max_parsing_requests_per_min will be used by indexing
525+
pipeline as the global limit.
518526
"""
519527

520528
processor_name: str
521529
max_parsing_requests_per_min: Optional[int] = None
530+
global_max_parsing_requests_per_min: Optional[int] = None
522531

523532

524533
@dataclasses.dataclass

0 commit comments

Comments
 (0)
0