Skip to content

Commit

Permalink
Don't create a new k8s configuration object every time a model is pro…
Browse files Browse the repository at this point in the history
…gramatically created (#24635)

Speedscope results suggest that we are spending excessive amounts of
time just acquiring a logging lock that we don't use here. Pass in a
configuration object to prevent that from happening.

Test Plan: GK

Changelog
- [x] `NEW` _(added new feature or capability)_
  • Loading branch information
gibsondan authored Sep 26, 2024
1 parent 14036ea commit e49e025
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 1 deletion.
14 changes: 13 additions & 1 deletion python_modules/libraries/dagster-k8s/dagster_k8s/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,11 +7,20 @@
import kubernetes.client.models
from dagster._vendored.dateutil.parser import parse
from kubernetes.client.api_client import ApiClient
from kubernetes.client.configuration import Configuration

# Unclear what the correct type is to use for a bound here.
T_KubernetesModel = TypeVar("T_KubernetesModel")


# Create a single Configuration object to pass through to each model creation -
# the default otherwise in the OpenAPI version currently in use by the k8s
# client will create one on each model creation otherwise, which can cause
# lock contention since it acquires the global python logger lock
# see: https://github.com/kubernetes-client/python/issues/1921
shared_k8s_model_configuration = Configuration()


def _get_k8s_class(classname: str) -> Type[Any]:
if classname in ApiClient.NATIVE_TYPES_MAPPING:
return ApiClient.NATIVE_TYPES_MAPPING[classname]
Expand Down Expand Up @@ -149,7 +158,10 @@ def k8s_model_from_dict(
if len(invalid_keys):
raise Exception(f"Unexpected keys in model class {model_class.__name__}: {invalid_keys}")

kwargs = {}
# Pass through the configuration object since the default implementation creates a new one
# in the constructor, which can create lock contention if multiple threads are calling this
# simultaneously
kwargs = {"local_vars_configuration": shared_k8s_model_configuration}
for attr, attr_type in model_class.openapi_types.items(): # type: ignore
# e.g. config_map => configMap
if attr in model_dict:
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import re
from unittest import mock

import kubernetes
import pytest
Expand Down Expand Up @@ -134,3 +135,18 @@ def test_snake_case_extra_key():
}
with pytest.raises(Exception, match="Unexpected keys in model class V1Volume: {'extraKey'}"):
k8s_snake_case_dict(kubernetes.client.V1Volume, volume_dict)


def test_logger_calls_bounded():
with mock.patch("logging.Logger.setLevel") as mock_set_level:
volume_dict = {
"name": "my_volume",
"csi": {
"driver": "my_driver",
"volume_attributes": {"foo_key": "foo_val", "bar_key": "bar_val"},
},
}
k8s_model_from_dict(kubernetes.client.V1Volume, volume_dict)

# Creating a model should not use the logger lock
assert mock_set_level.call_count == 0

0 comments on commit e49e025

Please sign in to comment.