feature: add framework upgrade tool (#1762)

* feature: add framework upgrade tool

* add framework upgrade test

* add algo image upgrade logic

* support for region update

* move framework upgrade tool to cli folder

* fix typo

* fix auto changes

* fix auto changes

* fix auto changes

* correct typos and simplify methods

* refactor method and update docstring

* simplify and rename method

* fix dict KeyError and and unit tests

Co-authored-by: Chuyang Deng <[email protected]>
Co-authored-by: Lauren Yu <[email protected]>

Author: 3 people

src/sagemaker/cli/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying athis file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Tools for automating code updates"""
+from __future__ import absolute_import

src/sagemaker/cli/framework_upgrade.py

@@ -0,0 +1,248 @@
+# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""A Python script to upgrade framework versions and regions"""
+from __future__ import absolute_import
+
+import argparse
+import json
+import os
+
+from sagemaker.image_uris import config_for_framework
+
+IMAGE_URI_CONFIG_DIR = os.path.join("..", "image_uri_config")
+
+
+def get_latest_values(existing_content, scope=None):
+    """Get the latest "registries", "py_versions" and "repository" values
+
+    Args:
+        existing_content (dict): Dictionary of complete framework image information.
+        scope (str): Type of the image, required if the target is DLC
+            framework (Default: None).
+    """
+    if scope in existing_content:
+        existing_content = existing_content[scope]
+    else:
+        if "versions" not in existing_content:
+            raise ValueError(
+                "Invalid image scope: {}. Valid options: {}.".format(
+                    scope, ", ".join(existing_content.key())
+                )
+            )
+
+    latest_version = list(existing_content["versions"].keys())[-1]
+    registries = existing_content["versions"][latest_version]["registries"]
+    py_versions = existing_content["versions"][latest_version]["py_versions"]
+    repository = existing_content["versions"][latest_version]["repository"]
+
+    return registries, py_versions, repository
+
+
+def _write_dict_to_json(filename, existing_content):
+    """Write a Python dictionary to a json file.
+
+    Args:
+        filename (str): Name of the target json file.
+        existing_content (dict): Dictionary to be written to the json file.
+    """
+    with open(filename, "w") as f:
+        json.dump(existing_content, f, sort_keys=True, indent=4)
+
+
+def add_dlc_framework_version(
+    existing_content,
+    short_version,
+    full_version,
+    scope,
+    processors,
+    py_versions,
+    registries,
+    repository,
+):
+    """Update DLC framework image uri json file with new version information.
+
+    Args:
+        existing_content (dict): Existing framework image uri information read from
+            "<framework>.json" file.
+        framework (str): Framework name (e.g. tensorflow, pytorch, mxnet)
+        short_version (str): Abbreviated framework version (e.g. 1.0, 1.5)
+        full_version (str): Complete framework version (e.g. 1.0.0, 1.5.2)
+        scope (str): Framework image type, it could be "training", "inference"
+            or "eia"
+        processors (list): Supported processors (e.g. ["cpu", "gpu"])
+        py_versions (list): Supported Python versions (e.g. ["py3", "py37"])
+        registries (dict): Framework image's region to account mapping.
+        repository (str): Framework image's ECR repository.
+    """
+    for processor in processors:
+        if processor not in existing_content[scope]["processors"]:
+            existing_content[scope]["processors"].append(processor)
+    existing_content[scope]["version_aliases"][short_version] = full_version
+
+    new_version = {
+        "registries": registries,
+        "repository": repository,
+        "py_versions": py_versions,
+    }
+    existing_content[scope]["versions"][full_version] = new_version
+
+
+def add_algo_version(
+    existing_content,
+    processors,
+    scopes,
+    full_version,
+    py_versions,
+    registries,
+    repository,
+    tag_prefix,
+):
+    """Update Algorithm image uri json file with new version information.
+
+    Args:
+        existing_content (dict): Existing algorithm image uri information read from
+            "<algorithm>.json" file.
+        processors (list): Supported processors (e.g. ["cpu", "gpu"])
+        scopes (list): Framework image type, it could be "training", "inference
+        full_version (str): Complete framework version (e.g. 1.0.0, 1.5.2)
+        py_versions (list): Supported Python versions (e.g. ["py3", "py37"])
+        registries (dict): Algorithm image's region to account mapping.
+        repository (str): Algorithm's corresponding repository name.
+        tag_prefix (str): Algorithm image's tag prefix.
+    """
+    for processor in processors:
+        if processor not in existing_content["processors"]:
+            existing_content["processors"].append(processor)
+    for scope in scopes:
+        if scope not in existing_content["scope"]:
+            existing_content["scope"].append(scope)
+
+    new_version = {
+        "py_versions": py_versions,
+        "registries": registries,
+        "repository": repository,
+    }
+    if tag_prefix:
+        new_version["tag_prefix"] = tag_prefix
+    existing_content["versions"][full_version] = new_version
+
+
+def add_region(existing_content, region, account):
+    """Add region account to framework/algorithm registries.
+
+    Args:
+        existing_content (dict): Existing framework/algorithm image uri information read from
+            json file.
+        region (str): New region to be added to framework/algorithm registries (e.g. us-west-2).
+        account (str): Region registry account number.
+    """
+    if "scope" not in existing_content:
+        for scope in existing_content:
+            for version in existing_content[scope]["versions"]:
+                existing_content[scope]["versions"][version]["registries"][region] = account
+    else:
+        for version in existing_content["versions"]:
+            existing_content["versions"][version]["registries"][region] = account
+
+
+def add_version(
+    existing_content, short_version, full_version, scope, processors, py_versions, tag_prefix,
+):
+    """Read framework image uri information from json file to a dictionary, update it with new
+    framework version information, then write the dictionary back to json file.
+
+    Args:
+        existing_content (dict): Existing framework/algorithm image uri information read from
+            json file.
+        short_version (str): Abbreviated framework version (e.g. 1.0, 1.5).
+        full_version (str): Complete framework version (e.g. 1.0.0, 1.5.2).
+        scope (str): Framework image type, it could be "training", "inference".
+        processors (str): Supported processors (e.g. "cpu,gpu").
+        py_versions (str): Supported Python versions (e.g. "py3,py37").
+        tag_prefix (str): Algorithm image's tag prefix.
+     """
+    py_versions = py_versions.split(",")
+    processors = processors.split(",")
+    latest_registries, latest_py_versions, latest_repository = get_latest_values(
+        existing_content, scope
+    )
+    if not py_versions:
+        py_versions = latest_py_versions
+
+    if scope in existing_content:
+        add_dlc_framework_version(
+            existing_content,
+            short_version,
+            full_version,
+            scope,
+            processors,
+            py_versions,
+            latest_registries,
+            latest_repository,
+        )
+    else:
+        scopes = scope.split(",")
+        add_algo_version(
+            existing_content,
+            processors,
+            scopes,
+            full_version,
+            py_versions,
+            latest_registries,
+            latest_repository,
+            tag_prefix,
+        )
+
+
+def main():
+    """Parse command line arguments, call corresponding methods."""
+    parser = argparse.ArgumentParser(description="Framework upgrade tool.")
+    parser.add_argument(
+        "--framework", required=True, help="Name of the framework (e.g. tensorflow, mxnet, etc.)"
+    )
+    parser.add_argument("--short-version", help="Abbreviated framework version (e.g. 2.0)")
+    parser.add_argument("--full-version", help="Full framework version (e.g. 2.0.1)")
+    parser.add_argument("--processors", help="Suppoted processors (e.g. cpu, gpu)")
+    parser.add_argument("--py-versions", help="Supported Python versions (e.g. py3,py37)")
+    parser.add_argument("--scope", help="Scope for the Algorithm image (e.g. inference, training)")
+    parser.add_argument(
+        "--tag-prefix", help="Tag prefix of the Algorithm image (e.g. ray-0.8.5-torch)"
+    )
+    parser.add_argument("--region", help="New region to be added (e.g. us-west-2)")
+    parser.add_argument("--account", help="Registry account of new region")
+
+    args = parser.parse_args()
+
+    content = config_for_framework(args.framework)
+
+    if args.region or args.account:
+        if args.region and not args.account or args.account and not args.region:
+            raise ValueError("--region and --account must be used together to expand region.")
+        add_region(content, args.region, args.account)
+    else:
+        add_version(
+            content,
+            args.short_version,
+            args.full_version,
+            args.scope,
+            args.processors,
+            args.py_versions,
+            args.tag_prefix,
+        )
+
+    file = os.path.join(IMAGE_URI_CONFIG_DIR, "{}.json".format(args.framework))
+    _write_dict_to_json(file, content)
+
+
+if __name__ == "__main__":
+    main()