Report VMSize via Heartbeat telemetry event (Azure#2462)

azurelinuxagent/ga/update.py

@@ -33,6 +33,7 @@
 
 import azurelinuxagent.common.conf as conf
 import azurelinuxagent.common.logger as logger
+from azurelinuxagent.common.protocol.imds import get_imds_client
 import azurelinuxagent.common.utils.fileutil as fileutil
 import azurelinuxagent.common.utils.restutil as restutil
 import azurelinuxagent.common.utils.textutil as textutil
@@ -161,6 +162,9 @@ def __init__(self):
  self._heartbeat_id = str(uuid.uuid4()).upper()
  self._heartbeat_counter = 0
 
+ # VM Size is reported via the heartbeat, default it here.
+ self._vm_size = None
+
  # these members are used to avoid reporting errors too frequently
  self._heartbeat_update_goal_state_error_count = 0
  self._last_try_update_goal_state_failed = False
@@ -407,6 +411,25 @@ def run(self, debug=False):
  self._shutdown()
  sys.exit(0)
 
+ def _get_vm_size(self, protocol):
+ """
+ Including VMSize is meant to capture the architecture of the VM (i.e. arm64 VMs will
+ have arm64 included in their vmsize field and amd64 will have no architecture indicated).
+ """
+ if self._vm_size is None:
+
+ imds_client = get_imds_client(protocol.get_endpoint())
+
+ try:
+ imds_info = imds_client.get_compute()
+ self._vm_size = imds_info.vmSize
+ except Exception as e:
+ err_msg = "Attempts to retrieve VM size information from IMDS are failing: {0}".format(textutil.format_exception(e))
+ logger.periodic_warn(logger.EVERY_SIX_HOURS, "[PERIODIC] {0}".format(err_msg))
+ return "unknown"
+
+ return self._vm_size
+
  def _check_daemon_running(self, debug):
  # Check that the parent process (the agent's daemon) is still running
  if not debug and self._is_orphaned:
@@ -1152,10 +1175,13 @@ def _send_heartbeat_telemetry(self, protocol):
  if datetime.utcnow() >= (self._last_telemetry_heartbeat + UpdateHandler.TELEMETRY_HEARTBEAT_PERIOD):
  dropped_packets = self.osutil.get_firewall_dropped_packets(protocol.get_endpoint())
  auto_update_enabled = 1 if conf.get_autoupdate_enabled() else 0
+ # Include VMSize in the heartbeat message because the kusto table does not have 
+ # a separate column for it (or architecture).
+ vmsize = self._get_vm_size(protocol)
 
- telemetry_msg = "{0};{1};{2};{3};{4}".format(self._heartbeat_counter, self._heartbeat_id, dropped_packets,
+ telemetry_msg = "{0};{1};{2};{3};{4};{5}".format(self._heartbeat_counter, self._heartbeat_id, dropped_packets,
  self._heartbeat_update_goal_state_error_count,
- auto_update_enabled)
+ auto_update_enabled, vmsize)
  debug_log_msg = "[DEBUG HeartbeatCounter: {0};HeartbeatId: {1};DroppedPackets: {2};" \
  "UpdateGSErrors: {3};AutoUpdate: {4}]".format(self._heartbeat_counter,
  self._heartbeat_id, dropped_packets,

tests/ga/test_update.py

@@ -20,6 +20,7 @@
 
 from datetime import datetime, timedelta
 from threading import currentThread
+from azurelinuxagent.common.protocol.imds import ComputeInfo
 from tests.common.osutil.test_default import TestOSUtil
 import azurelinuxagent.common.osutil.default as osutil
 
@@ -1565,6 +1566,48 @@ def test_telemetry_heartbeat_creates_event(self, patch_add_event, patch_info, *_
  self.assertTrue(any(call_args[0] == "[HEARTBEAT] Agent {0} is running as the goal state agent {1}"
  for call_args in patch_info.call_args), "The heartbeat was not written to the agent's log")
 
+ @patch("azurelinuxagent.ga.update.add_event")
+ @patch("azurelinuxagent.common.protocol.imds.ImdsClient")
+ def test_telemetry_heartbeat_retries_failed_vm_size_fetch(self, mock_imds_factory, patch_add_event, *_):
+
+ def validate_single_heartbeat_event_matches_vm_size(vm_size):
+ heartbeat_event_kwargs = [
+ kwargs for _, kwargs in patch_add_event.call_args_list
+ if kwargs.get('op', None) == WALAEventOperation.HeartBeat
+ ]
+
+ self.assertEqual(1, len(heartbeat_event_kwargs), "Expected exactly one HeartBeat event, got {0}"\
+ .format(heartbeat_event_kwargs))
+
+ telemetry_message = heartbeat_event_kwargs[0].get("message", "")
+ self.assertTrue(telemetry_message.endswith(vm_size),
+ "Expected HeartBeat message ('{0}') to end with the test vmSize value, {1}."\
+ .format(telemetry_message, vm_size))
+
+ with mock_wire_protocol(mockwiredata.DATA_FILE) as mock_protocol:
+ update_handler = get_update_handler()
+ update_handler.protocol_util.get_protocol = Mock(return_value=mock_protocol)
+
+ # Zero out the _vm_size parameter for test resiliency
+ update_handler._vm_size = None
+
+ mock_imds_client = mock_imds_factory.return_value = Mock()
+
+ # First force a vmSize retrieval failure
+ mock_imds_client.get_compute.side_effect = HttpError(msg="HTTP Test Failure")
+ update_handler._last_telemetry_heartbeat = datetime.utcnow() - timedelta(hours=1)
+ update_handler._send_heartbeat_telemetry(mock_protocol)
+
+ validate_single_heartbeat_event_matches_vm_size("unknown")
+ patch_add_event.reset_mock()
+
+ # Now provide a vmSize
+ mock_imds_client.get_compute = lambda: ComputeInfo(vmSize="TestVmSizeValue")
+ update_handler._last_telemetry_heartbeat = datetime.utcnow() - timedelta(hours=1)
+ update_handler._send_heartbeat_telemetry(mock_protocol)
+
+ validate_single_heartbeat_event_matches_vm_size("TestVmSizeValue")
+
  @staticmethod
  def _get_test_ext_handler_instance(protocol, name="OSTCExtensions.ExampleHandlerLinux", version="1.0.0"):
  eh = Extension(name=name)