Skip to content

Create a single module for generating image URIs #1464

@laurenyu

Description

@laurenyu

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.

Metadata

Metadata

Type

No type

Projects

No projects

Milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions