Add support for custom log file configuration in CW agent

Signed-off-by: chenwany <[email protected]>

Author: chenwany

cookbooks/aws-parallelcluster-config/files/default/cloudwatch_agent/write_cloudwatch_agent_json.py

@@ -11,6 +11,7 @@
 import os
 import socket
 
+import yaml
 
 AWS_CLOUDWATCH_CFG_PATH = '/opt/aws/amazon-cloudwatch-agent/etc/amazon-cloudwatch-agent.json'
 
@@ -39,6 +40,10 @@ def parse_args():
                         required=True,
                         choices=['slurm', 'awsbatch', 'plugin'],
                         help='Scheduler')
+    parser.add_argument('--cluster-config-path',
+                        required=False,
+                        help='Cluster configuration path',
+                        )
     return parser.parse_args()
 
 
@@ -132,6 +137,44 @@ def select_logs(configs, args):
     return selected_configs
 
 
+def get_node_roles(scheudler_plugin_node_roles):
+    node_type_roles_map = {"ALL": ["ComputeFleet", "HeadNode"], "HEAD": ["HeadNode"], "COMPUTE": ["ComputeFleet"]}
+    return node_type_roles_map.get(scheudler_plugin_node_roles)
+
+
+def load_config(cluster_config_path):
+    with open(cluster_config_path) as input_file:
+        return yaml.load(input_file, Loader=yaml.SafeLoader)
+
+
+def add_scheduler_plugin_log(config_data, cluster_config_path):
+    """Add custom log files to config data if log files specified in scheduler plugin."""
+    cluster_config = load_config(cluster_config_path)
+    if (
+            get_dict_value(cluster_config, "Scheduling.SchedulerSettings.SchedulerDefinition.Monitoring.Logs.Files")
+            and get_dict_value(cluster_config, "Scheduling.Scheduler") == "plugin"
+    ):
+        log_files = get_dict_value(
+            cluster_config, "Scheduling.SchedulerSettings.SchedulerDefinition.Monitoring.Logs.Files"
+        )
+        for log_file in log_files:
+            # Add log config
+            log_config = {
+                "timestamp_format_key": log_file.get("LogStreamName"),
+                "file_path": log_file.get("FilePath"),
+                "log_stream_name": log_file.get("LogStreamName"),
+                "schedulers": ["plugin"],
+                "platforms": ["centos", "ubuntu", "amazon"],
+                "node_roles": get_node_roles(log_file.get("NodeType")),
+                "feature_conditions": [],
+            }
+            config_data["log_configs"].append(log_config)
+
+            # Add timestamp formats
+            config_data["timestamp_formats"][log_file.get("LogStreamName")] = log_file.get("TimestampFormat")
+    return config_data
+
+
 def add_timestamps(configs, timestamps_dict):
     """For each config, set its timestamp_format field based on its timestamp_format_key field."""
     for config in configs:
@@ -159,10 +202,21 @@ def create_config(log_configs):
     }
 
 
+def get_dict_value(value, attributes, default=None):
+    """Get key value from dictionary and return default if the key does not exist."""
+    for key in attributes.split("."):
+        value = value.get(key, None)
+        if value is None:
+            return default
+    return value
+
+
 def main():
     """Create cloudwatch agent config file."""
     args = parse_args()
     config_data = read_data(args.config)
+    if args.cluster_config_path:
+        config_data = add_scheduler_plugin_log(config_data, args.cluster_config_path)
     log_configs = select_logs(config_data['log_configs'], args)
     log_configs = add_timestamps(log_configs, config_data['timestamp_formats'])
     log_configs = add_log_group_name_params(args.log_group, log_configs)

cookbooks/aws-parallelcluster-config/recipes/cloudwatch_agent.rb

@@ -73,9 +73,12 @@
     'CONFIG_DATA_PATH' => config_data_path
   )
   not_if { ::File.exist?('/opt/aws/amazon-cloudwatch-agent/etc/amazon-cloudwatch-agent.json') }
+
+  cluster_config_path = node['cluster']['scheduler'] == 'plugin' ? "--cluster-config-path #{node['cluster']['cluster_config_path']}" : ""
+
   command "#{node.default['cluster']['cookbook_virtualenv_path']}/bin/python #{config_script_path} "\
-          "--platform #{node['platform']} --config $CONFIG_DATA_PATH --log-group $LOG_GROUP_NAME "\
-          "--scheduler $SCHEDULER --node-role $NODE_ROLE"
+      "--platform #{node['platform']} --config $CONFIG_DATA_PATH --log-group $LOG_GROUP_NAME "\
+      "--scheduler $SCHEDULER --node-role $NODE_ROLE #{cluster_config_path}"
 end
 
 execute "cloudwatch-agent-start" do

cookbooks/aws-parallelcluster-config/recipes/init.rb

@@ -52,16 +52,16 @@
   mode "0755"
 end
 
+include_recipe "aws-parallelcluster-config::mount_shared" if node['cluster']['node_type'] == "ComputeFleet"
+
+fetch_config 'Fetch and load cluster configs' unless node['cluster']['scheduler'] == 'awsbatch'
+
 # Install cloudwatch, write configuration and start it.
 include_recipe "aws-parallelcluster-config::cloudwatch_agent"
 
 # Configure additional Networking Interfaces (if present)
 include_recipe "aws-parallelcluster-config::network_interfaces" unless virtualized?
 
-include_recipe "aws-parallelcluster-config::mount_shared" if node['cluster']['node_type'] == "ComputeFleet"
-
-fetch_config 'Fetch and load cluster configs' unless node['cluster']['scheduler'] == 'awsbatch'
-
 include_recipe "aws-parallelcluster-slurm::init" if node['cluster']['scheduler'] == 'slurm'
 include_recipe "aws-parallelcluster-scheduler-plugin::init" if node['cluster']['scheduler'] == 'plugin'