From 9218095c2da800c946ad2090c8a87f10d35340b2 Mon Sep 17 00:00:00 2001 From: Nikhil Kulkarni Date: Fri, 2 Oct 2020 10:30:36 -0700 Subject: [PATCH 1/6] Add a condition to retrieve correct image URI for xgboost --- src/sagemaker/model.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/src/sagemaker/model.py b/src/sagemaker/model.py index 247b04cf79..ecbdcf343f 100644 --- a/src/sagemaker/model.py +++ b/src/sagemaker/model.py @@ -274,9 +274,17 @@ def _compilation_image_uri(self, region, target_instance_type, framework, framew framework (str): The framework name. framework_version (str): The framework version. """ - framework_prefix = "inferentia-" if target_instance_type.startswith("ml_inf") else "neo-" + framework_prefix = "" + framework_suffix = "" + + if framework == "xgboost": + framework_suffix = "-neo" + else: + framework_prefix = "inferentia-" if target_instance_type.startswith("ml_inf") else "neo-" + + return image_uris.retrieve( - "{}{}".format(framework_prefix, framework), + "{}{}".format(framework_prefix, framework, framework_suffix), region, instance_type=target_instance_type, version=framework_version, From 7ff4d556589c26fff1514351c3d38919bf9510af Mon Sep 17 00:00:00 2001 From: Nikhil Kulkarni Date: Fri, 2 Oct 2020 11:53:27 -0700 Subject: [PATCH 2/6] Fix f string --- src/sagemaker/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/sagemaker/model.py b/src/sagemaker/model.py index ecbdcf343f..933df5659d 100644 --- a/src/sagemaker/model.py +++ b/src/sagemaker/model.py @@ -284,7 +284,7 @@ def _compilation_image_uri(self, region, target_instance_type, framework, framew return image_uris.retrieve( - "{}{}".format(framework_prefix, framework, framework_suffix), + "{}{}{}".format(framework_prefix, framework, framework_suffix), region, instance_type=target_instance_type, version=framework_version, From af73379ef52d3e188ae5f9938e02e0a02eac97eb Mon Sep 17 00:00:00 2001 From: Nikhil Kulkarni Date: Fri, 2 Oct 2020 13:52:01 -0700 Subject: [PATCH 3/6] fix pylint --- src/sagemaker/model.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/sagemaker/model.py b/src/sagemaker/model.py index 933df5659d..d24ba19b08 100644 --- a/src/sagemaker/model.py +++ b/src/sagemaker/model.py @@ -276,11 +276,12 @@ def _compilation_image_uri(self, region, target_instance_type, framework, framew """ framework_prefix = "" framework_suffix = "" - + if framework == "xgboost": framework_suffix = "-neo" else: - framework_prefix = "inferentia-" if target_instance_type.startswith("ml_inf") else "neo-" + framework_prefix = "inferentia-" if \ + target_instance_type.startswith("ml_inf") else "neo-" return image_uris.retrieve( From 70da482e5de184d1dc39b9f24d11b4da2b0b4ddd Mon Sep 17 00:00:00 2001 From: Nikhil Kulkarni Date: Fri, 2 Oct 2020 14:08:15 -0700 Subject: [PATCH 4/6] fix black-check --- src/sagemaker/model.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/sagemaker/model.py b/src/sagemaker/model.py index d24ba19b08..9c75f430ed 100644 --- a/src/sagemaker/model.py +++ b/src/sagemaker/model.py @@ -280,9 +280,10 @@ def _compilation_image_uri(self, region, target_instance_type, framework, framew if framework == "xgboost": framework_suffix = "-neo" else: - framework_prefix = "inferentia-" if \ - target_instance_type.startswith("ml_inf") else "neo-" - + if target_instance_type.startswith("ml_inf"): + framework_prefix = "inferentia" + else: + framework_prefix = "neo-" return image_uris.retrieve( "{}{}{}".format(framework_prefix, framework, framework_suffix), From 783908a60939a56cd708afb9fede5a1834aa8109 Mon Sep 17 00:00:00 2001 From: Nikhil Kulkarni Date: Fri, 2 Oct 2020 14:18:50 -0700 Subject: [PATCH 5/6] Fix inferentia- --- src/sagemaker/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/sagemaker/model.py b/src/sagemaker/model.py index 9c75f430ed..2c770a8c63 100644 --- a/src/sagemaker/model.py +++ b/src/sagemaker/model.py @@ -281,7 +281,7 @@ def _compilation_image_uri(self, region, target_instance_type, framework, framew framework_suffix = "-neo" else: if target_instance_type.startswith("ml_inf"): - framework_prefix = "inferentia" + framework_prefix = "inferentia-" else: framework_prefix = "neo-" From 1494bb06b1b32193ab4549619ae7d0444a641ca1 Mon Sep 17 00:00:00 2001 From: Nikhil Kulkarni Date: Fri, 2 Oct 2020 15:16:21 -0700 Subject: [PATCH 6/6] Fix if condition --- src/sagemaker/model.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/sagemaker/model.py b/src/sagemaker/model.py index 2c770a8c63..38006ea136 100644 --- a/src/sagemaker/model.py +++ b/src/sagemaker/model.py @@ -279,11 +279,10 @@ def _compilation_image_uri(self, region, target_instance_type, framework, framew if framework == "xgboost": framework_suffix = "-neo" + elif target_instance_type.startswith("ml_inf"): + framework_prefix = "inferentia-" else: - if target_instance_type.startswith("ml_inf"): - framework_prefix = "inferentia-" - else: - framework_prefix = "neo-" + framework_prefix = "neo-" return image_uris.retrieve( "{}{}{}".format(framework_prefix, framework, framework_suffix),