⚡️ Speed up method PipelineJob.get by 53%

[codeflash-ai]

📄 53% (0.53x) speedup for PipelineJob.get in google/cloud/aiplatform/pipeline_jobs.py

⏱️ Runtime : 29.3 milliseconds → 19.1 milliseconds (best of 5 runs)

📝 Explanation and details

The optimization achieves a 53% speedup by eliminating redundant dictionary lookups and computations in the PipelineJob.__init__ method.

Key optimizations applied:

1. Cached dictionary lookups: Instead of repeatedly calling pipeline_json.get("pipelineSpec") and pipeline_json.get("runtimeConfig"), these values are retrieved once and stored in local variables (pipeline_spec and runtime_config). This eliminates multiple dictionary key lookups on the same object.

2. Reduced nested attribute access: The deeply nested access pipeline_job["pipelineSpec"]["pipelineInfo"]["name"] is broken down into intermediate variables (pipeline_info and pipeline_name_value), reducing the chain of dictionary lookups.

3. Pre-computed regex operation: The expensive regex substitution re.sub("[^-0-9a-z]+", "-", pipeline_name_value.lower()).lstrip("-").rstrip("-") is computed once and stored in pipeline_name_key, avoiding redundant string processing.

4. Streamlined pipeline root resolution: The cascading fallback logic for determining pipeline_root is restructured to use cached values (default_pipeline_root, runtime_gcs_output_dir) instead of repeated dictionary access.

5. Variable renaming for clarity: Using gca_runtime_config instead of runtime_config to avoid naming conflicts and improve code readability.

Why this works: Dictionary lookups and nested attribute access are relatively expensive operations in Python. By caching frequently accessed values in local variables, the optimizer reduces the number of hash table lookups and attribute resolution calls, leading to faster execution.

Test case performance: The optimizations show consistent improvements across all test cases, with the most significant gains (100%+ speedup) in error-handling scenarios where the reduced overhead in setup code before exceptions are raised provides substantial benefits.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from typing import Any, Dict, Optional

# imports
import pytest
from aiplatform.pipeline_jobs import PipelineJob

# --- Minimal stubs and mocks to make the test self-contained ---

# Simulate proto.Message
class DummyProto:
    def __init__(self, name):
        self.name = name

# Simulate the GAPIC client class with parse/format methods
class DummyGapicClient:
    @staticmethod
    def parse_pipeline_job_path(resource_name):
        # Accepts only fully qualified names of the form:
        # "projects/{project}/locations/{location}/pipelineJobs/{job_id}"
        import re
        m = re.match(r"^projects/(?P<project>[^/]+)/locations/(?P<location>[^/]+)/pipelineJobs/(?P<job_id>[^/]+)$", resource_name)
        if m:
            return m.groupdict()
        return {}

    @staticmethod
    def pipeline_job_path(project, location, pipeline_job):
        return f"projects/{project}/locations/{location}/pipelineJobs/{pipeline_job}"

# Simulate the client class wrapper
class DummyClientClass:
    @staticmethod
    def get_gapic_client_class():
        return DummyGapicClient

# Simulate the API client with a get_pipeline_job method
class DummyApiClient:
    def __init__(self):
        self.called = []
    def get_pipeline_job(self, name, retry=None):
        # Just return a DummyProto with the name
        self.called.append(name)
        return DummyProto(name)

# Simulate global config
class DummyGlobalConfig:
    project = "test-proj"
    location = "us-central1"
    credentials = "dummy-credentials"
    def create_client(self, client_class, credentials, location_override, appended_user_agent):
        return DummyApiClient()
    def get_encryption_spec(self, encryption_spec_key_name):
        return None
    def common_location_path(self, project, location):
        return f"projects/{project}/locations/{location}"

def dummy_full_resource_name(resource_name, resource_noun, parse_resource_name_method, format_resource_name_method, parent_resource_name_fields=None, project=None, location=None, resource_id_validator=None):
    fields = parse_resource_name_method(resource_name)
    if fields:
        return resource_name
    # Validate the ID
    if resource_id_validator:
        resource_id_validator(resource_name)
    else:
        if not resource_name or not isinstance(resource_name, str):
            raise ValueError("Invalid resource id.")
    return format_resource_name_method(project=project, location=location, pipeline_job=resource_name)

def dummy_resource_id_validator(resource_id):
    if not resource_id or not isinstance(resource_id, str) or "/" in resource_id:
        raise ValueError("Invalid resource id.")
from aiplatform.pipeline_jobs import PipelineJob

# --- Pytest unit tests for PipelineJob.get ---

# 1. Basic Test Cases

















#------------------------------------------------
import builtins
import re
# Patch the initializer and utils for tests
import sys
import types

# imports
import pytest  # used for our unit tests
from aiplatform.pipeline_jobs import PipelineJob


# Mocks and helpers for testing
class DummyCredentials:
    pass

class DummyProtoMessage:
    def __init__(self, name):
        self.name = name

class DummyClientClass:
    @staticmethod
    def get_gapic_client_class():
        return DummyClientClass

    @staticmethod
    def parse_pipeline_job_path(resource_name):
        # Simulate parsing a fully qualified resource name
        if resource_name.startswith("projects/") and "/pipelineJobs/" in resource_name:
            parts = resource_name.split("/")
            return {
                "project": parts[1],
                "location": parts[3],
                "pipelineJobs": parts[5]
            }
        return {}

    @staticmethod
    def pipeline_job_path(project, location, pipeline_jobs):
        # Simulate formatting a resource name
        return f"projects/{project}/locations/{location}/pipelineJobs/{pipeline_jobs}"

    @staticmethod
    def get_pipeline_job(name, retry=None):
        # Simulate API client returning a proto message
        return DummyProtoMessage(name)

class DummyInitializer:
    class global_config:
        project = "test-project"
        location = "us-central1"
        credentials = DummyCredentials()
        @staticmethod
        def create_client(client_class, credentials=None, location_override=None, appended_user_agent=None):
            return DummyClientClass

        @staticmethod
        def common_location_path(project=None, location=None):
            return f"projects/{project or DummyInitializer.global_config.project}/locations/{location or DummyInitializer.global_config.location}"

        @staticmethod
        def get_encryption_spec(encryption_spec_key_name=None):
            return None

        staging_bucket = "gs://test-staging-bucket"

# Patch pipeline constants for tests
class DummyPipelineConstants:
    _PIPELINE_COMPLETE_STATES = ["SUCCEEDED", "FAILED", "CANCELLED"]
    _VALID_NAME_PATTERN = re.compile(r"^[a-z0-9-]{1,128}$")
    _VALID_AR_URL = re.compile(r"^https://.*pkg\.dev/.*$")
    _VALID_HTTPS_URL = re.compile(r"^https://.*\.(json|yaml|yml)$")

pipeline_constants = DummyPipelineConstants

# Patch base for tests
class DummyBase:
    class VertexAiStatefulResource:
        def __init__(self, project=None, location=None, credentials=None):
            self.project = project
            self.location = location
            self.credentials = credentials

base = DummyBase

# Patch experiment_resources for tests
class DummyExperimentResources:
    class _ExperimentLoggable:
        pass
    class _ExperimentLoggableSchema:
        def __init__(self, title):
            self.title = title

experiment_resources = DummyExperimentResources

# Patch metadata_constants for tests
class DummyMetadataConstants:
    SYSTEM_PIPELINE_RUN = "system-pipeline-run"

metadata_constants = DummyMetadataConstants

# Patch retry for tests
class DummyRetry:
    class Retry:
        pass

retry = DummyRetry
from aiplatform.pipeline_jobs import PipelineJob

# Unit tests for PipelineJob.get

# 1. Basic Test Cases







def test_get_with_non_string_resource_name():
    """Test get with non-string resource name (should raise)."""
    with pytest.raises(Exception):
        PipelineJob.get(12345) # 5.44ms -> 2.23ms (143% faster)

def test_get_with_resource_name_missing_job_id():
    """Test get with resource name missing job id (should raise)."""
    resource_name = "projects/test-project/locations/us-central1/pipelineJobs/"
    with pytest.raises(Exception):
        PipelineJob.get(resource_name) # 5.91ms -> 2.88ms (105% faster)

def test_get_with_resource_name_with_special_characters():
    """Test get with resource name containing special characters (should raise)."""
    resource_name = "projects/test-project/locations/us-central1/pipelineJobs/test@job"
    with pytest.raises(Exception):
        PipelineJob.get(resource_name) # 5.69ms -> 2.84ms (100% faster)

# 3. Large Scale Test Cases







def test_get_with_job_id_with_uppercase_should_fail():
    """Test get with job id containing uppercase (should fail regex)."""
    resource_name = "projects/test-project/locations/us-central1/pipelineJobs/JOBUPPER"
    with pytest.raises(Exception):
        PipelineJob.get(resource_name) # 6.36ms -> 5.59ms (13.8% faster)

def test_get_with_job_id_with_space_should_fail():
    """Test get with job id containing space (should fail regex)."""
    resource_name = "projects/test-project/locations/us-central1/pipelineJobs/job with space"
    with pytest.raises(Exception):
        PipelineJob.get(resource_name) # 5.88ms -> 5.55ms (5.85% faster)
# 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-PipelineJob.get-mglgcvym and push.