[LoginNodes] Add DCV support for login nodes (#6363)

* Add CustomActions configuration for login nodes

- Added schema for login nodes custom actions
- Added CustomAction to LoginNodesPool resource in cluster_config.py

Signed-off-by: Chris Makin <cmakin@amazon.com>

* Update unit tests to include login node custom action changes

- Modified login node dna_json unit test to include pool_name

- Ran tox autoformatter

* Update test_dcv integration test to include login nodes

* Merge head and login nodes custom action schema unit test

- Update changelog
- Fix DescribeStacks resource error in unit test

* Add DCV params to login node config and schema

* Modify `cluster_config.py` to support login node DCV.
    * Added `has_dcv_enabled()` to the `LoginNodesPool` resource.
        * Checks if dcv is enabled in a login node pool.

    * Added `has_dcv_configured()` to `LoginNodes` resource
        * Checks for dcv configuration in login node pool.
        * For use in validating DCV to avoid validating region multiple times.

    * Added `architecture` and `instance_type_info` properties to `LoginNodesPool`  resource to check DCV eligibility.
    * Added `DcvValidator` for login node pools in `SlurmClusterConfig`.

* Modify `cluster_schema.py` to support login node DCV
    * Add Dcv field to `LoginNodePoolSchema`

Signed-off-by: Chris Makin <cmakin@amazon.com>

* Add login node DCV param to config unit tests

* Modify dcv-connect and login node security group to support DCV on login nodes

* Modify `dcv-connect.py`
    * Added `--login-node-ip` parameter
    * Added `_validate_login_node_ip_arg()`
        * Checks whether the `--login-node-ip` argument is a valid IPv4 address.
        * Checks that `--login-node-ip` is a login node public or private IP address.

* Modify `cluster.py`
    * Add `login_node_instances`  property to Cluster.
    * For use in `dcv_connect.py` to get public and private IPs of login nodes in the cluster.

* Add DCV access to login node security group
    * Update `add_login_nodes_security_group` in `cluster_stack.py` to include ingress property for DCV.

* Update unit tests to support login node DCV configuration

* Add S3 access policy to login node IAM resources

This is needed so DCV can determine whether a valid license is available.

* Add login node support to test_cloudwatch_logging integration test

* Modified `test_cloudwatch_logging` to support login nodes.
  * The current changes will not support testing multiple pools with different configurations (e.g. DCV enabled on one pool but not another).
* Added `login_node_ip` param to `RemoteCommandExecutor`.
  * This is used in `test_cloudwatch_logging` to run commands on all login nodes in a cluster.

* Add S3 bucket access to login node IAM unit test

* Update CHANGELOGOC

* Fix dcv-connect help message

---------

Signed-off-by: Chris Makin <cmakin@amazon.com>

Author: cjmakin

CHANGELOG.md

@@ -7,6 +7,7 @@ CHANGELOG
 **ENHANCEMENTS**
 
 - Add support for custom actions on login nodes.
+- Allow DCV connection on login nodes.
 
 **BUG FIXES**
 - Fix validator `EfaPlacementGroupValidator` so that it does not suggest to configure a Placement Group when Capacity Blocks are used.

cli/src/pcluster/cli/commands/dcv_connect.py

@@ -27,6 +27,7 @@
 
 DCV_CONNECT_SCRIPT = "/opt/parallelcluster/scripts/pcluster_dcv_connect.sh"
 LOGGER = logging.getLogger(__name__)
+IPV4_REGEX = r"^((25[0-5]|2[0-4][0-9]|1[0-9][0-9]|[1-9]?[0-9])\.){3}(25[0-5]|2[0-4][0-9]|1[0-9][0-9]|[1-9]?[0-9])$"
 
 
 class DCVConnectionError(Exception):
@@ -42,28 +43,55 @@ def _check_command_output(cmd):
     return sub.check_output(cmd, shell=True, universal_newlines=True, stderr=sub.STDOUT).strip()  # nosec B602 nosemgrep
 
 
+def _validate_login_node_ip_arg(ip, login_nodes):
+    """Check if ip is a valid IPv4 address and belongs to a login node."""
+    if not re.match(IPV4_REGEX, ip):
+        error(f"{ip} is not a valid IPv4 address.")
+
+    login_node_ips = set()
+
+    for node in login_nodes:
+        login_node_ips.update([node.public_ip, node.private_ip])
+
+    if ip not in login_node_ips:
+        error(f"{ip} is not an IP address of a login node.")
+
+
 def _dcv_connect(args):
     """
     Execute pcluster dcv connect command.
 
     :param args: pcluster cli arguments.
     """
     try:
-        head_node = Cluster(args.cluster_name).head_node_instance
+        cluster = Cluster(args.cluster_name)
+
+        if args.login_node_ip:
+            login_nodes = cluster.login_node_instances
+            _validate_login_node_ip_arg(args.login_node_ip, login_nodes)
+
+            node_ip = args.login_node_ip
+            default_user = login_nodes[0].default_user
+        else:
+            head_node = cluster.head_node_instance
+
+            node_ip = head_node.public_ip or head_node.private_ip
+            default_user = head_node.default_user
+
     except Exception as e:
         error(f"Unable to connect to the cluster.\n{e}")
     else:
-        head_node_ip = head_node.public_ip or head_node.private_ip
-        # Prepare ssh command to execute in the head node instance
-        cmd = 'ssh {CFN_USER}@{HEAD_NODE_IP} {KEY} "{REMOTE_COMMAND} /home/{CFN_USER}"'.format(
-            CFN_USER=head_node.default_user,
-            HEAD_NODE_IP=head_node_ip,
+
+        # Prepare ssh command to execute in the node instance
+        cmd = 'ssh {CFN_USER}@{NODE_IP} {KEY} "{REMOTE_COMMAND} /home/{CFN_USER}"'.format(
+            CFN_USER=default_user,
+            NODE_IP=node_ip,
             KEY="-i {0}".format(args.key_path) if args.key_path else "",
             REMOTE_COMMAND=DCV_CONNECT_SCRIPT,
         )
 
         try:
-            url = _retry(_retrieve_dcv_session_url, func_args=[cmd, args.cluster_name, head_node_ip], attempts=4)
+            url = _retry(_retrieve_dcv_session_url, func_args=[cmd, args.cluster_name, node_ip], attempts=4)
             url_message = f"Please use the following one-time URL in your browser within 30 seconds:\n{url}"
 
             if args.show_url:
@@ -80,12 +108,12 @@ def _dcv_connect(args):
             error(
                 "Something went wrong during DCV connection.\n{0}"
                 "Please check the logs in the /var/log/parallelcluster/ folder "
-                "of the head node and submit an issue {1}\n".format(e, PCLUSTER_ISSUES_LINK)
+                "of the node and submit an issue {1}\n".format(e, PCLUSTER_ISSUES_LINK)
             )
 
 
-def _retrieve_dcv_session_url(ssh_cmd, cluster_name, head_node_ip):
-    """Connect by ssh to the head node instance, prepare DCV session and return the DCV session URL."""
+def _retrieve_dcv_session_url(ssh_cmd, cluster_name, node_ip):
+    """Connect by ssh to the head or login node instance, prepare DCV session and return the DCV session URL."""
     try:
         LOGGER.debug("SSH command: %s", ssh_cmd)
         output = _check_command_output(ssh_cmd)
@@ -104,7 +132,7 @@ def _retrieve_dcv_session_url(ssh_cmd, cluster_name, head_node_ip):
             error(
                 "Something went wrong during DCV connection. Please manually execute the command:\n{0}\n"
                 "If the problem persists, please check the logs in the /var/log/parallelcluster/ folder "
-                "of the head node and submit an issue {1}".format(ssh_cmd, PCLUSTER_ISSUES_LINK)
+                "of the node and submit an issue {1}".format(ssh_cmd, PCLUSTER_ISSUES_LINK)
             )
 
     except sub.CalledProcessError as e:
@@ -117,7 +145,7 @@ def _retrieve_dcv_session_url(ssh_cmd, cluster_name, head_node_ip):
             raise DCVConnectionError(e.output)
 
     return "https://{IP}:{PORT}?authToken={TOKEN}#{SESSION_ID}".format(
-        IP=head_node_ip,
+        IP=node_ip,
         PORT=dcv_server_port,  # pylint: disable=E0606
         TOKEN=dcv_session_token,  # pylint: disable=E0606
         SESSION_ID=dcv_session_id,  # pylint: disable=E0606
@@ -152,7 +180,7 @@ class DcvConnectCommand(CliCommand):
 
     # CLI
     name = "dcv-connect"
-    help = "Permits to connect to the head node through an interactive session by using NICE DCV."
+    help = "Permits connection to the head or login nodes through an interactive session by using NICE DCV."
     description = help
 
     def __init__(self, subparsers):
@@ -162,6 +190,7 @@ def register_command_args(self, parser: ArgumentParser) -> None:  # noqa: D102
         parser.add_argument("-n", "--cluster-name", help="Name of the cluster to connect to", required=True)
         parser.add_argument("--key-path", dest="key_path", help="Key path of the SSH key to use for the connection")
         parser.add_argument("--show-url", action="store_true", default=False, help="Print URL and exit")
+        parser.add_argument("--login-node-ip", dest="login_node_ip", help="IP address of a login node to connect to")
 
     def execute(self, args: Namespace, extra_args: List[str]) -> None:  # noqa: D102  #pylint: disable=unused-argument
         _dcv_connect(args)

cli/src/pcluster/config/cluster_config.py

@@ -1346,6 +1346,7 @@ def __init__(
         networking: LoginNodesNetworking = None,
         count: int = None,
         ssh: LoginNodesSsh = None,
+        dcv: Dcv = None,
         custom_actions: CustomActions = None,
         iam: LoginNodesIam = None,
         gracetime_period: int = None,
@@ -1358,9 +1359,11 @@ def __init__(
         self.networking = networking
         self.count = Resource.init_param(count, default=1)
         self.ssh = ssh
+        self.dcv = dcv
         self.custom_actions = custom_actions
         self.iam = iam or LoginNodesIam(implied=True)
         self.gracetime_period = Resource.init_param(gracetime_period, default=10)
+        self.__instance_type_info = None
 
     @property
     def instance_profile(self):
@@ -1372,6 +1375,23 @@ def instance_role(self):
         """Return the IAM role for login nodes, if set."""
         return self.iam.instance_role if self.iam else None
 
+    @property
+    def architecture(self):
+        """Return login node pool architecture based on instance type."""
+        return self.instance_type_info.supported_architecture()[0]
+
+    @property
+    def instance_type_info(self) -> InstanceTypeInfo:
+        """Return login node pool instance type information as returned from aws ec2 describe-instance-types."""
+        if not self.__instance_type_info:
+            self.__instance_type_info = AWSApi.instance().ec2.get_instance_type_info(self.instance_type)
+        return self.__instance_type_info
+
+    @property
+    def has_dcv_enabled(self):
+        """Return True if DCV is enabled."""
+        return self.dcv and self.dcv.enabled
+
     def _register_validators(self, context: ValidatorContext = None):  # noqa: D102 #pylint: disable=unused-argument
         self._register_validator(InstanceTypeValidator, instance_type=self.instance_type)
         self._register_validator(NameValidator, name=self.name)
@@ -1388,6 +1408,14 @@ def __init__(
         super().__init__(**kwargs)
         self.pools = pools
 
+    @property
+    def has_dcv_enabled(self):
+        """Returns True if there is a pool in the cluster with DCV enabled."""
+        for pool in self.pools:
+            if pool.has_dcv_enabled:
+                return True
+        return False
+
 
 class HeadNode(Resource):
     """Represent the Head Node resource."""
@@ -2938,6 +2966,46 @@ def login_nodes_subnet_ids(self):
                 subnet_ids_set.add(subnet_id)
         return list(subnet_ids_set)
 
+    def _register_login_node_validators(self):
+        """Register all login node validators to ensure that the resource parameters are valid."""
+        has_dcv_configured = False
+        # Check if all subnets(head node, Login nodes, compute nodes) are in the same VPC and support DNS.
+        self._register_validator(
+            SubnetsValidator,
+            subnet_ids=self.login_nodes_subnet_ids + self.compute_subnet_ids + [self.head_node.networking.subnet_id],
+        )
+        self._register_validator(LoginNodesSchedulerValidator, scheduler=self.scheduling.scheduler)
+
+        for pool in self.login_nodes.pools:
+            # Check the LoginNodes CustomAMI must be an ami of the same os family and the same arch.
+            if pool.image and pool.image.custom_ami:
+                self._register_validator(AmiOsCompatibleValidator, os=self.image.os, image_id=pool.image.custom_ami)
+                self._register_validator(
+                    InstanceTypeBaseAMICompatibleValidator,
+                    instance_type=pool.instance_type,
+                    image=self.login_nodes_ami[pool.name],
+                )
+                self._register_validator(
+                    InstanceTypeOSCompatibleValidator,
+                    instance_type=pool.instance_type,
+                    os=self.image.os,
+                )
+            # Check pool instance compatability with NICE DCV.
+            if pool.dcv:
+                self._register_validator(
+                    DcvValidator,
+                    instance_type=pool.instance_type,
+                    dcv_enabled=pool.dcv.enabled,
+                    allowed_ips=pool.dcv.allowed_ips,
+                    port=pool.dcv.port,
+                    os=self.image.os,
+                    architecture=pool.architecture,
+                )
+                has_dcv_configured = True
+
+        if has_dcv_configured:
+            self._register_validator(FeatureRegionValidator, feature=Feature.DCV, region=self.region)
+
     def _register_validators(self, context: ValidatorContext = None):  # noqa: C901
         super()._register_validators(context)
         self._register_validator(
@@ -2946,33 +3014,8 @@ def _register_validators(self, context: ValidatorContext = None):  # noqa: C901
             queues=self.scheduling.queues,
         )
 
-        # Check if all subnets(head node, Login nodes, compute nodes) are in the same VPC and support DNS.
         if self.login_nodes:
-            self._register_validator(
-                SubnetsValidator,
-                subnet_ids=self.login_nodes_subnet_ids
-                + self.compute_subnet_ids
-                + [self.head_node.networking.subnet_id],
-            )
-
-        if self.login_nodes:
-            self._register_validator(LoginNodesSchedulerValidator, scheduler=self.scheduling.scheduler)
-
-        # Check the LoginNodes CustomAMI must be an ami of the same os family and the same arch.
-        if self.login_nodes:
-            for pool in self.login_nodes.pools:
-                if pool.image and pool.image.custom_ami:
-                    self._register_validator(AmiOsCompatibleValidator, os=self.image.os, image_id=pool.image.custom_ami)
-                    self._register_validator(
-                        InstanceTypeBaseAMICompatibleValidator,
-                        instance_type=pool.instance_type,
-                        image=self.login_nodes_ami[pool.name],
-                    )
-                    self._register_validator(
-                        InstanceTypeOSCompatibleValidator,
-                        instance_type=pool.instance_type,
-                        os=self.image.os,
-                    )
+            self._register_login_node_validators()
 
         if self.scheduling.settings and self.scheduling.settings.dns and self.scheduling.settings.dns.hosted_zone_id:
             self._register_validator(

cli/src/pcluster/models/cluster.py

@@ -713,6 +713,15 @@ def head_node_instance(self) -> ClusterInstance:
         else:
             raise ClusterActionError("Unable to retrieve head node information.")
 
+    @property
+    def login_node_instances(self) -> List[ClusterInstance]:
+        """Get login node instances."""
+        instances, _ = self.describe_instances(node_type=NodeType.LOGIN_NODE)
+        if instances:
+            return instances
+        else:
+            raise ClusterActionError("Unable to retrieve login node information.")
+
     def _get_instance_filters(self, node_type: NodeType, queue_name: str = None):
         filters = [
             {"Name": f"tag:{PCLUSTER_CLUSTER_NAME_TAG}", "Values": [self.stack_name]},

cli/src/pcluster/schemas/cluster_schema.py

@@ -1434,6 +1434,7 @@ class LoginNodesPoolSchema(BaseSchema):
         ),
         metadata={"update_policy": UpdatePolicy.SUPPORTED},
     )
+    dcv = fields.Nested(DcvSchema, metadata={"update_policy": UpdatePolicy.UNSUPPORTED})
     ssh = fields.Nested(LoginNodesSshSchema, metadata={"update_policy": UpdatePolicy.LOGIN_NODES_STOP})
     custom_actions = fields.Nested(LoginNodesCustomActionsSchema, metadata={"update_policy": UpdatePolicy.IGNORED})
     iam = fields.Nested(LoginNodesIamSchema, metadata={"update_policy": UpdatePolicy.LOGIN_NODES_STOP})

cli/src/pcluster/templates/cdk_builder_utils.py

@@ -958,6 +958,21 @@ def _build_policy(self) -> List[iam.PolicyStatement]:
                     )
                 ],
             ),
+            iam.PolicyStatement(
+                sid="DcvLicense",
+                actions=[
+                    "s3:GetObject",
+                ],
+                effect=iam.Effect.ALLOW,
+                resources=[
+                    self._format_arn(
+                        service="s3",
+                        resource="dcv-license.{0}/*".format(Stack.of(self).region),
+                        region="",
+                        account="",
+                    )
+                ],
+            ),
         ]
 
 

cli/src/pcluster/templates/cluster_stack.py

@@ -887,11 +887,23 @@ def _get_source_ingress_rule(self, setting):
             return ec2.CfnSecurityGroup.IngressProperty(ip_protocol="tcp", from_port=22, to_port=22, cidr_ip=setting)
 
     def _add_login_nodes_security_group(self):
+        # TODO review this once we allow more pools to be defined in the LoginNodes section
         login_nodes_security_group_ingress = [
             # SSH access
-            # TODO review this once we allow more pools to be defined in the LoginNodes section
             self._get_source_ingress_rule(self.config.login_nodes.pools[0].ssh.allowed_ips)
         ]
+
+        if self.config.login_nodes.has_dcv_enabled:
+            login_nodes_security_group_ingress.append(
+                # DCV access
+                ec2.CfnSecurityGroup.IngressProperty(
+                    ip_protocol="tcp",
+                    from_port=self.config.login_nodes.pools[0].dcv.port,
+                    to_port=self.config.login_nodes.pools[0].dcv.port,
+                    cidr_ip=self.config.login_nodes.pools[0].dcv.allowed_ips,
+                )
+            )
+
         return ec2.CfnSecurityGroup(
             self.stack,
             "LoginNodesSecurityGroup",

cli/src/pcluster/templates/login_nodes_stack.py

@@ -256,6 +256,8 @@ def _add_login_nodes_pool_launch_template(self):
                     "head_node_private_ip": self._head_eni.attr_primary_private_ip_address,
                     "dns_domain": (str(self._cluster_hosted_zone.name) if self._cluster_hosted_zone else ""),
                     "hosted_zone": (str(self._cluster_hosted_zone.ref) if self._cluster_hosted_zone else ""),
+                    "dcv_enabled": "login_node" if self._pool.has_dcv_enabled else "false",
+                    "dcv_port": self._pool.dcv.port if self._pool.dcv else "NONE",
                     "log_group_name": self._log_group.log_group_name,
                     "log_rotation_enabled": "true" if self._config.is_log_rotation_enabled else "false",
                     "pool_name": self._pool.name,

cli/tests/pcluster/cli/test_dcv_connect/TestDcvConnectCommand/test_helper/pcluster-help.txt

@@ -1,8 +1,9 @@
 usage: pcluster dcv-connect [-h] [--debug] [-r REGION] -n CLUSTER_NAME
                             [--key-path KEY_PATH] [--show-url]
+                            [--login-node-ip LOGIN_NODE_IP]
 
-Permits to connect to the head node through an interactive session by using
-NICE DCV.
+Permits connection to the head or login nodes through an interactive session
+by using NICE DCV.
 
 options:
   -h, --help            show this help message and exit
@@ -13,3 +14,5 @@ options:
                         Name of the cluster to connect to
   --key-path KEY_PATH   Key path of the SSH key to use for the connection
   --show-url            Print URL and exit
+  --login-node-ip LOGIN_NODE_IP
+                        IP address of a login node to connect to

cli/tests/pcluster/cli/test_entrypoint/TestParallelClusterCli/test_helper/pcluster-help.txt

@@ -49,8 +49,8 @@ COMMANDS:
     list-official-images
                         List Official ParallelCluster AMIs.
     configure           Start the AWS ParallelCluster configuration.
-    dcv-connect         Permits to connect to the head node through an
-                        interactive session by using NICE DCV.
+    dcv-connect         Permits connection to the head or login nodes through
+                        an interactive session by using NICE DCV.
     export-cluster-logs
                         Export the logs of the cluster to a local tar.gz
                         archive by passing through an Amazon S3 Bucket.