-
Notifications
You must be signed in to change notification settings - Fork 1.2k
Description
Is your feature request related to a problem? Please describe.
The code for determining an image URI for a given algorithm or framework is scattered across the Python SDK. In addition, the documentation on it is confusing. We used to have our own page for URIs, which now simply points to the DLC page, and that can cause its own confusion because there’s no longer documentation for older images or the SageMaker-style tags (e.g. “2.1.0-gpu-py3” as opposed to “2.1.0-gpu-py36-cu101-ubuntu18.04”).
Describe the solution you'd like
I’d like to create a new class, ImageURIProvider
(open to naming suggestions), that houses all of our logic for creating an image URI. This would move create_image_uri
(and associated methods/constants) out of fw_utils.py
and get_image_uri
(and associated methods/constants) out of the Amazon algorithm section. Similar logic from SKLearn, XGBoost, fw_registry.py
, etc. would also be moved to this new class.
Proposed interface:
class ImageURIProvider(object):
@staticmethod
def framework_image_uri(
framework,
framework_version,
accelerator_type=None,
optimized_families=None,
region=None,
instance_type=None,
py_version=None,
account=None,
):
@staticmethod
def algorithm_image_uri(algorithm, algorithm_version=1):
"""Might need a better name to avoid confusion with
AlgorithmEstimator (Marketplace)
"""
The methods are generic, as opposed to per framework/algorithm (e.g. tensorflow_image_uri
, xgboost_image_uri
, etc.), but we can potentially define an enum to help validate the framework/algorithm name. There is also potential here to combine the two methods, and simply just have one image_uri
method for everything.
This will be useful for Amazon algorithms that don't have a dedicated estimator (currently a use case of get_image_uri
in many notebooks), getting a framework image URI to use as a base image or with boto3, and our own internal code cleanliness.