@@ -36,10 +36,10 @@ def retrieve(
3636 """Retrieves the ECR URI for the Docker image matching the given arguments.
3737
3838 Args:
39- framework (str): The name of the framework.
39+ framework (str): The name of the framework or algorithm .
4040 region (str): The AWS region.
41- version (str): The framework version. This is required if there is
42- more than one supported version for the given framework.
41+ version (str): The framework or algorithm version. This is required if there is
42+ more than one supported version for the given framework or algorithm .
4343 py_version (str): The Python version. This is required if there is
4444 more than one supported Python version for the given framework version.
4545 instance_type (str): The SageMaker instance type. For supported types, see
@@ -58,7 +58,9 @@ def retrieve(
5858 ValueError: If the combination of arguments specified is not supported.
5959 """
6060 config = _config_for_framework_and_scope (framework , image_scope , accelerator_type )
61- version_config = config ["versions" ][_version_for_config (version , config , framework )]
61+
62+ version = _validate_version_and_set_if_needed (version , config , framework )
63+ version_config = config ["versions" ][_version_for_config (version , config )]
6264
6365 py_version = _validate_py_version_and_set_if_needed (py_version , version_config )
6466 version_config = version_config .get (py_version ) or version_config
@@ -67,7 +69,7 @@ def retrieve(
6769 hostname = utils ._botocore_resolver ().construct_endpoint ("ecr" , region )["hostname" ]
6870
6971 repo = version_config ["repository" ]
70- tag = _format_tag (version , _processor (instance_type , config [ "processors" ] ), py_version )
72+ tag = _format_tag (version , _processor (instance_type , config . get ( "processors" ) ), py_version )
7173
7274 return ECR_URI_TEMPLATE .format (registry = registry , hostname = hostname , repository = repo , tag = tag )
7375
@@ -94,13 +96,33 @@ def config_for_framework(framework):
9496 return json .load (f )
9597
9698
97- def _version_for_config (version , config , framework ):
99+ def _validate_version_and_set_if_needed (version , config , framework ):
100+ """Checks if the framework/algorithm version is one of the supported versions."""
101+ available_versions = list (config ["versions" ].keys ())
102+
103+ if len (available_versions ) == 1 :
104+ log_message = "Defaulting to the only supported framework/algorithm version: {}." .format (
105+ available_versions [0 ]
106+ )
107+ if version and version != available_versions [0 ]:
108+ logger .warning ("%s Ignoring framework/algorithm version: %s." , log_message , version )
109+ elif not version :
110+ logger .info (log_message )
111+
112+ return available_versions [0 ]
113+
114+ available_versions += list (config .get ("version_aliases" , {}).keys ())
115+ _validate_arg ("{} version" .format (framework ), version , available_versions )
116+
117+ return version
118+
119+
120+ def _version_for_config (version , config ):
98121 """Returns the version string for retrieving a framework version's specific config."""
99122 if "version_aliases" in config :
100123 if version in config ["version_aliases" ].keys ():
101124 return config ["version_aliases" ][version ]
102125
103- _validate_arg ("{} version" .format (framework ), version , config ["versions" ].keys ())
104126 return version
105127
106128
@@ -112,6 +134,10 @@ def _registry_from_region(region, registry_dict):
112134
113135def _processor (instance_type , available_processors ):
114136 """Returns the processor type for the given instance type."""
137+ if not available_processors :
138+ logger .info ("Ignoring unnecessary instance type: %s." , instance_type )
139+ return None
140+
115141 if instance_type .startswith ("local" ):
116142 processor = "cpu" if instance_type == "local" else "gpu"
117143 elif not instance_type .startswith ("ml." ):
@@ -129,9 +155,12 @@ def _processor(instance_type, available_processors):
129155
130156def _validate_py_version_and_set_if_needed (py_version , version_config ):
131157 """Checks if the Python version is one of the supported versions."""
132- available_versions = version_config .get ("py_versions" , version_config .keys ())
158+ if "repository" in version_config :
159+ available_versions = version_config .get ("py_versions" )
160+ else :
161+ available_versions = list (version_config .keys ())
133162
134- if len ( available_versions ) == 0 :
163+ if not available_versions :
135164 if py_version :
136165 logger .info ("Ignoring unnecessary Python version: %s." , py_version )
137166 return None
0 commit comments