Skip to content

Commit 2635ac4

Browse files
hanwen-clusterhanwen-cluster
committed
Add options to install FSx Lustre and Nvidia software
Signed-off-by: Hanwen <[email protected]>
1 parent c389847 commit 2635ac4

File tree

7 files changed

+172
-49
lines changed

7 files changed

+172
-49
lines changed

‎cli/src/pcluster/config/imagebuilder_config.py‎

Lines changed: 45 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,41 @@ def __init__(
143143
self.enabled = enabled
144144

145145

146+
class LustreClient(Resource):
147+
"""Represent the LustreClient configuration for the ImageBuilder."""
148+
149+
def __init__(
150+
self,
151+
enabled: bool = None,
152+
):
153+
super().__init__()
154+
self.enabled = Resource.init_param(enabled, default=True)
155+
156+
157+
class NvidiaSoftware(Resource):
158+
"""Represent the NvidiaSoftware configuration for the ImageBuilder."""
159+
160+
def __init__(
161+
self,
162+
enabled: bool = None,
163+
):
164+
super().__init__()
165+
self.enabled = Resource.init_param(enabled, default=False)
166+
167+
168+
class Installation(Resource):
169+
"""Represent the installation configuration for the ImageBuilder."""
170+
171+
def __init__(
172+
self,
173+
lustre_client: LustreClient = None,
174+
nvidia_software: NvidiaSoftware = None,
175+
):
176+
super().__init__()
177+
self.lustre_client = lustre_client or LustreClient()
178+
self.nvidia_software = nvidia_software or NvidiaSoftware()
179+
180+
146181
class Build(Resource):
147182
"""Represent the build configuration for the ImageBuilder."""
148183

@@ -157,6 +192,7 @@ def __init__(
157192
components: List[Component] = None,
158193
update_os_packages: UpdateOsPackages = None,
159194
imds: Imds = None,
195+
installation: Installation = None,
160196
):
161197
super().__init__()
162198
self.instance_type = Resource.init_param(instance_type)
@@ -168,6 +204,7 @@ def __init__(
168204
self.components = components
169205
self.update_os_packages = update_os_packages
170206
self.imds = imds or Imds(implied="v2.0")
207+
self.installation = installation or Installation()
171208

172209
def _register_validators(self, context: ValidatorContext = None): # noqa: D102 #pylint: disable=unused-argument
173210
self._register_validator(
@@ -282,21 +319,24 @@ def lambda_functions_vpc_config(self):
282319
class ImageBuilderExtraChefAttributes(ExtraChefAttributes):
283320
"""Extra Attributes for ImageBuilder Chef Client."""
284321

285-
def __init__(self, dev_settings: ImagebuilderDevSettings):
286-
super().__init__(dev_settings)
322+
def __init__(self, config: ImageBuilderConfig):
323+
super().__init__(config.dev_settings)
287324
self.region = None
288325
self.nvidia = None
326+
self.lustre = None
289327
self.is_official_ami_build = None
290328
self.custom_node_package = None
291329
self.custom_awsbatchcli_package = None
292330
self.base_os = None
293331
self.disable_kernel_update = None
294332
self.slurm_patches_s3_archive = None
295-
self._set_default(dev_settings)
333+
self._set_default(config)
296334

297-
def _set_default(self, dev_settings: ImagebuilderDevSettings):
335+
def _set_default(self, config: ImageBuilderConfig):
336+
dev_settings = config.dev_settings
298337
self.region = "{{ build.AWSRegion.outputs.stdout }}"
299-
self.nvidia = {"enabled": "no"}
338+
self.nvidia = {"enabled": "yes"} if config.build.installation.nvidia_software.enabled else {"enabled": "no"}
339+
self.lustre = {"enabled": "yes"} if config.build.installation.lustre_client.enabled else {"enabled": "no"}
300340
self.is_official_ami_build = "false"
301341
self.custom_node_package = dev_settings.node_package if dev_settings and dev_settings.node_package else ""
302342
self.custom_awsbatchcli_package = (

‎cli/src/pcluster/schemas/imagebuilder_schema.py‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,9 @@
2626
ImageBuilderConfig,
2727
ImagebuilderDeploymentSettings,
2828
ImagebuilderDevSettings,
29+
Installation,
30+
LustreClient,
31+
NvidiaSoftware,
2932
UpdateOsPackages,
3033
Volume,
3134
)
@@ -167,6 +170,40 @@ def make_resource(self, data, **kwargs):
167170
return UpdateOsPackages(**data)
168171

169172

173+
class LustreClientSchema(BaseSchema):
174+
"""Represent the schema of the ImageBuilder NvidiaSoftware."""
175+
176+
enabled = fields.Bool()
177+
178+
@post_load
179+
def make_resource(self, data, **kwargs):
180+
"""Generate resource."""
181+
return LustreClient(**data)
182+
183+
184+
class NvidiaSoftwareSchema(BaseSchema):
185+
"""Represent the schema of the ImageBuilder NvidiaSoftware."""
186+
187+
enabled = fields.Bool()
188+
189+
@post_load
190+
def make_resource(self, data, **kwargs):
191+
"""Generate resource."""
192+
return NvidiaSoftware(**data)
193+
194+
195+
class InstallationSchema(BaseSchema):
196+
"""Represent the schema of the ImageBuilder Installation."""
197+
198+
lustre_client = fields.Nested(LustreClientSchema)
199+
nvidia_software = fields.Nested(NvidiaSoftwareSchema)
200+
201+
@post_load
202+
def make_resource(self, data, **kwargs):
203+
"""Generate resource."""
204+
return Installation(**data)
205+
206+
170207
class BuildSchema(BaseSchema):
171208
"""Represent the schema of the ImageBuilder Build."""
172209

@@ -179,6 +216,7 @@ class BuildSchema(BaseSchema):
179216
subnet_id = fields.Str(validate=get_field_validator("subnet_id"))
180217
update_os_packages = fields.Nested(UpdateOsPackagesSchema)
181218
imds = fields.Nested(ImdsSchema)
219+
installation = fields.Nested(InstallationSchema)
182220

183221
@post_load
184222
def make_resource(self, data, **kwargs):

‎cli/src/pcluster/templates/imagebuilder_stack.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -197,7 +197,7 @@ def _add_cfn_parameters(self):
197197
self,
198198
"CfnParamChefDnaJson",
199199
type="String",
200-
default=ImageBuilderExtraChefAttributes(self.config.dev_settings).dump_json(),
200+
default=ImageBuilderExtraChefAttributes(self.config).dump_json(),
201201
description="ChefAttributes",
202202
)
203203
CfnParameter(

‎cli/tests/pcluster/config/dummy_imagebuilder_config.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,9 @@
1919
ImageBuilderConfig,
2020
ImagebuilderDeploymentSettings,
2121
ImagebuilderDevSettings,
22+
Installation,
23+
LustreClient,
24+
NvidiaSoftware,
2225
UpdateOsPackages,
2326
Volume,
2427
)
@@ -39,6 +42,9 @@
3942
"additional_iam_policies": AdditionalIamPolicy,
4043
"update_os_packages": UpdateOsPackages,
4144
"imds": Imds,
45+
"installation": Installation,
46+
"lustre_client": LustreClient,
47+
"nvidia_software": NvidiaSoftware,
4248
}
4349

4450

‎cli/tests/pcluster/models/test_imagebuilder.py‎

Lines changed: 74 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -157,16 +157,18 @@ def test_imagebuilder_url_validator(
157157
[
158158
(
159159
{
160-
"build": {
161-
"parent_image": "ami-0185634c5a8a37250",
162-
"instance_type": "c5.xlarge",
163-
"update_os_packages": {"enabled": True},
164-
},
165-
"dev_settings": {
166-
"node_package": "s3://test/aws-parallelcluster-node-3.0.tgz",
167-
"aws_batch_cli_package": "https://test/aws-parallelcluster-3.0.tgz",
168-
"disable_kernel_update": "true",
169-
},
160+
"imagebuilder": {
161+
"build": {
162+
"parent_image": "ami-0185634c5a8a37250",
163+
"instance_type": "c5.xlarge",
164+
"update_os_packages": {"enabled": True},
165+
},
166+
"dev_settings": {
167+
"node_package": "s3://test/aws-parallelcluster-node-3.0.tgz",
168+
"aws_batch_cli_package": "https://test/aws-parallelcluster-3.0.tgz",
169+
"disable_kernel_update": "true",
170+
},
171+
}
170172
},
171173
{
172174
"cluster": {
@@ -176,20 +178,28 @@ def test_imagebuilder_url_validator(
176178
"disable_kernel_update": "true",
177179
"is_official_ami_build": "false",
178180
"nvidia": {"enabled": "no"},
181+
"lustre": {"enabled": "yes"},
179182
"region": "{{ build.AWSRegion.outputs.stdout }}",
180183
"slurm_patches_s3_archive": "",
181184
}
182185
},
183186
),
184187
(
185188
{
186-
"dev_settings": {
187-
"cookbook": {
188-
"chef_cookbook": "https://test/aws-parallelcluster-cookbook-3.0.tgz",
189-
"extra_chef_attributes": '{"cluster": {"nvidia": { "enabled" : "yes" }, "dcv" :"no"}}',
189+
"imagebuilder": {
190+
"build": {
191+
"parent_image": "ami-0185634c5a8a37250",
192+
"instance_type": "c5.xlarge",
193+
"installation": {"lustre_client": {"enabled": "False"}},
190194
},
191-
"node_package": "s3://test/aws-parallelcluster-node-3.0.tgz",
192-
},
195+
"dev_settings": {
196+
"cookbook": {
197+
"chef_cookbook": "https://test/aws-parallelcluster-cookbook-3.0.tgz",
198+
"extra_chef_attributes": '{"cluster": {"nvidia": { "enabled" : "yes" }, "dcv" :"no"}}',
199+
},
200+
"node_package": "s3://test/aws-parallelcluster-node-3.0.tgz",
201+
},
202+
}
193203
},
194204
{
195205
"cluster": {
@@ -199,6 +209,7 @@ def test_imagebuilder_url_validator(
199209
"dcv": "no",
200210
"disable_kernel_update": "false",
201211
"is_official_ami_build": "false",
212+
"lustre": {"enabled": "no"},
202213
"nvidia": {"enabled": "yes"},
203214
"region": "{{ build.AWSRegion.outputs.stdout }}",
204215
"slurm_patches_s3_archive": "",
@@ -207,13 +218,20 @@ def test_imagebuilder_url_validator(
207218
),
208219
(
209220
{
210-
"dev_settings": {
211-
"cookbook": {
212-
"chef_cookbook": "https://test/aws-parallelcluster-cookbook-3.0.tgz",
213-
"extra_chef_attributes": '{"cluster": {"nvidia": { "enabled" : "yes" }, "dcv" :"no"}, '
214-
'"nfs": "true"}',
221+
"imagebuilder": {
222+
"build": {
223+
"parent_image": "ami-0185634c5a8a37250",
224+
"instance_type": "c5.xlarge",
225+
"installation": {"lustre_client": {"enabled": "True"}},
226+
},
227+
"dev_settings": {
228+
"cookbook": {
229+
"chef_cookbook": "https://test/aws-parallelcluster-cookbook-3.0.tgz",
230+
"extra_chef_attributes": '{"cluster": {"nvidia": { "enabled" : "yes" }, "dcv" :"no"}, '
231+
'"nfs": "true"}',
232+
},
233+
"aws_batch_cli_package": "https://test/aws-parallelcluster-3.0.tgz",
215234
},
216-
"aws_batch_cli_package": "https://test/aws-parallelcluster-3.0.tgz",
217235
},
218236
},
219237
{
@@ -225,6 +243,7 @@ def test_imagebuilder_url_validator(
225243
"disable_kernel_update": "false",
226244
"is_official_ami_build": "false",
227245
"nvidia": {"enabled": "yes"},
246+
"lustre": {"enabled": "yes"},
228247
"region": "{{ build.AWSRegion.outputs.stdout }}",
229248
"slurm_patches_s3_archive": "",
230249
},
@@ -233,12 +252,19 @@ def test_imagebuilder_url_validator(
233252
),
234253
(
235254
{
236-
"dev_settings": {
237-
"cookbook": {
238-
"chef_cookbook": "https://test/aws-parallelcluster-cookbook-3.0.tgz",
239-
"extra_chef_attributes": '{"cluster": {"is_official_ami_build": "true"},"nfs": "true"}',
255+
"imagebuilder": {
256+
"build": {
257+
"parent_image": "ami-0185634c5a8a37250",
258+
"instance_type": "c5.xlarge",
259+
"installation": {"nvidia_software": {"enabled": True}},
260+
},
261+
"dev_settings": {
262+
"cookbook": {
263+
"chef_cookbook": "https://test/aws-parallelcluster-cookbook-3.0.tgz",
264+
"extra_chef_attributes": '{"cluster": {"is_official_ami_build": "true"},"nfs": "true"}',
265+
},
266+
"aws_batch_cli_package": "https://test/aws-parallelcluster-3.0.tgz",
240267
},
241-
"aws_batch_cli_package": "https://test/aws-parallelcluster-3.0.tgz",
242268
},
243269
},
244270
{
@@ -248,7 +274,8 @@ def test_imagebuilder_url_validator(
248274
"custom_node_package": "",
249275
"disable_kernel_update": "false",
250276
"is_official_ami_build": "true",
251-
"nvidia": {"enabled": "no"},
277+
"nvidia": {"enabled": "yes"},
278+
"lustre": {"enabled": "yes"},
252279
"region": "{{ build.AWSRegion.outputs.stdout }}",
253280
"slurm_patches_s3_archive": "",
254281
},
@@ -258,15 +285,22 @@ def test_imagebuilder_url_validator(
258285
# Test case with URL for Slurm patches from S3
259286
(
260287
{
261-
"dev_settings": {
262-
"cookbook": {
263-
"extra_chef_attributes": "{"
264-
'"cluster": {'
265-
'"slurm_patches_s3_archive": "s3://example-s3-bucket/example-archive.tgz"'
266-
"}"
267-
"}"
288+
"imagebuilder": {
289+
"build": {
290+
"parent_image": "ami-0185634c5a8a37250",
291+
"instance_type": "c5.xlarge",
292+
"installation": {"nvidia_software": {"enabled": False}},
268293
},
269-
},
294+
"dev_settings": {
295+
"cookbook": {
296+
"extra_chef_attributes": "{"
297+
'"cluster": {'
298+
'"slurm_patches_s3_archive": "s3://example-s3-bucket/example-archive.tgz"'
299+
"}"
300+
"}"
301+
},
302+
},
303+
}
270304
},
271305
{
272306
"cluster": {
@@ -276,6 +310,7 @@ def test_imagebuilder_url_validator(
276310
"disable_kernel_update": "false",
277311
"is_official_ami_build": "false",
278312
"nvidia": {"enabled": "no"},
313+
"lustre": {"enabled": "yes"},
279314
"region": "{{ build.AWSRegion.outputs.stdout }}",
280315
"slurm_patches_s3_archive": "s3://example-s3-bucket/example-archive.tgz",
281316
}
@@ -284,9 +319,9 @@ def test_imagebuilder_url_validator(
284319
],
285320
)
286321
def test_imagebuilder_extra_chef_attributes(resource, dna_json):
287-
dev_settings = imagebuilder_factory(resource).get("dev_settings")
288-
chef_attributes = ImageBuilderExtraChefAttributes(dev_settings).dump_json()
289-
assert_that(chef_attributes).is_equal_to(json.dumps(dna_json))
322+
config = imagebuilder_factory(resource).get("imagebuilder")
323+
chef_attributes = ImageBuilderExtraChefAttributes(config).dump_json()
324+
assert_that(json.loads(chef_attributes)).is_equal_to(dna_json)
290325

291326

292327
def _test_imagebuilder(

‎cli/tests/pcluster/schemas/test_imagebuilder_schema/test_imagebuilder_schema/imagebuilder_schema_all.yaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,11 @@ Build:
3737
SubnetId: subnet-0d03dc52
3838
UpdateOsPackages:
3939
Enabled: true
40+
Installation:
41+
NvidiaSoftware:
42+
Enabled: true
43+
LustreClient:
44+
Enabled: true
4045

4146
DevSettings:
4247
DisablePclusterComponent: False

‎tests/integration-tests/tests/createami/test_createami/test_build_image/image.config.yaml‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,9 @@ Build:
2121
{% endif %}
2222
{% if os in ["ubuntu2204", "rhel9", "rocky9"] %}
2323
# Disable Lustre installation because these newer operating systems release new kernels more often. Lustre usually does not support the latest kernels
24-
DevSettings:
25-
Cookbook:
26-
ExtraChefAttributes: |
27-
{"cluster": {"lustre": {"enabled": "no" }}}
24+
Installation:
25+
LustreClient:
26+
Enabled: true
2827
{% endif %}
2928

3029
CustomS3Bucket: {{ bucket_name }}

0 commit comments

Comments
 (0)