Merge branch 'microsoft:main' into paull/ippower

lisa/sut_orchestrator/azure/common.py

@@ -523,7 +523,24 @@ class VhdSchema(AzureImageSchema):
     vmgs_path: Optional[str] = None
 
     def load_from_platform(self, platform: "AzurePlatform") -> None:
-        return
+        # There are no platform tags to parse, but we can assume the
+        # security profile based on the presence of a VMGS path.
+        if self.vmgs_path:
+            self.security_profile = search_space.SetSpace(
+                True,
+                [
+                    SecurityProfileType.CVM,
+                    SecurityProfileType.Stateless,
+                ],
+            )
+        else:
+            self.security_profile = search_space.SetSpace(
+                True,
+                [
+                    SecurityProfileType.Standard,
+                    SecurityProfileType.SecureBoot,
+                ],
+            )
 
 
 @dataclass_json()

lisa/sut_orchestrator/azure/transformers.py

@@ -116,6 +116,7 @@ class VhdTransformer(Transformer):
     """
 
     __url_name = "url"
+    __vmgs_url_name = "vmgs_url"
 
     @classmethod
     def type_name(cls) -> str:
@@ -127,7 +128,7 @@ def type_schema(cls) -> Type[schema.TypedSchema]:
 
     @property
     def _output_names(self) -> List[str]:
-        return [self.__url_name]
+        return [self.__url_name, self.__vmgs_url_name]
 
     def _internal_run(self) -> Dict[str, Any]:
         runbook: VhdTransformerSchema = self.runbook
@@ -157,11 +158,11 @@ def _internal_run(self) -> Dict[str, Any]:
 
         virtual_machine = get_vm(platform, node)
 
-        vhd_location = self._export_vhd(platform, virtual_machine)
+        outputs = self._export_vhd(platform, virtual_machine)
 
         self._restore_vm(platform, virtual_machine, node)
 
-        return {self.__url_name: vhd_location}
+        return outputs
 
     def _prepare_virtual_machine(self, node: RemoteNode) -> None:
         runbook: VhdTransformerSchema = self.runbook
@@ -178,22 +179,38 @@ def _prepare_virtual_machine(self, node: RemoteNode) -> None:
         startstop = node.features[StartStop]
         startstop.stop()
 
-    def _export_vhd(self, platform: AzurePlatform, virtual_machine: Any) -> str:
+    def _export_vhd(
+        self, platform: AzurePlatform, virtual_machine: Any
+    ) -> Dict[str, str]:
         runbook: VhdTransformerSchema = self.runbook
         compute_client = get_compute_client(platform)
 
         # generate sas url from os disk, so it can be copied.
         self._log.debug("generating sas url...")
         location = virtual_machine.location
         os_disk_name = virtual_machine.storage_profile.os_disk.name
+
+        has_vmgs = (
+            virtual_machine.storage_profile.os_disk.managed_disk.security_profile
+            is not None
+        )
+
         operation = compute_client.disks.begin_grant_access(
             resource_group_name=runbook.resource_group_name,
             disk_name=os_disk_name,
-            grant_access_data=GrantAccessData(access="Read", duration_in_seconds=86400),
+            grant_access_data=GrantAccessData(
+                access="Read",
+                duration_in_seconds=86400,
+                get_secure_vm_guest_state_sas=has_vmgs,
+            ),
         )
         wait_operation(operation)
-        sas_url = operation.result().access_sas
+        result = operation.result()
+        sas_url = result.access_sas
+        vmgs_sas_url = result.security_data_access_sas or ""
         assert sas_url, "cannot get sas_url from os disk"
+        assert isinstance(sas_url, str), "sas_url is not a string"
+        assert isinstance(vmgs_sas_url, str), "vmgs_sas_url is not a string"
 
         self._log.debug("getting or creating storage account and container...")
         # get vhd container
@@ -219,16 +236,35 @@ def _export_vhd(self, platform: AzurePlatform, virtual_machine: Any) -> str:
         )
 
         if runbook.custom_blob_name:
-            path = runbook.custom_blob_name
+            vhd_path = runbook.custom_blob_name
+        else:
+            vhd_path = _generate_vhd_path(container_client, runbook.file_name_part)
+        vhd_url_path = f"{container_client.url}/{vhd_path}"
+        vhd_blob_client = container_client.get_blob_client(vhd_path)
+        vhd_blob_client.start_copy_from_url(
+            sas_url, metadata=None, incremental_copy=False
+        )
+
+        if vmgs_sas_url:
+            vmgs_path = runbook.custom_blob_name.replace(".vhd", "_vmgs.vhd")
+            assert vmgs_path != vhd_path, (
+                "vmgs url path is the same as vhd url path. "
+                "Make sure the custom_blob_name ends in .vhd"
+            )
+            vmgs_url_path = f"{container_client.url}/{vmgs_path}"
+            vmgs_blob_client = container_client.get_blob_client(vmgs_path)
+            vmgs_blob_client.start_copy_from_url(
+                vmgs_sas_url, metadata=None, incremental_copy=False
+            )
+            wait_copy_blob(vmgs_blob_client, vmgs_url_path, self._log)
         else:
-            path = _generate_vhd_path(container_client, runbook.file_name_part)
-        vhd_path = f"{container_client.url}/{path}"
-        blob_client = container_client.get_blob_client(path)
-        blob_client.start_copy_from_url(sas_url, metadata=None, incremental_copy=False)
+            vmgs_url_path = ""
 
-        wait_copy_blob(blob_client, vhd_path, self._log)
+        # Wait for VHD to copy after copying VMGS
+        # so they can copy in parallel
+        wait_copy_blob(vhd_blob_client, vhd_url_path, self._log)
 
-        return vhd_path
+        return {self.__url_name: vhd_url_path, self.__vmgs_url_name: vmgs_url_path}
 
     def _restore_vm(
         self, platform: AzurePlatform, virtual_machine: Any, node: Node

lisa/tools/lsvmbus.py

@@ -5,7 +5,7 @@
 from lisa.base_tools.wget import Wget
 from lisa.executable import Tool
 from lisa.operating_system import Redhat, Suse, Ubuntu
-from lisa.tools import Dmesg
+from lisa.tools import Dmesg, Echo
 from lisa.util import LisaException, find_groups_in_lines
 
 from .python import Python
@@ -189,9 +189,16 @@ def _check_exists(self) -> bool:
 
     def _install_from_src(self) -> None:
         wget_tool = self.node.tools[Wget]
-        file_path = wget_tool.get(
-            self._lsvmbus_repo, "$HOME/.local/bin", executable=True
-        )
+        echo = self.node.tools[Echo]
+        tool_path = echo.run(
+            "$HOME/.local/bin",
+            shell=True,
+            expected_exit_code=0,
+            expected_exit_code_failure_message=(
+                "failed to get $HOME/.local/bin via echo"
+            ),
+        ).stdout
+        file_path = wget_tool.get(self._lsvmbus_repo, tool_path, executable=True)
         self._command = file_path
 
     def install(self) -> bool: