Skip to content

Commit d50b88f

Browse files
committed
change: remove usage of pkg_resources
pkg_resources is a notorious drag on import time because it does complex initialization and IO on import. Fortunately, there are modern replacements for its functionality, namely importlib_resources and packaging.
1 parent 26acec7 commit d50b88f

File tree

5 files changed

+14
-11
lines changed

5 files changed

+14
-11
lines changed

setup.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,8 @@ def read_version():
4141
"protobuf3-to-dict>=0.1.5",
4242
"requests>=2.20.0, <3",
4343
"smdebug-rulesconfig==0.1.2",
44+
"importlib-metadata>=1.4.0",
45+
"packaging>=20.0",
4446
]
4547

4648
# Specific use case dependencies

src/sagemaker/__init__.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
"""Placeholder docstring"""
1414
from __future__ import absolute_import
1515

16-
import pkg_resources
16+
import importlib_metadata
1717

1818
from sagemaker import estimator, parameter, tuner # noqa: F401
1919
from sagemaker.amazon.kmeans import KMeans, KMeansModel, KMeansPredictor # noqa: F401
@@ -60,4 +60,4 @@
6060
from sagemaker.automl.automl import AutoML, AutoMLJob, AutoMLInput # noqa: F401
6161
from sagemaker.automl.candidate_estimator import CandidateEstimator, CandidateStep # noqa: F401
6262

63-
__version__ = pkg_resources.require("sagemaker")[0].version
63+
__version__ = importlib_metadata.version("sagemaker")

src/sagemaker/mxnet/model.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,8 @@
1515

1616
import logging
1717

18-
from pkg_resources import parse_version
18+
import packaging.version
19+
1920
from sagemaker import fw_utils
2021

2122
import sagemaker
@@ -137,9 +138,9 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
137138
dict[str, str]: A container definition object usable with the
138139
CreateModel API.
139140
"""
140-
is_mms_version = parse_version(self.framework_version) >= parse_version(
141-
self._LOWEST_MMS_VERSION
142-
)
141+
is_mms_version = packaging.version.Version(
142+
self.framework_version
143+
) >= packaging.version.Version(self._LOWEST_MMS_VERSION)
143144

144145
deploy_image = self.image
145146
if not deploy_image:

src/sagemaker/pytorch/model.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
from __future__ import absolute_import
1515

1616
import logging
17-
import pkg_resources
17+
import packaging.version
1818
from sagemaker import fw_utils
1919

2020
import sagemaker
@@ -138,8 +138,8 @@ def prepare_container_def(self, instance_type, accelerator_type=None):
138138
dict[str, str]: A container definition object usable with the
139139
CreateModel API.
140140
"""
141-
lowest_mms_version = pkg_resources.parse_version(self._LOWEST_MMS_VERSION)
142-
framework_version = pkg_resources.parse_version(self.framework_version)
141+
lowest_mms_version = packaging.version.Version(self._LOWEST_MMS_VERSION)
142+
framework_version = packaging.version.Version(self.framework_version)
143143
is_mms_version = framework_version >= lowest_mms_version
144144

145145
deploy_image = self.image

src/sagemaker/user_agent.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,12 @@
1515

1616
import platform
1717
import sys
18-
import pkg_resources
1918

2019
import boto3
2120
import botocore
21+
import importlib_metadata
2222

23-
SDK_VERSION = pkg_resources.require("sagemaker")[0].version
23+
SDK_VERSION = importlib_metadata.version("sagemaker")
2424
OS_NAME = platform.system() or "UnresolvedOS"
2525
OS_VERSION = platform.release() or "UnresolvedOSVersion"
2626
PYTHON_VERSION = "{}.{}.{}".format(

0 commit comments

Comments
 (0)