Add support for Data Repository Associations with Pcluster Managed FSx for Lustre (#5693)

* Add support for DRAs as an option for s3 association with an FSx file system. FSx released Persistent_2 which no longer supports the use of ExportPath and ImportPath. Instead, it uses DRAs to configure S3 association.

* Add integration and unit tests to cover the creation, modification, and deletion of DRAs

Author: hgreebe

CHANGELOG.md

@@ -5,7 +5,7 @@ CHANGELOG
 ------
 
 **ENHANCEMENTS**
-
+- Add support for Data Repository Associations when using PERSISTENT_2 as DeploymentType for a managed FSx for Lustre.
 **CHANGES**
 
 **BUG FIXES**

cli/src/pcluster/aws/fsx.py

@@ -21,6 +21,7 @@ def __init__(self):
         self.svm_cache = {}
         self.volume_cache = {}
         self.fc_cache = {}
+        self.dra_cache = {}
 
     @AWSExceptionHandler.handle_client_exception
     def get_file_systems_info(self, fsx_fs_ids):
@@ -84,6 +85,24 @@ def describe_volumes(self, volume_ids):
                 result.append(volume)
         return result
 
+    @AWSExceptionHandler.handle_client_exception
+    def describe_data_repository_associations(self, dra_ids):
+        """Describe FSx data repository associations."""
+        result = []
+        missed_dra_ids = []
+        for dra_id in dra_ids:
+            cached_data = self.dra_cache.get(dra_id)
+            if cached_data:
+                result.append(cached_data)
+            else:
+                missed_dra_ids.append(dra_id)
+        if missed_dra_ids:
+            response = self._client.describe_data_repository_associations(AssociationIds=missed_dra_ids)["Associations"]
+            for dra in response:
+                self.dra_cache[dra.get("AssociationId")] = dra
+                result.append(dra)
+        return result
+
     @AWSExceptionHandler.handle_client_exception
     def describe_backup(self, backup_id):
         """Describe backup."""

cli/src/pcluster/config/cluster_config.py

@@ -142,6 +142,7 @@
     FsxAutoImportValidator,
     FsxBackupIdValidator,
     FsxBackupOptionsValidator,
+    FsxDraValidator,
     FsxPersistentOptionsValidator,
     FsxS3Validator,
     FsxStorageCapacityValidator,
@@ -431,6 +432,38 @@ def existing_dns_name(self):
         return self.file_system_data.dns_name if self.is_unmanaged else ""
 
 
+class DataRepositoryAssociation(Resource):
+    """Represent the Data Repository Association resource."""
+
+    def __init__(
+        self,
+        name: str,
+        data_repository_path: str,
+        file_system_path: str,
+        batch_import_meta_data_on_create: bool = None,
+        imported_file_chunk_size: int = None,
+        auto_export_policy: List[str] = None,
+        auto_import_policy: List[str] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.name = Resource.init_param(name)
+        self.batch_import_meta_data_on_create = Resource.init_param(batch_import_meta_data_on_create, default=False)
+        self.data_repository_path = Resource.init_param(data_repository_path)
+        self.file_system_path = Resource.init_param(file_system_path)
+        self.imported_file_chunk_size = Resource.init_param(imported_file_chunk_size, default=1024)
+        self.auto_export_policy = Resource.init_param(auto_export_policy)
+        self.auto_import_policy = Resource.init_param(auto_import_policy)
+
+    def _register_validators(self, context: ValidatorContext = None):
+        super()._register_validators(context)
+        self._register_validator(SharedStorageNameValidator, name=self.name)
+        self._register_validator(S3BucketUriValidator, url=self.data_repository_path)
+        self._register_validator(
+            FsxAutoImportValidator, auto_import_policy=self.auto_import_policy, import_path=self.data_repository_path
+        )
+
+
 class SharedFsxLustre(BaseSharedFsx):
     """Represent the shared FSX resource."""
 
@@ -454,6 +487,7 @@ def __init__(
         drive_cache_type: str = None,
         fsx_storage_type: str = None,
         deletion_policy: str = None,
+        data_repository_associations: List[DataRepositoryAssociation] = None,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -481,9 +515,16 @@ def __init__(
         self.deletion_policy = Resource.init_param(
             deletion_policy, default=DELETE_POLICY if not file_system_id else None
         )
+        self.data_repository_associations = Resource.init_param(data_repository_associations)
 
     def _register_validators(self, context: ValidatorContext = None):
         super()._register_validators(context)
+        self._register_validator(
+            FsxDraValidator,
+            data_repository_associations=self.data_repository_associations,
+            import_path=self.import_path,
+            export_path=self.export_path,
+        )
         self._register_validator(
             FsxS3Validator,
             import_path=self.import_path,

cli/src/pcluster/config/config_patch.py

@@ -189,18 +189,19 @@ def _compare_list(self, base_section, target_section, param_path, data_key, fiel
                     )
                 )
         # Then, compare all non visited base sections vs target config.
-        for base_nested_section in base_section.get(data_key, []):
-            if not base_nested_section.get("visited", False):
-                self.changes.append(
-                    Change(
-                        param_path,
-                        data_key,
-                        base_nested_section,
-                        None,
-                        change_update_policy,
-                        is_list=True,
+        if base_section:
+            for base_nested_section in base_section.get(data_key, []):
+                if not base_nested_section.get("visited", False):
+                    self.changes.append(
+                        Change(
+                            param_path,
+                            data_key,
+                            base_nested_section,
+                            None,
+                            change_update_policy,
+                            is_list=True,
+                        )
                     )
-                )
 
     @property
     def update_policy_level(self):

cli/src/pcluster/schemas/cluster_schema.py

@@ -42,6 +42,7 @@
     CustomActions,
     Dashboards,
     Database,
+    DataRepositoryAssociation,
     Dcv,
     DirectoryService,
     Dns,
@@ -355,6 +356,38 @@ def validate_existence_of_mode_throughput(self, data, **kwargs):
             )
 
 
+class DataRepositoryAssociationSchema(BaseSchema):
+    """Represent the DRA schema."""
+
+    name = fields.Str(required=True, metadata={"update_policy": UpdatePolicy.UNSUPPORTED})
+    batch_import_meta_data_on_create = fields.Bool(metadata={"update_policy": UpdatePolicy.UNSUPPORTED})
+    data_repository_path = fields.Str(required=True, metadata={"update_policy": UpdatePolicy.UNSUPPORTED})
+    file_system_path = fields.Str(
+        required=True,
+        validate=validate.Regexp(r"^[^\u0000\u0085\u2028\u2029\r\n]{1,4096}$"),
+        metadata={"update_policy": UpdatePolicy.UNSUPPORTED},
+    )
+    imported_file_chunk_size = fields.Int(
+        validate=validate.Range(min=1, max=512000, error="has a minimum size of 1 MiB, and max size of 512,000 MiB"),
+        metadata={"update_policy": UpdatePolicy.SUPPORTED},
+    )
+    auto_export_policy = fields.List(
+        fields.Str(validate=validate.OneOf(["NEW", "CHANGED", "DELETED"])),
+        validate=validate.Length(max=3),
+        metadata={"update_policy": UpdatePolicy.SUPPORTED},
+    )
+    auto_import_policy = fields.List(
+        fields.Str(validate=validate.OneOf(["NEW", "CHANGED", "DELETED"])),
+        validate=validate.Length(max=3),
+        metadata={"update_policy": UpdatePolicy.SUPPORTED},
+    )
+
+    @post_load
+    def make_resource(self, data, **kwargs):
+        """Generate resource."""
+        return DataRepositoryAssociation(**data)
+
+
 class FsxLustreSettingsSchema(BaseSchema):
     """Represent the FSX schema."""
 
@@ -408,6 +441,14 @@ class FsxLustreSettingsSchema(BaseSchema):
     deletion_policy = fields.Str(
         validate=validate.OneOf(DELETION_POLICIES), metadata={"update_policy": UpdatePolicy.SUPPORTED}
     )
+    data_repository_associations = fields.Nested(
+        DataRepositoryAssociationSchema,
+        many=True,
+        metadata={
+            "update_policy": UpdatePolicy.SUPPORTED,
+            "update_key": "Name",
+        },
+    )
 
     @validates_schema
     def validate_file_system_id_ignored_parameters(self, data, **kwargs):

cli/src/pcluster/templates/cluster_stack.py

@@ -962,6 +962,9 @@ def _add_fsx_storage(self, id: str, shared_fsx: BaseSharedFsx):
             ) = convert_deletion_policy(shared_fsx.deletion_policy)
 
             fsx_id = fsx_resource.ref
+
+            self._add_dra(fsx_id, shared_fsx)
+
             # Get MountName for new filesystem. DNSName cannot be retrieved from CFN and will be generated in cookbook
             mount_name = fsx_resource.attr_lustre_mount_name
 
@@ -976,6 +979,30 @@ def _add_fsx_storage(self, id: str, shared_fsx: BaseSharedFsx):
 
         return fsx_id
 
+    def _add_dra(self, fsx_id: str, shared_fsx: BaseSharedFsx):
+        """Add Cfn Data Repository Association Resources."""
+        if shared_fsx.data_repository_associations:
+            for dra in shared_fsx.data_repository_associations:
+                dra_id = "{0}{1}".format(dra.name, create_hash_suffix(dra.name))
+                fsx.CfnDataRepositoryAssociation(
+                    self.stack,
+                    dra_id,
+                    batch_import_meta_data_on_create=dra.batch_import_meta_data_on_create,
+                    data_repository_path=dra.data_repository_path,
+                    file_system_id=fsx_id,
+                    file_system_path=dra.file_system_path,
+                    imported_file_chunk_size=dra.imported_file_chunk_size,
+                    s3=fsx.CfnDataRepositoryAssociation.S3Property(
+                        auto_export_policy=fsx.CfnDataRepositoryAssociation.AutoExportPolicyProperty(
+                            events=dra.auto_export_policy
+                        ),
+                        auto_import_policy=fsx.CfnDataRepositoryAssociation.AutoImportPolicyProperty(
+                            events=dra.auto_import_policy
+                        ),
+                    ),
+                    tags=[CfnTag(key="Name", value=dra.name)],
+                )
+
     def _add_efs_storage(self, id: str, shared_efs: SharedEfs):
         """Add specific Cfn Resources to map the EFS storage."""
         # EFS FileSystem

cli/src/pcluster/validators/fsx_validators.py

@@ -48,6 +48,28 @@ def _validate(
             )
 
 
+class FsxDraValidator(Validator):
+    """
+    FSX Dra validator.
+
+    Verify compatibility of given S3 association options for FSX.
+    """
+
+    def _validate(self, data_repository_associations, import_path, export_path):
+        if data_repository_associations and (import_path or export_path):
+            self._add_failure(
+                "When specifying data repository associations, import and export path "
+                "can not be used on the same file system.",
+                FailureLevel.ERROR,
+            )
+
+        if data_repository_associations and len(data_repository_associations) > 8:
+            self._add_failure(
+                "The number of data repository association used for one file system cannot be greater that 8.",
+                FailureLevel.ERROR,
+            )
+
+
 class FsxPersistentOptionsValidator(Validator):
     """
     FSX persistent options validator.

cli/tests/pcluster/aws/test_fsx.py

@@ -134,6 +134,45 @@ def get_describe_file_caches_mocked_request(file_cache_ids):
     )
 
 
+def get_describe_data_repository_associations_mocked_request(dra_ids, lifecycle):
+    return MockedBoto3Request(
+        method="describe_data_repository_associations",
+        response={"Associations": [{"AssociationId": dra_id, "Lifecycle": lifecycle} for dra_id in dra_ids]},
+        expected_params={"AssociationIds": dra_ids},
+    )
+
+
+def test_describe_data_repository_associations(boto3_stubber):
+    dra_id = "dra-012345678"
+    additional_dra_id = "dra-876543210"
+    # The first mocked request and the third are about the same fsx. However, the lifecycle of the fsx changes
+    # from CREATING to AVAILABLE. The second mocked request is about another fsx
+    mocked_requests = [
+        get_describe_data_repository_associations_mocked_request([dra_id], "CREATING"),
+        get_describe_data_repository_associations_mocked_request([additional_dra_id], "CREATING"),
+        get_describe_data_repository_associations_mocked_request([dra_id], "AVAILABLE"),
+    ]
+    boto3_stubber("fsx", mocked_requests)
+    assert_that(AWSApi.instance().fsx.describe_data_repository_associations([dra_id])[0]["Lifecycle"]).is_equal_to(
+        "CREATING"
+    )
+
+    # Second boto3 call with more volumes. The volume already cached should not be included in the boto3 call.
+    response = AWSApi.instance().fsx.describe_data_repository_associations([dra_id, additional_dra_id])
+    assert_that(response).is_length(2)
+
+    # Third boto3 call. The result should be from cache even if the lifecycle of the SVM is different
+    assert_that(AWSApi.instance().fsx.describe_data_repository_associations([dra_id])[0]["Lifecycle"]).is_equal_to(
+        "CREATING"
+    )
+
+    # Fourth boto3 call after resetting the AWSApi instance. The latest fsx lifecycle should be retrieved from boto3
+    AWSApi.reset()
+    assert_that(AWSApi.instance().fsx.describe_data_repository_associations([dra_id])[0]["Lifecycle"]).is_equal_to(
+        "AVAILABLE"
+    )
+
+
 def test_describe_file_caches(boto3_stubber):
     file_cache_id = "fc-12345678"
     additional_file_cache_id = "fc-123456789012345678"

cli/tests/pcluster/example_configs/slurm.full.yaml

@@ -182,10 +182,10 @@ SharedStorage:
     FsxLustreSettings:
       StorageCapacity: 3600
       DeploymentType: PERSISTENT_1  # PERSISTENT_1 | PERSISTENT_2 | SCRATCH_1 | SCRATCH_2
-      ImportedFileChunkSize: 1024
+      # ImportedFileChunkSize: 1024 # ImportedFileChunkSize cannot coexist with some of the fields
       DataCompressionType: LZ4
-      ExportPath: s3://bucket/folder
-      ImportPath: s3://bucket
+      # ExportPath: s3://bucket/folder # ExportPath cannot coexist with some of the fields
+      # ImportPath: s3://bucket # ImportPath cannot coexist with some of the fields
       WeeklyMaintenanceStartTime: "1:00:00"
       AutomaticBackupRetentionDays: 0
       CopyTagsToBackups: true
@@ -194,9 +194,17 @@ SharedStorage:
       # BackupId: backup-fedcba98 # BackupId cannot coexist with some of the fields
       KmsKeyId: xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx
       # FileSystemId: fs-12345678123456789 # FileSystemId cannot coexist with other fields
-      AutoImportPolicy: NEW  # NEW | NEW_CHANGED | NEW_CHANGED_DELETED
+      # AutoImportPolicy: NEW  # NEW | NEW_CHANGED | NEW_CHANGED_DELETED # AutoImportPolicy cannot coexist with some of the fields
       DriveCacheType: READ  # READ
       StorageType: HDD  # HDD | SSD
+      DataRepositoryAssociations:
+        - Name: dra
+          BatchImportMetaDataOnCreate: false
+          DataRepositoryPath: s3://bucket/folder
+          FileSystemPath: /
+          ImportedFileChunkSize: 1024
+          AutoExportPolicy: [ NEW, CHANGED, DELETED ]
+          AutoImportPolicy: [ NEW, CHANGED, DELETED ]
   - MountDir: /my/mount/point4
     Name: name4
     StorageType: FsxOpenZfs