⚡️ Speed up function _get_local_rank by 10%

[codeflash-ai]

📄 10% (0.10x) speedup for _get_local_rank in src/spdl/pipeline/_profile.py

⏱️ Runtime : 3.00 milliseconds → 2.72 milliseconds (best of 19 runs)

📝 Explanation and details

The optimization replaces os.environ.get("LOCAL_RANK", "0") with a try-except block that directly accesses os.environ["LOCAL_RANK"] and catches KeyError exceptions.

Key changes:

• Direct dictionary access os.environ["LOCAL_RANK"] instead of the .get() method
• Exception handling for the missing key case instead of default parameter

Why it's faster:

1. Eliminates default value creation overhead: os.environ.get() creates a new string object "0" on every call, even when LOCAL_RANK exists. The optimized version only creates the integer 0 in the rare exception case.

2. Reduces method call overhead: Direct dictionary access [] is faster than calling the .get() method, which has additional function call and parameter processing overhead.

3. Optimizes for the common case: In distributed PyTorch environments, LOCAL_RANK is typically set, making the exception path rare. The line profiler shows only 6 out of 3052 calls hit the exception handler.

Performance characteristics:

• Best gains (15-38%) when LOCAL_RANK is unset (exception path is more efficient than string default creation)
• Consistent 10-20% improvements for normal cases when the environment variable exists
• Maintains identical behavior and error handling for invalid values
• Particularly effective for high-frequency calls in distributed training scenarios

The 10% overall speedup demonstrates the cumulative benefit of avoiding unnecessary string allocations and method call overhead in this frequently-called function.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import os

# imports
import pytest
from spdl.pipeline._profile import _get_local_rank

# unit tests

@pytest.mark.basic
def test_default_local_rank(monkeypatch):
    """
    Test that _get_local_rank returns 0 when LOCAL_RANK is not set.
    """
    # Unset LOCAL_RANK if it exists
    monkeypatch.delenv("LOCAL_RANK", raising=False)
    codeflash_output = _get_local_rank() # 1.69μs -> 1.23μs (38.2% faster)

@pytest.mark.basic
@pytest.mark.parametrize("rank_str,expected", [
    ("0", 0),
    ("1", 1),
    ("42", 42),
    ("999", 999),
])
def test_local_rank_basic_valid(monkeypatch, rank_str, expected):
    """
    Test _get_local_rank with several normal integer string values.
    """
    monkeypatch.setenv("LOCAL_RANK", rank_str)
    codeflash_output = _get_local_rank() # 5.34μs -> 4.74μs (12.7% faster)

@pytest.mark.edge
@pytest.mark.parametrize("rank_str,expected", [
    ("-1", -1),    # negative rank
    ("0003", 3),   # leading zeros
    ("+5", 5),     # explicit plus sign
])
def test_local_rank_edge_numeric(monkeypatch, rank_str, expected):
    """
    Test _get_local_rank with edge-case numeric string values.
    """
    monkeypatch.setenv("LOCAL_RANK", rank_str)
    codeflash_output = _get_local_rank() # 4.22μs -> 3.73μs (13.0% faster)

@pytest.mark.edge
@pytest.mark.parametrize("rank_str", [
    ("notanumber"),    # non-numeric string
    ("3.14"),          # float string
    ("1e2"),           # scientific notation
    (" "),             # whitespace
    ("0x10"),          # hex string
])
def test_local_rank_invalid(monkeypatch, rank_str):
    """
    Test _get_local_rank raises ValueError when LOCAL_RANK is not a valid int.
    """
    monkeypatch.setenv("LOCAL_RANK", rank_str)
    with pytest.raises(ValueError):
        _get_local_rank() # 17.4μs -> 16.5μs (5.36% faster)

@pytest.mark.edge
def test_local_rank_empty_string(monkeypatch):
    """
    Test _get_local_rank raises ValueError when LOCAL_RANK is set to empty string.
    """
    monkeypatch.setenv("LOCAL_RANK", "")
    with pytest.raises(ValueError):
        _get_local_rank() # 3.19μs -> 2.81μs (13.5% faster)

@pytest.mark.edge
def test_local_rank_whitespace(monkeypatch):
    """
    Test _get_local_rank raises ValueError when LOCAL_RANK is whitespace.
    """
    monkeypatch.setenv("LOCAL_RANK", "   ")
    with pytest.raises(ValueError):
        _get_local_rank() # 3.47μs -> 3.21μs (8.33% faster)

@pytest.mark.edge
def test_local_rank_large_integer(monkeypatch):
    """
    Test _get_local_rank with a very large integer value.
    """
    large_value = str(2**63 - 1)  # Largest signed 64-bit int
    monkeypatch.setenv("LOCAL_RANK", large_value)
    codeflash_output = _get_local_rank() # 1.54μs -> 1.30μs (18.2% faster)

@pytest.mark.edge
def test_local_rank_min_integer(monkeypatch):
    """
    Test _get_local_rank with a very small (negative) integer value.
    """
    min_value = str(-2**63)
    monkeypatch.setenv("LOCAL_RANK", min_value)
    codeflash_output = _get_local_rank() # 1.54μs -> 1.29μs (19.0% faster)

@pytest.mark.large
def test_local_rank_many_values(monkeypatch):
    """
    Test _get_local_rank with a range of valid integer values (0 to 999).
    """
    for i in range(1000):
        monkeypatch.setenv("LOCAL_RANK", str(i))
        codeflash_output = _get_local_rank() # 638μs -> 571μs (11.8% faster)

@pytest.mark.large
def test_local_rank_many_invalid(monkeypatch):
    """
    Test _get_local_rank with many invalid values to check robustness.
    """
    invalid_values = ["", "a", "!", "1.1", "NaN", "None", "[]", "{}", "0x1", " "]
    for val in invalid_values:
        monkeypatch.setenv("LOCAL_RANK", val)
        with pytest.raises(ValueError):
            _get_local_rank()

@pytest.mark.large
def test_local_rank_stress_env(monkeypatch):
    """
    Test _get_local_rank under rapid changes to the LOCAL_RANK environment variable.
    """
    # Alternate between valid and invalid values
    for i in range(500):
        monkeypatch.setenv("LOCAL_RANK", str(i))
        codeflash_output = _get_local_rank() # 336μs -> 307μs (9.52% faster)
        monkeypatch.setenv("LOCAL_RANK", "invalid")
        with pytest.raises(ValueError):
            _get_local_rank()
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import os

# imports
import pytest
from spdl.pipeline._profile import _get_local_rank

# unit tests

@pytest.mark.parametrize(
    "env_val,expected",
    [
        # Basic: LOCAL_RANK not set, should default to 0
        (None, 0),
        # Basic: LOCAL_RANK set to 0
        ("0", 0),
        # Basic: LOCAL_RANK set to 1
        ("1", 1),
        # Basic: LOCAL_RANK set to 42
        ("42", 42),
        # Basic: LOCAL_RANK set to a negative integer
        ("-3", -3),
        # Edge: LOCAL_RANK set to a large positive integer
        (str(2**31 - 1), 2**31 - 1),
        # Edge: LOCAL_RANK set to a large negative integer
        (str(-(2**31)), -(2**31)),
        # Edge: LOCAL_RANK set to a string with leading/trailing spaces
        (" 7 ", 7),
        # Edge: LOCAL_RANK set to a string with plus sign
        ("+8", 8),
        # Edge: LOCAL_RANK set to zero with leading zeros
        ("0000123", 123),
    ]
)
def test_get_local_rank_valid(monkeypatch, env_val, expected):
    """
    Test _get_local_rank with valid LOCAL_RANK values and unset variable.
    """
    # Setup: Set or unset the environment variable
    if env_val is not None:
        monkeypatch.setenv("LOCAL_RANK", env_val)
    else:
        monkeypatch.delenv("LOCAL_RANK", raising=False)
    # Assert: Function returns expected integer
    codeflash_output = _get_local_rank() # 14.0μs -> 12.0μs (16.5% faster)

@pytest.mark.parametrize(
    "env_val",
    [
        # Edge: LOCAL_RANK set to a non-integer string
        "abc",
        # Edge: LOCAL_RANK set to an empty string
        "",
        # Edge: LOCAL_RANK set to a float string
        "3.14",
        # Edge: LOCAL_RANK set to a string with spaces and non-numeric
        "  5x ",
        # Edge: LOCAL_RANK set to a hexadecimal string
        "0x10",
        # Edge: LOCAL_RANK set to a binary string
        "0b101",
        # Edge: LOCAL_RANK set to a string with only spaces
        "   ",
        # Edge: LOCAL_RANK set to a special character
        "@",
    ]
)
def test_get_local_rank_invalid(monkeypatch, env_val):
    """
    Test _get_local_rank raises ValueError for invalid LOCAL_RANK values.
    """
    monkeypatch.setenv("LOCAL_RANK", env_val)
    with pytest.raises(ValueError):
        _get_local_rank() # 25.7μs -> 24.0μs (6.92% faster)

def test_get_local_rank_env_var_case_sensitive(monkeypatch):
    """
    Edge: Ensure the function does not pick up 'local_rank' or 'Local_Rank' (case sensitivity).
    """
    monkeypatch.delenv("LOCAL_RANK", raising=False)
    monkeypatch.setenv("local_rank", "99")
    monkeypatch.setenv("Local_Rank", "100")
    codeflash_output = _get_local_rank() # 1.43μs -> 1.12μs (27.5% faster)

def test_get_local_rank_env_var_other_vars(monkeypatch):
    """
    Edge: Ensure other unrelated env vars do not affect the result.
    """
    monkeypatch.delenv("LOCAL_RANK", raising=False)
    monkeypatch.setenv("RANK", "123")
    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1")
    codeflash_output = _get_local_rank() # 1.42μs -> 1.10μs (29.5% faster)

def test_get_local_rank_env_var_removed(monkeypatch):
    """
    Edge: LOCAL_RANK is set, then removed; should default to 0.
    """
    monkeypatch.setenv("LOCAL_RANK", "23")
    codeflash_output = _get_local_rank() # 1.29μs -> 1.09μs (17.8% faster)
    monkeypatch.delenv("LOCAL_RANK", raising=True)
    codeflash_output = _get_local_rank() # 1.10μs -> 920ns (19.1% faster)

@pytest.mark.parametrize(
    "val",
    [str(i) for i in range(1000)]  # Large scale: 0 to 999
)
def test_get_local_rank_large_scale(monkeypatch, val):
    """
    Large Scale: Test _get_local_rank with LOCAL_RANK set to many different valid values.
    """
    monkeypatch.setenv("LOCAL_RANK", val)
    codeflash_output = _get_local_rank() # 1.29ms -> 1.13ms (13.4% faster)

def test_get_local_rank_large_negative(monkeypatch):
    """
    Large Scale: Test with a large negative value.
    """
    val = str(-(10**9))
    monkeypatch.setenv("LOCAL_RANK", val)
    codeflash_output = _get_local_rank() # 1.57μs -> 1.36μs (15.9% faster)

def test_get_local_rank_large_positive(monkeypatch):
    """
    Large Scale: Test with a large positive value.
    """
    val = str(10**9)
    monkeypatch.setenv("LOCAL_RANK", val)
    codeflash_output = _get_local_rank() # 1.44μs -> 1.25μs (15.7% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from spdl.pipeline._profile import _get_local_rank

def test__get_local_rank():
    _get_local_rank()

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_uafn4wd5/tmpc0azatp6/test_concolic_coverage.py::test__get_local_rank	2.42μs	1.93μs	24.9%✅

To edit these changes git checkout codeflash/optimize-_get_local_rank-mgrea4ys and push.