⚡️ Speed up function is_prebuilt_prediction_container_uri by 52%

[codeflash-ai]

📄 52% (0.52x) speedup for is_prebuilt_prediction_container_uri in google/cloud/aiplatform/helpers/container_uri_builders.py

⏱️ Runtime : 2.34 milliseconds → 1.54 milliseconds (best of 245 runs)

📝 Explanation and details

The optimization precompiles the regex pattern at module import time instead of recompiling it on every function call.

Key Change:

• Added _CONTAINER_URI_REGEX = re.compile(prediction.CONTAINER_URI_REGEX) at module level
• Changed from re.fullmatch(prediction.CONTAINER_URI_REGEX, image_uri) to _CONTAINER_URI_REGEX.fullmatch(image_uri)

Why It's Faster:
Python's re.fullmatch() with a string pattern internally compiles the regex on every call, which is expensive. By precompiling the regex once at import time, each function call skips the compilation step and directly uses the compiled pattern object. The line profiler shows the per-hit time dropped from 1961ns to 614ns (68% improvement per call).

Performance Benefits:
This optimization excels in scenarios with repeated calls to the same function, which is evident in the test results:

• Single calls: 30-70% faster (e.g., basic URI validation)
• Batch operations: 48-60% faster (e.g., validating 500-1000 URIs)
• High-frequency validation scenarios benefit most, as the compilation cost is amortized across many calls

The 51% overall speedup demonstrates that regex compilation was a significant bottleneck, making this a highly effective optimization for any code path that validates container URIs repeatedly.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 3563 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import re

# imports
import pytest  # used for our unit tests
from aiplatform.helpers.container_uri_builders import \
    is_prebuilt_prediction_container_uri

# Simulate the relevant constant for the test environment
# This regex is assumed from the context of Vertex AI's prebuilt prediction container URIs.
# Example: 'us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-11:latest'
CONTAINER_URI_REGEX = (
    r"^([a-z]+-docker.pkg.dev/vertex-ai/prediction/([a-z0-9\-]+)\.[a-z0-9\-]+:[a-z0-9\-\.]+)$"
)
from aiplatform.helpers.container_uri_builders import \
    is_prebuilt_prediction_container_uri

# ------------------ UNIT TESTS ------------------

# 1. Basic Test Cases

def test_valid_prebuilt_uris_basic():
    # Test with typical valid prebuilt URIs
    codeflash_output = is_prebuilt_prediction_container_uri(
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-11:latest"
    ) # 2.99μs -> 1.72μs (73.8% faster)
    codeflash_output = is_prebuilt_prediction_container_uri(
        "europe-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-24:latest"
    ) # 1.19μs -> 890ns (33.7% faster)
    codeflash_output = is_prebuilt_prediction_container_uri(
        "asia-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-4:latest"
    ) # 834ns -> 588ns (41.8% faster)
    # Test with valid version tags
    codeflash_output = is_prebuilt_prediction_container_uri(
        "us-docker.pkg.dev/vertex-ai/prediction/pytorch-cpu.1-10:1.0.0"
    ) # 733ns -> 519ns (41.2% faster)
    codeflash_output = is_prebuilt_prediction_container_uri(
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-11:beta"
    ) # 706ns -> 504ns (40.1% faster)

def test_invalid_prebuilt_uris_basic():
    # Test with URIs that are not prebuilt prediction containers
    codeflash_output = not is_prebuilt_prediction_container_uri(
        "gcr.io/my-project/my-custom-image:latest"
    ) # 1.57μs -> 760ns (106% faster)
    codeflash_output = not is_prebuilt_prediction_container_uri(
        "docker.io/library/tensorflow:2.11"
    ) # 680ns -> 420ns (61.9% faster)
    codeflash_output = not is_prebuilt_prediction_container_uri(
        "us-docker.pkg.dev/vertex-ai/training/tf2-cpu.2-11:latest"
    ) # 853ns -> 656ns (30.0% faster)
    codeflash_output = not is_prebuilt_prediction_container_uri(
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-11"
    ) # 961ns -> 704ns (36.5% faster)
    codeflash_output = not is_prebuilt_prediction_container_uri(
        "us-docker.pkg.dev/vertex-ai/prediction/:latest"
    ) # 693ns -> 497ns (39.4% faster)

# 2. Edge Test Cases

def test_empty_and_none_inputs():
    # Empty string should not match
    codeflash_output = not is_prebuilt_prediction_container_uri("") # 1.49μs -> 761ns (96.1% faster)
    # None should raise TypeError
    with pytest.raises(TypeError):
        is_prebuilt_prediction_container_uri(None) # 1.71μs -> 1.10μs (55.6% faster)

def test_whitespace_and_case_sensitivity():
    # Leading/trailing whitespace should not match
    codeflash_output = not is_prebuilt_prediction_container_uri(
        " us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-11:latest"
    ) # 1.75μs -> 916ns (91.2% faster)
    codeflash_output = not is_prebuilt_prediction_container_uri(
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-11:latest "
    ) # 1.42μs -> 1.03μs (37.7% faster)
    # Case sensitivity: should fail if region is capitalized
    codeflash_output = not is_prebuilt_prediction_container_uri(
        "US-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-11:latest"
    ) # 595ns -> 353ns (68.6% faster)
    # Internal spaces should fail
    codeflash_output = not is_prebuilt_prediction_container_uri(
        "us-docker.pkg.dev/vertex-ai/prediction/tf2 cpu.2-11:latest"
    ) # 754ns -> 555ns (35.9% faster)

def test_unusual_but_valid_tags_and_names():
    # Test with unusual but valid tag formats
    codeflash_output = is_prebuilt_prediction_container_uri(
        "us-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-24:release-2023.01"
    ) # 2.12μs -> 1.41μs (49.6% faster)
    codeflash_output = is_prebuilt_prediction_container_uri(
        "us-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-4:rc-1"
    ) # 1.05μs -> 825ns (27.3% faster)
    # Test with hyphens and numbers in image name
    codeflash_output = is_prebuilt_prediction_container_uri(
        "us-docker.pkg.dev/vertex-ai/prediction/my-custom-cpu.1-2:latest"
    ) # 724ns -> 496ns (46.0% faster)

def test_invalid_uri_formats():
    # Missing region prefix
    codeflash_output = not is_prebuilt_prediction_container_uri(
        "docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-11:latest"
    ) # 1.53μs -> 866ns (76.7% faster)
    # Wrong registry domain
    codeflash_output = not is_prebuilt_prediction_container_uri(
        "us-docker.pkg.com/vertex-ai/prediction/tf2-cpu.2-11:latest"
    ) # 915ns -> 647ns (41.4% faster)
    # Wrong path
    codeflash_output = not is_prebuilt_prediction_container_uri(
        "us-docker.pkg.dev/vertex-ai/predict/tf2-cpu.2-11:latest"
    ) # 680ns -> 473ns (43.8% faster)
    # Extra path segments
    codeflash_output = not is_prebuilt_prediction_container_uri(
        "us-docker.pkg.dev/vertex-ai/prediction/extra/tf2-cpu.2-11:latest"
    ) # 687ns -> 493ns (39.4% faster)
    # Tag with invalid characters
    codeflash_output = not is_prebuilt_prediction_container_uri(
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-11:la test"
    ) # 916ns -> 726ns (26.2% faster)

def test_partial_and_substring_matches():
    # Should not match partial substrings
    codeflash_output = not is_prebuilt_prediction_container_uri(
        "us-docker.pkg.dev/vertex-ai/prediction/"
    ) # 1.44μs -> 737ns (95.3% faster)
    codeflash_output = not is_prebuilt_prediction_container_uri(
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-11"
    ) # 1.38μs -> 1.09μs (26.7% faster)
    # Should not match if tag is missing
    codeflash_output = not is_prebuilt_prediction_container_uri(
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-11:"
    ) # 778ns -> 553ns (40.7% faster)

def test_long_image_names_and_tags():
    # Very long but valid image name and tag
    long_name = "my-model-" + "x" * 50 + ".cpu.1-0"
    long_tag = "release-" + "y" * 50
    uri = f"us-docker.pkg.dev/vertex-ai/prediction/{long_name}:{long_tag}"
    codeflash_output = is_prebuilt_prediction_container_uri(uri) # 1.92μs -> 1.03μs (86.1% faster)
    # Very long tag with numbers and dots
    uri2 = f"us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-11:{'1.'*100}latest"
    codeflash_output = is_prebuilt_prediction_container_uri(uri2) # 1.32μs -> 938ns (40.7% faster)

# 3. Large Scale Test Cases

def test_many_valid_and_invalid_uris():
    # Generate 500 valid URIs
    for i in range(500):
        region = "us" if i % 2 == 0 else "europe"
        framework = f"tf{i % 10}-cpu"
        version = f"{i % 5}-{i % 20}"
        tag = f"release-{i}"
        uri = f"{region}-docker.pkg.dev/vertex-ai/prediction/{framework}.{version}:{tag}"
        codeflash_output = is_prebuilt_prediction_container_uri(uri) # 335μs -> 225μs (48.6% faster)
    # Generate 500 invalid URIs (wrong path)
    for i in range(500):
        region = "us" if i % 2 == 0 else "europe"
        framework = f"tf{i % 10}-cpu"
        version = f"{i % 5}-{i % 20}"
        tag = f"release-{i}"
        uri = f"{region}-docker.pkg.dev/vertex-ai/train/{framework}.{version}:{tag}"
        codeflash_output = not is_prebuilt_prediction_container_uri(uri) # 281μs -> 175μs (60.6% faster)

def test_performance_large_batch():
    # Test efficiency with a batch of 1000 URIs
    valid_uris = [
        f"us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.{i}:latest"
        for i in range(500)
    ]
    invalid_uris = [
        f"us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.{i}latest"
        for i in range(500)
    ]  # missing colon
    # All valid URIs should match
    for uri in valid_uris:
        codeflash_output = is_prebuilt_prediction_container_uri(uri) # 320μs -> 215μs (48.4% faster)
    # All invalid URIs should not match
    for uri in invalid_uris:
        codeflash_output = not is_prebuilt_prediction_container_uri(uri) # 313μs -> 210μs (48.5% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import re

# imports
import pytest  # used for our unit tests
from aiplatform.helpers.container_uri_builders import \
    is_prebuilt_prediction_container_uri

# Simulate google.cloud.aiplatform.constants.prediction.CONTAINER_URI_REGEX for testing purposes.
# This regex is based on typical Vertex AI prebuilt container URIs.
# Example: 'us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:latest'
CONTAINER_URI_REGEX = (
    r"(?P<region>[a-z]+)-docker\.pkg\.dev/vertex-ai/prediction/"
    r"(?P<framework>[a-z0-9\-]+)\.(?P<hardware>[a-z0-9\-]+)\.(?P<version>[0-9\-]+):(?P<tag>[a-z0-9\-\.]+)"
)
from aiplatform.helpers.container_uri_builders import \
    is_prebuilt_prediction_container_uri

# ---------------------- UNIT TESTS ----------------------

# 1. Basic Test Cases
@pytest.mark.parametrize(
    "image_uri",
    [
        # Typical valid URIs
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:latest",
        "europe-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.0-24:latest",
        "asia-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-5:latest",
        "us-docker.pkg.dev/vertex-ai/prediction/pytorch-cpu.1-10:latest",
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-8:latest",
    ],
)
def test_basic_valid_prebuilt_uris(image_uri):
    """Test valid prebuilt container URIs."""
    codeflash_output = is_prebuilt_prediction_container_uri(image_uri) # 12.1μs -> 7.81μs (55.0% faster)

@pytest.mark.parametrize(
    "image_uri",
    [
        # Invalid: wrong domain
        "us-docker.pkg.dev/vertex-ai/training/tf2-cpu.2-8:latest",
        "us-docker.pkg.dev/vertex-ai/prediction/tf2cpu.2-8:latest",  # missing dash
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu2-8:latest",  # missing dot
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8-latest", # missing colon
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8",        # missing tag
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:",       # empty tag
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:lat est",# space in tag
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:latest!",# invalid char in tag
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:latest:extra", # extra colon
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8",        # no tag
        "us-docker.pkg.dev/vertex-ai/prediction/:latest",            # missing framework/hardware/version
    ],
)
def test_basic_invalid_prebuilt_uris(image_uri):
    """Test invalid prebuilt container URIs."""
    codeflash_output = is_prebuilt_prediction_container_uri(image_uri) # 24.8μs -> 15.9μs (56.1% faster)

# 2. Edge Test Cases

@pytest.mark.parametrize(
    "image_uri",
    [
        "",  # empty string
        " ", # whitespace only
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:latest\n", # trailing newline
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:latest ",  # trailing space
        "us-docker.pkg.dev/vertex-ai/prediction/Tf2-cpu.2-8:latest",   # uppercase framework
        "US-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:latest",   # uppercase region
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:la_test",  # underscore in tag
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:lat.est",  # dot in tag (should be valid)
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:la-test",  # dash in tag (should be valid)
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:123",      # numeric tag (should be valid)
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:la$test",  # special char in tag
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:la/test",  # slash in tag
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:la:test",  # colon in tag
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:la.test.extra", # multiple dots in tag
    ],
)
def test_edge_cases(image_uri):
    """Test edge cases for input strings."""
    # Only valid tags: lowercase letters, numbers, dash, dot
    # So, only those with valid tags should pass
    valid = re.fullmatch(CONTAINER_URI_REGEX, image_uri) is not None
    codeflash_output = is_prebuilt_prediction_container_uri(image_uri) # 22.7μs -> 16.6μs (37.2% faster)

@pytest.mark.parametrize(
    "image_uri",
    [
        # Minimal valid values
        "us-docker.pkg.dev/vertex-ai/prediction/a-b.1-1:tag",   # minimal framework/hardware/version/tag
        "us-docker.pkg.dev/vertex-ai/prediction/abc-def.0-0:latest", # minimal numbers
        # Maximal length values (within reason)
        (
            "us-docker.pkg.dev/vertex-ai/prediction/"
            + "a"*50 + "-" + "b"*50 + "." + "c"*50 + "-" + "d"*50 + "." + "1"*50 + "-" + "2"*50
            + ":latest"
        ),
        # Tag with multiple valid segments
        "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:latest-rc.1",
    ],
)
def test_edge_min_max(image_uri):
    """Test minimal and maximal valid values."""
    codeflash_output = is_prebuilt_prediction_container_uri(image_uri) # 8.43μs -> 5.44μs (54.8% faster)

# 3. Large Scale Test Cases

def test_large_scale_all_valid():
    """Test a large batch of valid URIs."""
    # Generate 500 valid URIs with different regions/frameworks/hardware/versions/tags
    for i in range(500):
        region = f"us"
        framework = f"tf{i % 10}-cpu"
        hardware = f"{i % 5}-{i % 10}"
        version = f"{i % 20}-{i % 30}"
        tag = f"latest-{i}"
        uri = f"{region}-docker.pkg.dev/vertex-ai/prediction/{framework}.{hardware}.{version}:{tag}"
        codeflash_output = is_prebuilt_prediction_container_uri(uri) # 333μs -> 221μs (50.5% faster)

def test_large_scale_all_invalid():
    """Test a large batch of invalid URIs."""
    # Generate 500 invalid URIs by breaking the pattern
    for i in range(500):
        region = f"us"
        framework = f"tf{i % 10}cpu"  # missing dash
        hardware = f"{i % 5}{i % 10}" # missing dash
        version = f"{i % 20}{i % 30}" # missing dash
        tag = f"latest_{i}"           # underscore (invalid)
        uri = f"{region}-docker.pkg.dev/vertex-ai/prediction/{framework}.{hardware}.{version}:{tag}"
        codeflash_output = is_prebuilt_prediction_container_uri(uri) # 332μs -> 219μs (51.2% faster)

def test_large_scale_random_mixed():
    """Test a large batch of mixed valid and invalid URIs."""
    # Mix valid and invalid URIs
    for i in range(500):
        if i % 2 == 0:
            # Valid
            region = "asia"
            framework = f"xgboost-cpu"
            hardware = f"{i % 5}-{i % 10}"
            version = f"{i % 20}-{i % 30}"
            tag = f"latest.{i}"
            uri = f"{region}-docker.pkg.dev/vertex-ai/prediction/{framework}.{hardware}.{version}:{tag}"
            codeflash_output = is_prebuilt_prediction_container_uri(uri)
        else:
            # Invalid: wrong domain
            region = "asia"
            framework = f"xgboostcpu"
            hardware = f"{i % 5}{i % 10}"
            version = f"{i % 20}{i % 30}"
            tag = f"latest_{i}"
            uri = f"{region}-docker.pkg.dev/vertex-ai/training/{framework}.{hardware}.{version}:{tag}"
            codeflash_output = is_prebuilt_prediction_container_uri(uri)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-is_prebuilt_prediction_container_uri-mglkk5rl and push.