⚡️ Speed up function get_default_conda_env by 43%

[codeflash-ai]

📄 43% (0.43x) speedup for get_default_conda_env in mlflow/diviner/__init__.py

⏱️ Runtime : 14.8 milliseconds → 10.3 milliseconds (best of 83 runs)

📝 Explanation and details

The optimization replaces the repeated YAML parsing bottleneck with a more efficient copy operation.

What was optimized:

• Pre-parsed YAML template: The _conda_header string is now parsed once at module import time into _conda_header_dict, eliminating repeated yaml.safe_load() calls
• Deep copy instead of YAML parsing: Each function call now uses copy.deepcopy(_conda_header_dict) instead of yaml.safe_load(_conda_header)

Why this is faster:
The line profiler shows the dramatic impact - the original yaml.safe_load(_conda_header) consumed 79.8% of function runtime (15.8ms out of 19.8ms total), while the optimized copy.deepcopy(_conda_header_dict) only takes 10.2% (0.4ms out of 4.0ms total). This represents a ~38x speedup for the environment creation step specifically.

YAML parsing involves lexical analysis, syntax tree building, and object construction from text, while copy.deepcopy() performs a straightforward recursive object duplication. Since _conda_header is a static string that never changes, parsing it repeatedly is pure waste.

Impact on workloads:
Based on the test results showing consistent 40-49% speedups across all test cases, this optimization particularly benefits scenarios that create multiple conda environments, such as:

• Batch model deployments
• MLflow model logging workflows
• Environment template generation for multiple projects

The optimization preserves all functionality while significantly reducing the computational overhead of environment dictionary creation.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import importlib
import sys
import types

imports

import pytest
import yaml
from mlflow.diviner.init import get_default_conda_env

--- Mocked dependencies for test isolation ---

Simulate mlflow.utils.PYTHON_VERSION and mlflow.version.VERSION

PYTHON_VERSION = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
MLFLOW_VERSION = "9.9.9" # Arbitrary version for test stability

Simulate importlib.metadata.version for pip and diviner

class DummyMetadata:
_versions = {
"pip": "23.1.2",
"diviner": "1.2.3",
"mlflow": MLFLOW_VERSION,
}
@staticmethod
def version(pkg):
if pkg in DummyMetadata._versions:
return DummyMetadata._versions[pkg]
raise Exception("PackageNotFoundError")

importlib.metadata = types.SimpleNamespace(version=DummyMetadata.version, PackageNotFoundError=Exception)
from mlflow.diviner.init import get_default_conda_env

--- Unit tests for get_default_conda_env ---

1. Basic Test Cases

def test_default_env_structure():
"""Test that the default conda env has the correct structure and keys."""
codeflash_output = get_default_conda_env(); env = codeflash_output # 755μs -> 522μs (44.5% faster)

def test_default_env_python_and_pip_versions():
"""Test that the default env pins python and pip versions correctly."""
codeflash_output = get_default_conda_env(); env = codeflash_output # 745μs -> 523μs (42.5% faster)
deps = env["dependencies"]

def test_default_env_contains_diviner_in_pip():
"""Test that diviner is present in pip dependencies with correct version."""
codeflash_output = get_default_conda_env(); env = codeflash_output # 733μs -> 522μs (40.2% faster)
pip_deps = None
for dep in env["dependencies"]:
if isinstance(dep, dict) and "pip" in dep:
pip_deps = dep["pip"]

def test_default_env_contains_mlflow_in_pip():
"""Test that mlflow is present in pip dependencies with correct version."""
codeflash_output = get_default_conda_env(); env = codeflash_output # 729μs -> 511μs (42.7% faster)
pip_deps = None
for dep in env["dependencies"]:
if isinstance(dep, dict) and "pip" in dep:
pip_deps = dep["pip"]

def test_default_env_no_extra_channels():
"""Test that no extra conda channels are present by default."""
codeflash_output = get_default_conda_env(); env = codeflash_output # 767μs -> 524μs (46.4% faster)

2. Edge Test Cases

def test_env_with_non_str_deps_raises():
"""Test that non-string dependencies raise an error."""
with pytest.raises(Exception):
_mlflow_conda_env(additional_conda_deps=[123])

#------------------------------------------------
import sys

imports

import pytest # used for our unit tests
from mlflow.diviner.init import get_default_conda_env

Helper function to extract pip requirements from the environment dict

def extract_pip_requirements(env):
for dep in env["dependencies"]:
if isinstance(dep, dict) and "pip" in dep:
return dep["pip"]
return []

Helper function to extract conda requirements from the environment dict

def extract_conda_requirements(env):
conda_reqs = []
for dep in env["dependencies"]:
if isinstance(dep, str) and not dep.startswith("python=") and not dep.startswith("pip"):
conda_reqs.append(dep)
elif isinstance(dep, str) and dep.startswith("pip"):
conda_reqs.append(dep)
return conda_reqs

Helper function to extract python version from the environment dict

def extract_python_version(env):
for dep in env["dependencies"]:
if isinstance(dep, str) and dep.startswith("python="):
return dep.split("=")[1]
return None

------------------ UNIT TESTS ------------------

1. Basic Test Cases

def test_env_structure_basic():
"""
Test that the returned environment is a dictionary with required keys.
"""
codeflash_output = get_default_conda_env(); env = codeflash_output # 771μs -> 534μs (44.3% faster)

def test_env_name_and_channels_basic():
"""
Test that the environment name and channels are set correctly.
"""
codeflash_output = get_default_conda_env(); env = codeflash_output # 723μs -> 512μs (41.2% faster)

def test_python_version_basic():
"""
Test that the Python version matches the system's version.
"""
codeflash_output = get_default_conda_env(); env = codeflash_output # 735μs -> 502μs (46.2% faster)
py_version = extract_python_version(env)
sys_version = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"

def test_pip_and_conda_requirements_basic():
"""
Test that pip and conda requirements include expected packages.
"""
codeflash_output = get_default_conda_env(); env = codeflash_output # 737μs -> 523μs (40.8% faster)
pip_reqs = extract_pip_requirements(env)
conda_reqs = extract_conda_requirements(env)

def test_pip_version_format_basic():
"""
Test that pip version is specified with <= in conda requirements.
"""
codeflash_output = get_default_conda_env(); env = codeflash_output # 750μs -> 522μs (43.6% faster)
conda_reqs = extract_conda_requirements(env)
pip_versions = [req for req in conda_reqs if req.startswith("pip")]

2. Edge Test Cases

def test_no_duplicate_pip_requirements_edge():
"""
Test that pip requirements do not contain duplicates.
"""
codeflash_output = get_default_conda_env(); env = codeflash_output # 731μs -> 524μs (39.5% faster)
pip_reqs = extract_pip_requirements(env)

def test_missing_pip_key_edge():
"""
Simulate an environment missing the pip key and check extraction.
"""
codeflash_output = get_default_conda_env(); env = codeflash_output # 747μs -> 510μs (46.4% faster)
# Remove the pip dict
env["dependencies"] = [dep for dep in env["dependencies"] if not (isinstance(dep, dict) and "pip" in dep)]
pip_reqs = extract_pip_requirements(env)

def test_invalid_python_version_edge():
"""
Simulate an environment with an invalid python version string.
"""
codeflash_output = get_default_conda_env(); env = codeflash_output # 729μs -> 519μs (40.5% faster)
# Replace python version with an invalid one
env["dependencies"][0] = "python=invalid_version"
py_version = extract_python_version(env)

def test_empty_dependencies_edge():
"""
Simulate an environment with empty dependencies.
"""
codeflash_output = get_default_conda_env(); env = codeflash_output # 737μs -> 505μs (46.0% faster)
env["dependencies"] = []

def test_missing_channels_edge():
"""
Simulate an environment missing channels key.
"""
codeflash_output = get_default_conda_env(); env = codeflash_output # 749μs -> 516μs (45.0% faster)
del env["channels"]

def test_extra_conda_channels_edge():
"""
Simulate adding extra conda channels.
"""
codeflash_output = get_default_conda_env(); env = codeflash_output # 745μs -> 519μs (43.4% faster)
env["channels"].append("custom-channel")

3. Large Scale Test Cases

def test_large_number_of_pip_requirements_large():
"""
Test environment with a large number of pip requirements (up to 1000).
"""
codeflash_output = get_default_conda_env(); env = codeflash_output # 747μs -> 522μs (43.1% faster)
# Add 1000 fake pip requirements
large_pip_reqs = [f"package{i}==1.0.0" for i in range(1000)]
for dep in env["dependencies"]:
if isinstance(dep, dict) and "pip" in dep:
dep["pip"].extend(large_pip_reqs)
pip_reqs = extract_pip_requirements(env)

def test_large_number_of_conda_requirements_large():
"""
Test environment with a large number of conda requirements (up to 1000).
"""
codeflash_output = get_default_conda_env(); env = codeflash_output # 719μs -> 507μs (41.8% faster)
# Add 1000 fake conda requirements
large_conda_reqs = [f"condapkg{i}=1.0.0" for i in range(1000)]
env["dependencies"][1:1] = large_conda_reqs # Insert after python version
conda_reqs = extract_conda_requirements(env)

def test_large_scale_env_structure_large():
"""
Test that environment structure remains valid with large numbers of dependencies.
"""
codeflash_output = get_default_conda_env(); env = codeflash_output # 733μs -> 518μs (41.4% faster)
# Add many pip and conda requirements
env["dependencies"].extend([f"condapkg{i}=1.0.0" for i in range(500)])
for dep in env["dependencies"]:
if isinstance(dep, dict) and "pip" in dep:
dep["pip"].extend([f"package{i}==1.0.0" for i in range(500)])

def test_performance_large_scale():
"""
Test that extraction functions run efficiently with large environments.
"""
codeflash_output = get_default_conda_env(); env = codeflash_output # 752μs -> 505μs (49.1% faster)
env["dependencies"].extend([f"condapkg{i}=1.0.0" for i in range(900)])
for dep in env["dependencies"]:
if isinstance(dep, dict) and "pip" in dep:
dep["pip"].extend([f"package{i}==1.0.0" for i in range(900)])
# Extraction should not take excessive time
pip_reqs = extract_pip_requirements(env)
conda_reqs = extract_conda_requirements(env)

codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

#------------------------------------------------
from mlflow.diviner.init import get_default_conda_env

def test_get_default_conda_env():
get_default_conda_env()

To edit these changes git checkout codeflash/optimize-get_default_conda_env-mhumci30 and push.