From fc35b78eb42b665c637d7b13969de57d88f2c6a0 Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 20:16:07 +0000 Subject: [PATCH] ci: drop aws partition hardcode gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-code-quality.yml | 3 - .../check_aws_partition_hardcodes.py | 140 ------------------ .../litellm_core_utils/test_aws_partition.py | 116 +-------------- 3 files changed, 1 insertion(+), 258 deletions(-) delete mode 100644 tests/code_coverage_tests/check_aws_partition_hardcodes.py diff --git a/.github/workflows/test-code-quality.yml b/.github/workflows/test-code-quality.yml index 44c3e97db91..987f66773f2 100644 --- a/.github/workflows/test-code-quality.yml +++ b/.github/workflows/test-code-quality.yml @@ -146,9 +146,6 @@ jobs: - name: check_migrations_no_data_rewrites run: uv run --no-sync python ./tests/code_coverage_tests/check_migrations_no_data_rewrites.py - - name: check_aws_partition_hardcodes - run: uv run --no-sync python ./tests/code_coverage_tests/check_aws_partition_hardcodes.py - - name: memory_test run: uv run --no-sync python ./tests/code_coverage_tests/memory_test.py diff --git a/tests/code_coverage_tests/check_aws_partition_hardcodes.py b/tests/code_coverage_tests/check_aws_partition_hardcodes.py deleted file mode 100644 index 0cbee4e7c80..00000000000 --- a/tests/code_coverage_tests/check_aws_partition_hardcodes.py +++ /dev/null @@ -1,140 +0,0 @@ -#!/usr/bin/env python3 -"""Ban hardcoded commercial-partition AWS hosts and ARN prefixes under `litellm/`. - -An endpoint or ARN built with a literal `amazonaws.com` or `arn:aws:` works in every -commercial region and breaks only for GovCloud (`us-gov-*`, `arn:aws-us-gov:`) and -China (`amazonaws.com.cn`, `arn:aws-cn:`) deployments, so the failure never shows up -in CI or on a developer laptop. `litellm/litellm_core_utils/aws_partition.py` derives -both from the region and is the only place those literals belong. Build hosts with -`get_aws_dns_suffix(region)` and ARNs with `get_aws_arn_prefix(region)`. - -Every string constant in every `litellm/**/*.py` file is scanned, including the -literal parts of f-strings and the strings inside `.format()` calls and -concatenations. Docstrings and comments are not, since they never reach a request. -`amazonaws.com.cn` passes because it is already the China partition. - -`ALLOWED` holds the (file, token, count) triples that are text rather than a request -target: a hosted logo, an IAM service principal, and hostnames quoted as examples -inside error messages and field descriptions. An entry only covers that many -occurrences of that exact token in that exact file, so a second copy of an allowed -literal is still caught, and an entry whose token is gone or whose count has changed -fails the check so the set only shrinks. -""" - -from __future__ import annotations - -import ast -import re -import sys -from collections import Counter -from pathlib import Path -from types import MappingProxyType -from typing import Final, NamedTuple - -REPO_ROOT: Final = Path(__file__).resolve().parents[2] -SCAN_ROOT: Final = REPO_ROOT / "litellm" -PARTITION_HELPER: Final = "litellm/litellm_core_utils/aws_partition.py" - -COMMERCIAL_TOKEN: Final = re.compile(r"[A-Za-z0-9.-]*amazonaws\.com(?!\.cn)|arn:aws:[A-Za-z0-9:/_.*-]*") - - -class Allowance(NamedTuple): - file: str - token: str - occurrences: int - - -ALLOWED: Final = frozenset( - { - Allowance("litellm/integrations/email_alerting.py", "litellm-listing.s3.amazonaws.com", 1), - Allowance("litellm/types/integrations/slack_alerting.py", "litellm-listing.s3.amazonaws.com", 1), - Allowance("litellm/rag/ingestion/bedrock_ingestion.py", "bedrock.amazonaws.com", 1), - Allowance( - "litellm/llms/bedrock/chat/agentcore/transformation.py", - "arn:aws:bedrock-agentcore:region:account:runtime/runtime_id", - 1, - ), - Allowance("litellm/llms/bedrock/search/transformation.py", ".amazonaws.com", 1), - Allowance( - "litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py", - "bucket.s3.amazonaws.com", - 1, - ), - Allowance("litellm/types/proxy/claude_code_endpoints.py", "bucket.s3.amazonaws.com", 1), - } -) -ALLOWED_COUNTS: Final = MappingProxyType({(entry.file, entry.token): entry.occurrences for entry in ALLOWED}) - - -class Hit(NamedTuple): - file: str - line: int - token: str - - -def _docstring_ids(tree: ast.Module) -> frozenset[int]: - return frozenset( - id(statement.value) - for node in ast.walk(tree) - if isinstance(node, (ast.Module, ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)) - for statement in node.body - if isinstance(statement, ast.Expr) - and isinstance(statement.value, ast.Constant) - and isinstance(statement.value.value, str) - ) - - -def _hits_in_file(path: Path) -> tuple[Hit, ...]: - tree: Final = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) - docstrings: Final = _docstring_ids(tree) - relative: Final = path.relative_to(REPO_ROOT).as_posix() - return tuple( - Hit(relative, node.lineno, match.group(0)) - for node in ast.walk(tree) - if isinstance(node, ast.Constant) and isinstance(node.value, str) and id(node) not in docstrings - for match in COMMERCIAL_TOKEN.finditer(node.value) - ) - - -def find_hits(scan_root: Path) -> tuple[Hit, ...]: - return tuple( - hit - for path in sorted(scan_root.rglob("*.py")) - if path.relative_to(REPO_ROOT).as_posix() != PARTITION_HELPER - for hit in _hits_in_file(path) - ) - - -def _violation_message(hit: Hit, found: int) -> str: - allowed: Final = ALLOWED_COUNTS.get((hit.file, hit.token)) - if allowed is None: - return f"{hit.file}:{hit.line}: hardcoded commercial AWS partition literal {hit.token!r}" - return ( - f"{hit.file}:{hit.line}: {hit.token!r} appears {found} times but ALLOWED covers {allowed}; " - "build it from the region helper or update the count" - ) - - -def main() -> int: - hits: Final = find_hits(SCAN_ROOT) - counts: Final = MappingProxyType(Counter((hit.file, hit.token) for hit in hits)) - violations: Final = tuple( - sorted(hit for hit in hits if Allowance(hit.file, hit.token, counts[hit.file, hit.token]) not in ALLOWED) - ) - stale: Final = tuple(entry for entry in sorted(ALLOWED) if (entry.file, entry.token) not in counts) - for hit in violations: - print(_violation_message(hit, counts[hit.file, hit.token])) - for allowance in stale: - print(f"{allowance.file}: ALLOWED entry {allowance.token!r} no longer matches anything, remove it") - if violations or stale: - print( - "\nBuild AWS hosts with get_aws_dns_suffix(region) and ARNs with get_aws_arn_prefix(region) " - "from litellm/litellm_core_utils/aws_partition.py so GovCloud and China regions resolve." - ) - return 1 - print(f"No hardcoded commercial AWS partition literals outside {PARTITION_HELPER}.") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/tests/test_litellm/litellm_core_utils/test_aws_partition.py b/tests/test_litellm/litellm_core_utils/test_aws_partition.py index 24a38268ae9..3594d3c354c 100644 --- a/tests/test_litellm/litellm_core_utils/test_aws_partition.py +++ b/tests/test_litellm/litellm_core_utils/test_aws_partition.py @@ -1,11 +1,9 @@ import ast from pathlib import Path -from types import MappingProxyType from typing import Final -from urllib.parse import unquote, urlparse +from urllib.parse import urlparse import pytest -from botocore.credentials import Credentials import litellm from litellm.integrations.s3_v2 import S3Logger @@ -22,20 +20,8 @@ from litellm.llms.aws_polly.text_to_speech.transformation import AWSPollyTextToS from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig -from litellm.llms.bedrock.chat.invoke_agent.transformation import AmazonInvokeAgentConfig from litellm.llms.bedrock.common_utils import init_bedrock_client -from litellm.llms.bedrock.files.transformation import BedrockFilesConfig -from litellm.llms.bedrock.rerank.handler import BedrockRerankHandler -from litellm.llms.bedrock.vector_stores.transformation import BedrockVectorStoreConfig from litellm.llms.sagemaker.chat.transformation import SagemakerChatConfig -from litellm.llms.sagemaker.completion.handler import SagemakerLLM -from litellm.proxy.auth.rds_iam_token import init_rds_client -from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail -from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2 - -STATIC_AWS_CREDENTIALS: Final = MappingProxyType( - {"aws_access_key_id": "test-key", "aws_secret_access_key": "test-secret"} -) @pytest.mark.parametrize( @@ -120,48 +106,6 @@ def _s3_object_url(region: str) -> str: return logger._build_object_url("2025-01-01/key.json") -def _bedrock_job_arn(region: str) -> str: - return f"{get_aws_arn_prefix(region)}bedrock:{region}:111122223333:model-invocation-job/abc1234567" - - -def _bedrock_files_upload_url(region: str) -> str: - return BedrockFilesConfig().get_complete_file_url( - api_base=None, - api_key=None, - model="amazon.nova-pro-v1:0", - optional_params={}, - litellm_params={"s3_bucket_name": "batch-bucket", "s3_region_name": region}, - data={"file": ("batch.jsonl", b"{}", "application/jsonl"), "purpose": "batch"}, - ) - - -def _bedrock_files_download_url(region: str) -> str: - return ( - BedrockFilesConfig() - ._s3_request_target(optional_params={}, litellm_params={"s3_region_name": region}) - .endpoint_url - ) - - -def _bedrock_guardrail_url(region: str) -> str: - guardrail = BedrockGuardrail(guardrailIdentifier="guardrail-id", guardrailVersion="1") - return guardrail._prepare_request( - credentials=Credentials("test-key", "test-secret"), - data={"source": "INPUT", "content": []}, - optional_params={}, - aws_region_name=region, - ).url - - -def _secrets_manager_url(region: str) -> str: - endpoint_url, _headers, _body = AWSSecretsManagerV2(aws_region_name=region)._prepare_request( - action="GetSecretValue", - secret_name="my-secret", - optional_params=dict(STATIC_AWS_CREDENTIALS), - ) - return endpoint_url - - ENDPOINT_BUILDERS: Final = { "bedrock_runtime_default": lambda region: BaseAWSLLM()._select_default_endpoint_url("runtime", region), "bedrock_agent_default": lambda region: BaseAWSLLM()._select_default_endpoint_url("agent", region), @@ -180,13 +124,6 @@ ENDPOINT_BUILDERS: Final = { litellm_params={}, data={"input_file_id": "s3://bucket/key.jsonl"}, ), - "bedrock_batches_retrieve": lambda region: BedrockBatchesConfig().transform_retrieve_batch_request( - batch_id=_bedrock_job_arn(region), - optional_params=dict(STATIC_AWS_CREDENTIALS), - litellm_params={}, - )["url"], - "bedrock_files_upload": _bedrock_files_upload_url, - "bedrock_files_download": _bedrock_files_download_url, "bedrock_agentcore_invoke": lambda region: AmazonAgentCoreConfig().get_complete_url( api_base=None, api_key=None, @@ -194,32 +131,6 @@ ENDPOINT_BUILDERS: Final = { optional_params={}, litellm_params={}, ), - "bedrock_invoke_agent": lambda region: AmazonInvokeAgentConfig().get_complete_url( - api_base=None, - api_key=None, - model="agent/AGENT123/ALIAS456", - optional_params={"aws_region_name": region}, - litellm_params={}, - ), - "bedrock_guardrail_apply": _bedrock_guardrail_url, - "bedrock_rerank": lambda region: BedrockRerankHandler()._prepare_request( - model="amazon.rerank-v1:0", - api_base=None, - extra_headers=None, - data={"queries": [], "sources": []}, - optional_params={"aws_region_name": region, **STATIC_AWS_CREDENTIALS}, - )["endpoint_url"], - "bedrock_knowledgebase_search": lambda region: BedrockVectorStoreConfig().get_complete_url( - api_base=None, litellm_params={"aws_region_name": region} - ), - "secrets_manager": _secrets_manager_url, - "rds_iam_client": lambda region: ( - init_rds_client( - aws_region_name=region, - aws_access_key_id="test-key", - aws_secret_access_key="test-secret", - ).meta.endpoint_url - ), "polly": lambda region: AWSPollyTextToSpeechConfig().get_complete_url( model="polly/neural", api_base=None, @@ -241,19 +152,6 @@ ENDPOINT_BUILDERS: Final = { litellm_params={}, stream=True, ), - "sagemaker_completion": lambda region: ( - SagemakerLLM() - ._prepare_request( - credentials=Credentials("test-key", "test-secret"), - model="my-endpoint", - data={}, - messages=[], - litellm_params={}, - optional_params={}, - aws_region_name=region, - ) - .url - ), "s3_object_url": _s3_object_url, } @@ -284,18 +182,6 @@ def test_every_endpoint_builder_keeps_amazonaws_com_outside_cn(builder_name: str assert hostname.endswith(".amazonaws.com"), url -@pytest.mark.parametrize("region", ["us-gov-west-1", "us-gov-east-1"]) -@pytest.mark.parametrize("builder_name", sorted(ENDPOINT_BUILDERS)) -def test_every_endpoint_builder_respects_us_gov_partition(builder_name: str, region: str) -> None: - url = unquote(ENDPOINT_BUILDERS[builder_name](region)) - hostname = urlparse(url).hostname - assert hostname is not None - assert hostname.endswith(f".{region}.amazonaws.com"), url - assert "arn:aws:" not in url, url - if "arn:" in url: - assert "arn:aws-us-gov:" in url, url - - def _fstring_literal_offenders(needle: str) -> list[str]: litellm_root = Path(litellm.__file__).parent return [