Add the region flag to the imagebuilder component custom script downl…

…oad action
aws · Sep 9, 2024 · c226014 · c226014
1 parent e647dc1
commit c226014
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 4 deletions.
diff --git a/cli/src/pcluster/imagebuilder_utils.py b/cli/src/pcluster/imagebuilder_utils.py
@@ -13,6 +13,7 @@
 import yaml
 
 from pcluster.aws.aws_api import AWSApi
+from pcluster.aws.common import get_region
 from pcluster.utils import get_url_scheme, yaml_load
 
 ROOT_VOLUME_TYPE = "gp3"
@@ -53,6 +54,8 @@ def wrap_script_to_component(url):
     with open(custom_component_script_template_file, "r", encoding="utf-8") as file:
         custom_component_script_template = yaml_load(file)
 
+    # region_action = _generate_action("Region", "set -v\necho {0}\n".format(get_region()))
+    # custom_component_script_template["phases"][0]["steps"].insert(0, region_action)
     script_url_action = _generate_action("ScriptUrl", "set -v\necho {0}\n".format(url))
     custom_component_script_template["phases"][0]["steps"].insert(0, script_url_action)
     script_scheme_action = _generate_action("ScriptUrlScheme", "set -v\necho {0}\n".format(scheme))

diff --git a/cli/src/pcluster/resources/imagebuilder/custom_script.yaml b/cli/src/pcluster/resources/imagebuilder/custom_script.yaml
@@ -21,7 +21,7 @@ phases:
               set -v
               mkdir /tmp/parallelcluster/
               temp_file=$(mktemp -p /tmp/parallelcluster/)
-              echo "${temp_file}"
+              echo "$temp_file"
 
       # Download script
       - name: Download
@@ -33,7 +33,7 @@ phases:
               if [[ {{ build.ScriptUrlScheme.outputs.stdout }} == "https" ]]; then
                 curl --retry 3 -L -o {{ build.TempScript.outputs.stdout }} {{ build.ScriptUrl.outputs.stdout }}
               elif [[ {{ build.ScriptUrlScheme.outputs.stdout }} == "s3" ]]; then
-                aws s3 cp {{ build.ScriptUrl.outputs.stdout }} {{ build.TempScript.outputs.stdout }}
+                aws s3 cp {{ build.ScriptUrl.outputs.stdout }} {{ build.TempScript.outputs.stdout }} --region ${AWS::Region}
               else
                 echo "Invalid script url"
                 exit {{ build.Fail.outputs.stdout }}

diff --git a/cli/src/pcluster/templates/imagebuilder_stack.py b/cli/src/pcluster/templates/imagebuilder_stack.py
@@ -30,7 +30,6 @@
 
 from pcluster import imagebuilder_utils, utils
 from pcluster.aws.aws_api import AWSApi
-from pcluster.aws.common import get_region
 from pcluster.config.common import BaseTag
 from pcluster.config.imagebuilder_config import ImageBuilderConfig, ImageBuilderExtraChefAttributes, Volume
 from pcluster.constants import (
@@ -979,7 +978,7 @@ def _add_custom_components(self, components, policy_statements, components_resou
                     description="This component is custom component for script, script name is {0}, script url is "
                     "{1}".format(component_script_name, custom_component.value),
                     platform="Linux",
-                    data=wrap_script_to_component(custom_component.value),
+                    data=Fn.sub(wrap_script_to_component(custom_component.value)),
                 )
                 components.append(
                     imagebuilder.CfnImageRecipe.ComponentConfigurationProperty(component_arn=Fn.ref(component_id))

diff --git a/tests/integration-tests/configs/mini.yaml b/tests/integration-tests/configs/mini.yaml
@@ -0,0 +1,10 @@
+test-suites:
+  createami:
+    test_createami.py::test_build_image_custom_components:
+      dimensions:
+        - instances:
+            - m6g.xlarge
+          oss:
+            - alinux2
+          regions:
+            - cn-northwest-1