⚡️ Speed up function configure_ipv6 by 351%

[codeflash-ai]

📄 351% (3.51x) speedup for configure_ipv6 in python/sglang/srt/utils/common.py

⏱️ Runtime : 11.0 milliseconds → 2.43 milliseconds (best of 127 runs)

📝 Explanation and details

The optimization achieves a 350% speedup by adding an lru_cache decorator to the is_valid_ipv6_address function and making minor micro-optimizations in configure_ipv6.

Key optimization: The @lru_cache(maxsize=4096) decorator caches IPv6 address validation results, eliminating redundant calls to the expensive ipaddress.IPv6Address() constructor. The line profiler shows this constructor consumed 95.2% of the original runtime (41.4ms out of 43.5ms), but with caching it only takes 26.2ms out of 39.4ms in the optimized version - a dramatic reduction.

Why this works: IPv6 address parsing is computationally expensive, involving string validation, format checking, and internal address representation creation. When the same addresses are validated repeatedly (as shown in the test cases where similar addresses like [2001:db8::1], [2001:db8::2], etc. are processed), the cache eliminates this overhead entirely after the first validation.

Impact on workloads: Based on the function reference, configure_ipv6 is called during distributed training setup when --dist-init-addr uses IPv6 format. In distributed ML workloads, the same IPv6 addresses are likely validated multiple times during initialization, making this cache highly effective. The test results show 200-700% speedups for repeated similar addresses, and even first-time validations see modest improvements due to the minor micro-optimizations.

Test case patterns: The optimization is particularly effective for scenarios with repeated or similar IPv6 addresses (like the large-scale tests showing 300-400% improvements) while maintaining performance for edge cases and error conditions.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

import ipaddress

# imports
import pytest  # used for our unit tests
import torch.distributed
from sglang.srt.utils.common import configure_ipv6

# unit tests

# 1. Basic Test Cases

def test_basic_valid_ipv6():
    # Test a standard valid IPv6 address with port
    port, host = configure_ipv6("[2001:db8::1]:12345") # 8.93μs -> 2.15μs (315% faster)

def test_basic_valid_ipv6_with_leading_zeros():
    # IPv6 with leading zeros in segments
    port, host = configure_ipv6("[2001:0db8:0000:0000:0000:0000:0000:0001]:80") # 11.2μs -> 2.33μs (382% faster)

def test_basic_valid_ipv6_loopback():
    # IPv6 loopback address
    port, host = configure_ipv6("[::1]:65535") # 8.11μs -> 2.16μs (275% faster)

def test_basic_valid_ipv6_with_uppercase():
    # IPv6 with uppercase letters
    port, host = configure_ipv6("[2001:DB8:AC10:FE01::]:8080") # 10.9μs -> 2.20μs (397% faster)

def test_basic_valid_ipv6_min_port():
    # Minimum port number
    port, host = configure_ipv6("[fe80::1]:0") # 8.83μs -> 2.30μs (284% faster)

def test_basic_valid_ipv6_max_port():
    # Maximum port number
    port, host = configure_ipv6("[fe80::1]:65535") # 8.29μs -> 2.12μs (290% faster)

# 2. Edge Test Cases

def test_edge_missing_bracket():
    # Missing closing bracket
    with pytest.raises(ValueError, match="invalid IPv6 address format: missing ']'"):
        configure_ipv6("[2001:db8::1:12345") # 1.10μs -> 1.12μs (1.61% slower)

def test_edge_missing_port():
    # No port specified
    with pytest.raises(ValueError, match="a port must be specified in IPv6 address"):
        configure_ipv6("[2001:db8::1]") # 14.1μs -> 2.55μs (452% faster)

def test_edge_missing_colon_after_bracket():
    # No colon after closing bracket
    with pytest.raises(ValueError, match="expected ':' after ']'"):
        configure_ipv6("[2001:db8::1]12345") # 10.3μs -> 2.29μs (351% faster)

def test_edge_non_numeric_port():
    # Port is not a number
    with pytest.raises(ValueError, match="invalid port in IPv6 address: 'abc'"):
        configure_ipv6("[2001:db8::1]:abc") # 12.7μs -> 5.14μs (148% faster)

def test_edge_invalid_ipv6_address():
    # Address is not a valid IPv6 address
    with pytest.raises(ValueError, match=r"invalid IPv6 address: \[2001:db8::zzz\]"):
        configure_ipv6("[2001:db8::zzz]:12345") # 12.8μs -> 1.99μs (541% faster)

def test_edge_empty_string():
    # Empty string as input
    with pytest.raises(ValueError):
        configure_ipv6("") # 995ns -> 1.06μs (6.57% slower)

def test_edge_port_with_leading_zeros():
    # Port with leading zeros should be accepted as integer
    port, host = configure_ipv6("[2001:db8::1]:00080") # 11.0μs -> 2.50μs (339% faster)

def test_edge_port_out_of_range():
    # Port number > 65535 (should still parse as int, but not restricted by function)
    port, host = configure_ipv6("[2001:db8::1]:70000") # 9.57μs -> 2.31μs (315% faster)

def test_edge_multiple_colons_in_port():
    # Extra colon in port (should fail)
    with pytest.raises(ValueError, match="invalid port in IPv6 address: '80:90'"):
        configure_ipv6("[2001:db8::1]:80:90") # 11.9μs -> 4.80μs (148% faster)

def test_edge_extra_characters_after_port():
    # Extra characters after port (should fail)
    with pytest.raises(ValueError, match="invalid port in IPv6 address: '12345abc'"):
        configure_ipv6("[2001:db8::1]:12345abc") # 11.8μs -> 4.32μs (174% faster)

def test_edge_empty_brackets():
    # Empty brackets
    with pytest.raises(ValueError, match=r"invalid IPv6 address: \[\]"):
        configure_ipv6("[]:12345") # 5.92μs -> 2.38μs (149% faster)

def test_edge_ipv4_in_brackets():
    # IPv4 in brackets is not valid IPv6
    with pytest.raises(ValueError, match=r"invalid IPv6 address: \[127\.0\.0\.1\]"):
        configure_ipv6("[127.0.0.1]:12345") # 6.29μs -> 2.09μs (201% faster)

def test_edge_ipv6_with_zone_id():
    # IPv6 with zone id (should be valid)
    port, host = configure_ipv6("[fe80::1%eth0]:1234") # 12.3μs -> 2.47μs (398% faster)

def test_edge_ipv6_with_embedded_ipv4():
    # IPv6 with embedded IPv4
    port, host = configure_ipv6("[::ffff:192.0.2.128]:9999") # 19.5μs -> 2.38μs (719% faster)

def test_edge_ipv6_with_shortest_form():
    # IPv6 with shortest possible form
    port, host = configure_ipv6("[::]:1") # 7.08μs -> 2.36μs (201% faster)

# 3. Large Scale Test Cases

def test_large_scale_many_addresses():
    # Test a large number of valid addresses
    for i in range(1, 1001):  # 1000 addresses
        addr = f"[2001:db8::{i}]:{10000 + i}"
        port, host = configure_ipv6(addr) # 2.87ms -> 562μs (410% faster)

def test_large_scale_long_ipv6_address():
    # Test a very long but valid IPv6 address (max length segments)
    addr = "[ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff]:65535"
    port, host = configure_ipv6(addr) # 13.1μs -> 2.59μs (407% faster)

def test_large_scale_long_port():
    # Test a very large port number (should still convert to int)
    addr = "[2001:db8::1]:999999999"
    port, host = configure_ipv6(addr) # 8.99μs -> 2.11μs (326% faster)

def test_large_scale_all_zero_ipv6():
    # Test all-zero IPv6
    addr = "[0:0:0:0:0:0:0:0]:123"
    port, host = configure_ipv6(addr) # 9.78μs -> 2.10μs (365% faster)

def test_large_scale_all_one_ipv6():
    # Test all-one IPv6
    addr = "[ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff]:321"
    port, host = configure_ipv6(addr) # 10.7μs -> 2.26μs (373% faster)

def test_large_scale_zone_id_variety():
    # Test many addresses with different zone ids
    for i in range(1, 101):
        addr = f"[fe80::1%eth{i}]:{i}"
        port, host = configure_ipv6(addr) # 259μs -> 61.3μs (324% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from __future__ import annotations

import ipaddress

# imports
import pytest
import torch.distributed
from sglang.srt.utils.common import configure_ipv6

# unit tests

# -------------------- Basic Test Cases --------------------

def test_basic_valid_ipv6_and_port():
    # Basic: Standard IPv6 with port
    addr = "[2001:0db8:85a3:0000:0000:8a2e:0370:7334]:8080"
    port, host = configure_ipv6(addr) # 13.0μs -> 2.17μs (496% faster)

def test_basic_valid_ipv6_and_port_shortened():
    # Basic: Shortened IPv6 notation
    addr = "[2001:db8::1]:1234"
    port, host = configure_ipv6(addr) # 9.05μs -> 1.98μs (356% faster)

def test_basic_valid_ipv6_loopback():
    # Basic: Loopback IPv6 address
    addr = "[::1]:65535"
    port, host = configure_ipv6(addr) # 7.34μs -> 2.13μs (244% faster)

def test_basic_valid_ipv6_and_port_leading_zeros():
    # Basic: IPv6 with leading zeros
    addr = "[0000:0000:0000:0000:0000:0000:0000:0001]:80"
    port, host = configure_ipv6(addr) # 10.4μs -> 2.14μs (388% faster)

def test_basic_valid_ipv6_and_port_with_capital_hex():
    # Basic: IPv6 with capital hex letters
    addr = "[2001:DB8:ACAD:1::1]:443"
    port, host = configure_ipv6(addr) # 10.4μs -> 2.02μs (414% faster)

# -------------------- Edge Test Cases --------------------

def test_edge_missing_closing_bracket():
    # Edge: Missing closing bracket
    addr = "[2001:db8::1:8080"
    with pytest.raises(ValueError) as excinfo:
        configure_ipv6(addr) # 1.07μs -> 1.06μs (0.847% faster)

def test_edge_missing_port():
    # Edge: Missing port
    addr = "[2001:db8::1]"
    with pytest.raises(ValueError) as excinfo:
        configure_ipv6(addr) # 14.3μs -> 2.63μs (444% faster)

def test_edge_colon_without_port():
    # Edge: Colon after bracket but no port
    addr = "[2001:db8::1]:"
    with pytest.raises(ValueError) as excinfo:
        configure_ipv6(addr) # 11.0μs -> 2.60μs (324% faster)

def test_edge_non_integer_port():
    # Edge: Non-integer port
    addr = "[2001:db8::1]:notaport"
    with pytest.raises(ValueError) as excinfo:
        configure_ipv6(addr) # 13.1μs -> 5.17μs (153% faster)

def test_edge_invalid_ipv6_address():
    # Edge: Not a valid IPv6 address
    addr = "[notanipv6]:8080"
    with pytest.raises(ValueError) as excinfo:
        configure_ipv6(addr) # 7.94μs -> 2.46μs (222% faster)

def test_edge_ipv4_address_in_brackets():
    # Edge: IPv4 address in IPv6 brackets
    addr = "[192.168.1.1]:8080"
    with pytest.raises(ValueError) as excinfo:
        configure_ipv6(addr) # 5.75μs -> 2.12μs (170% faster)

def test_edge_extra_characters_after_port():
    # Edge: Extra characters after port
    addr = "[2001:db8::1]:8080abc"
    with pytest.raises(ValueError) as excinfo:
        configure_ipv6(addr) # 15.8μs -> 5.36μs (195% faster)

def test_edge_colon_without_bracket():
    # Edge: Colon after address, but no closing bracket
    addr = "[2001:db8::1:8080"
    with pytest.raises(ValueError) as excinfo:
        configure_ipv6(addr) # 1.04μs -> 1.07μs (2.72% slower)

def test_edge_missing_colon_after_bracket():
    # Edge: No colon after closing bracket
    addr = "[2001:db8::1]8080"
    with pytest.raises(ValueError) as excinfo:
        configure_ipv6(addr) # 11.0μs -> 2.31μs (378% faster)

def test_edge_port_zero():
    # Edge: Port is zero (valid)
    addr = "[2001:db8::1]:0"
    port, host = configure_ipv6(addr) # 10.1μs -> 2.36μs (328% faster)

def test_edge_port_negative():
    # Edge: Port is negative (should be allowed by int(), but semantically odd)
    addr = "[2001:db8::1]:-1"
    port, host = configure_ipv6(addr) # 9.56μs -> 2.27μs (320% faster)

def test_edge_port_too_large():
    # Edge: Port is above 65535 (should parse as int, but semantically odd)
    addr = "[2001:db8::1]:70000"
    port, host = configure_ipv6(addr) # 9.33μs -> 2.19μs (327% faster)

def test_edge_ipv6_with_embedded_ipv4():
    # Edge: IPv6 with embedded IPv4
    addr = "[::ffff:192.168.1.1]:12345"
    port, host = configure_ipv6(addr) # 18.9μs -> 2.38μs (694% faster)

def test_edge_ipv6_with_multiple_colons_in_port():
    # Edge: Multiple colons after closing bracket
    addr = "[2001:db8::1]::8080"
    with pytest.raises(ValueError) as excinfo:
        configure_ipv6(addr) # 12.5μs -> 4.65μs (169% faster)

def test_edge_empty_string():
    # Edge: Empty string
    addr = ""
    with pytest.raises(ValueError) as excinfo:
        configure_ipv6(addr) # 1.29μs -> 1.31μs (2.13% slower)

def test_edge_only_brackets():
    # Edge: Only brackets
    addr = "[]:8080"
    with pytest.raises(ValueError) as excinfo:
        configure_ipv6(addr) # 5.68μs -> 2.21μs (157% faster)

# -------------------- Large Scale Test Cases --------------------

def test_large_scale_many_valid_ipv6_addresses():
    # Large Scale: Test with 1000 valid unique IPv6 addresses and ports
    for i in range(1000):
        # Construct a valid IPv6 address by incrementing the last block
        addr = f"[2001:db8::" + str(i) + f"]:{10000 + i}"
        port, host = configure_ipv6(addr) # 2.87ms -> 588μs (389% faster)

def test_large_scale_many_invalid_ipv6_addresses():
    # Large Scale: Test with 1000 invalid IPv6 addresses (invalid address part)
    for i in range(1000):
        addr = f"[invalid{i}]:{10000 + i}"
        with pytest.raises(ValueError) as excinfo:
            configure_ipv6(addr)

def test_large_scale_many_missing_ports():
    # Large Scale: Test with 1000 valid IPv6 addresses but missing port
    for i in range(1000):
        addr = f"[2001:db8::{i}]"
        with pytest.raises(ValueError) as excinfo:
            configure_ipv6(addr)

def test_large_scale_ports_all_digits():
    # Large Scale: Test with all possible single digit ports (0-9)
    for port in range(10):
        addr = f"[2001:db8::1]:{port}"
        out_port, host = configure_ipv6(addr) # 39.0μs -> 8.34μs (368% faster)

def test_large_scale_ports_edge_values():
    # Large Scale: Test with edge port values (0, 65535, -1, 70000)
    for port in [0, 65535, -1, 70000]:
        addr = f"[2001:db8::1]:{port}"
        out_port, host = configure_ipv6(addr) # 20.3μs -> 4.86μs (318% 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-configure_ipv6-mijm03hj and push.