Adds snapshot creation and recovery support for linked and full VMs

src/uds/services/Proxmox/helpers.py

@@ -41,11 +41,13 @@
 
 logger = logging.getLogger(__name__)
 
+
 def get_provider(parameters: typing.Any) -> 'ProxmoxProvider':
     return typing.cast(
         'ProxmoxProvider', models.Provider.objects.get(uuid=parameters['prov_uuid']).get_instance()
     )
 
+
 def get_storage(parameters: typing.Any) -> types.ui.CallbackResultType:
     logger.debug('Parameters received by getResources Helper: %s', parameters)
     provider = get_provider(parameters)

src/uds/services/Proxmox/jobs.py

@@ -56,8 +56,8 @@ class ProxmoxDeferredRemoval(jobs.Job):
     frecuency = 60 * 3  # Once every NN minutes
     friendly_name = 'Proxmox removal'
     counter = 0
-    
-    def get_vmid_stored_data_from(self, data: bytes) -> typing.Tuple[int, bool]:
+
+    def get_vmid_stored_data_from(self, data: bytes) -> tuple[int, bool]:
         vmdata = data.decode()
         if ':' in vmdata:
             vmid, try_graceful_shutdown_s = vmdata.split(':')
@@ -66,7 +66,6 @@ def get_vmid_stored_data_from(self, data: bytes) -> typing.Tuple[int, bool]:
             vmid = vmdata
             try_graceful_shutdown = False
         return int(vmid), try_graceful_shutdown
-
 
     # @staticmethod
     # def remove(provider_instance: 'provider.ProxmoxProvider', vmid: int, try_graceful_shutdown: bool) -> None:
@@ -79,19 +78,19 @@ def get_vmid_stored_data_from(self, data: bytes) -> typing.Tuple[int, bool]:
     #         vminfo = provider_instance.get_machine_info(vmid)
     #         if vminfo.status == 'running':
     #             if try_graceful_shutdown:
-    #                 # If running vm,  simply try to shutdown 
+    #                 # If running vm,  simply try to shutdown
     #                 provider_instance.shutdown_machine(vmid)
     #                 # Store for later removal
-    #             else: 
+    #             else:
     #                 # If running vm,  simply stops it and wait for next
     #                 provider_instance.stop_machine(vmid)
-                    
+
     #             store_for_deferred_removal()
     #             return
 
     #         provider_instance.remove_machine(vmid)  # Try to remove, launch removal, but check later
     #         store_for_deferred_removal()
-            
+
     #     except client.ProxmoxNotFound:
     #         return  # Machine does not exists
     #     except Exception as e:
@@ -132,7 +131,7 @@ def run(self) -> None:
                 vmid, _try_graceful_shutdown = self.get_vmid_stored_data_from(data[1])
                 # In fact, here, _try_graceful_shutdown is not used, but we keep it for mayby future use
                 # The soft shutdown has already being initiated by the remove method
-                   
+
                 try:
                     vm_info = instance.api.get_vm_info(vmid)
                     logger.debug('Found %s for removal %s', vmid, data)

src/uds/services/Proxmox/proxmox/client.py

@@ -311,9 +311,7 @@ def list_node_vgpus(self, node: str, **kwargs: typing.Any) -> list[types.VGPUInf
             for gpu in self.do_get(f'nodes/{node}/hardware/pci/{device}/mdev', node=node)['data']
         ]
 
-    def node_has_vgpus_available(
-        self, node: str, vgpu_type: str | None, **kwargs: typing.Any
-    ) -> bool:
+    def node_has_vgpus_available(self, node: str, vgpu_type: str | None, **kwargs: typing.Any) -> bool:
         return any(
             gpu.available and (vgpu_type is None or gpu.type == vgpu_type) for gpu in self.list_node_vgpus(node)
         )
@@ -383,9 +381,18 @@ def clone_vm(
         # Ensure exists target pool, (id is in fact the name of the pool)
         if target_pool and not any(p.id == target_pool for p in self.list_pools()):
             raise exceptions.ProxmoxDoesNotExists(f'Pool "{target_pool}" does not exist')
-        
+
         logger.debug('Cloning VM %s to %s', vmid, new_vmid)
-        logger.debug('Parameters: %s %s %s %s %s %s %s', name, description, as_linked_clone, target_node, target_storage, target_pool, must_have_vgpus)
+        logger.debug(
+            'Parameters: %s %s %s %s %s %s %s',
+            name,
+            description,
+            as_linked_clone,
+            target_node,
+            target_storage,
+            target_pool,
+            must_have_vgpus,
+        )
 
         src_node = vminfo.node
 
@@ -403,13 +410,13 @@ def clone_vm(
                 target_node = node.name
             else:
                 target_node = src_node
-                
-        logger.debug('Target node: %s', target_node)    
+
+        logger.debug('Target node: %s', target_node)
 
         # Ensure exists target node
         if not any(n.name == target_node for n in self.get_cluster_info().nodes):
             raise exceptions.ProxmoxDoesNotExists(f'Node "{target_node}" does not exist')
-        
+
         logger.debug('VM info: %s', vminfo)
         logger.debug('Type of vminfo.vgpu_type: %s', type(vminfo.vgpu_type))
         logger.debug('Value of vminfo.vgpu_type: %s', vminfo.vgpu_type)
@@ -498,9 +505,7 @@ def disable_vm_ha(self, vmid: int) -> None:
         except Exception:
             logger.exception('removeFromHA')
 
-    def set_vm_protection(
-        self, vmid: int, *, node: str | None = None, protection: bool = False
-    ) -> None:
+    def set_vm_protection(self, vmid: int, *, node: str | None = None, protection: bool = False) -> None:
         params: list[tuple[str, str]] = [
             ('protection', str(int(protection))),
         ]
@@ -537,7 +542,7 @@ def get_guest_ip_address(
             if found_ips:
                 sorted_ips = sorted(found_ips, key=lambda x: ':' in x)
                 return sorted_ips[0]
-            
+
         except Exception as e:
             logger.warning('Error getting guest ip address for machine %s: %s', vmid, e)
             raise exceptions.ProxmoxError(f'No ip address for vm {vmid}: {e}')
@@ -558,9 +563,7 @@ def list_snapshots(self, vmid: int, node: str | None = None) -> list[types.Snaps
         except Exception:
             return []  # If we can't get snapshots, just return empty list
 
-    def get_current_vm_snapshot(
-        self, vmid: int, node: str | None = None
-    ) -> types.SnapshotInfo | None:
+    def get_current_vm_snapshot(self, vmid: int, node: str | None = None) -> types.SnapshotInfo | None:
         return (
             sorted(
                 filter(lambda x: x.snaptime, self.list_snapshots(vmid, node)),
@@ -709,9 +712,7 @@ def get_vm_info(self, vmid: int, node: str | None = None, **kwargs: typing.Any)
         raise exceptions.ProxmoxNotFound(f'VM {vmid} not found')
 
     @cached('vmc', consts.CACHE_VM_INFO_DURATION, key_helper=caching_key_helper)
-    def get_vm_config(
-        self, vmid: int, node: str | None = None, **kwargs: typing.Any
-    ) -> types.VMConfiguration:
+    def get_vm_config(self, vmid: int, node: str | None = None, **kwargs: typing.Any) -> types.VMConfiguration:
         node = node or self.get_vm_info(vmid).node
         return types.VMConfiguration.from_dict(
             self.do_get(f'nodes/{node}/qemu/{vmid}/config', node=node)['data']
@@ -771,9 +772,7 @@ def suspend_vm(self, vmid: int, node: str | None = None) -> types.ExecResult:
             self.do_post(f'nodes/{node}/qemu/{vmid}/status/suspend', data=[('todisk', '1')], node=node)
         )
 
-    def shutdown_vm(
-        self, vmid: int, node: str | None = None, timeout: int | None = None
-    ) -> types.ExecResult:
+    def shutdown_vm(self, vmid: int, node: str | None = None, timeout: int | None = None) -> types.ExecResult:
         # if exitstatus is "OK" or contains "already running", all is fine
         node = node or self.get_vm_info(vmid).node
         return types.ExecResult.from_dict(

src/uds/services/Proxmox/proxmox/consts.py

@@ -40,4 +40,4 @@
 CACHE_INFO_DURATION: typing.Final[int] = consts.cache.SHORT_CACHE_TIMEOUT
 CACHE_VM_INFO_DURATION: typing.Final[int] = consts.cache.SHORTEST_CACHE_TIMEOUT
 # Cache duration is 3 minutes, so this is 60 mins * 24 = 1 day (default)
-CACHE_DURATION_LONG: typing.Final[int] = consts.cache.EXTREME_CACHE_TIMEOUT
+CACHE_DURATION_LONG: typing.Final[int] = consts.cache.EXTREME_CACHE_TIMEOUT

src/uds/services/Proxmox/proxmox/exceptions.py

@@ -43,9 +43,11 @@ class ProxmoxConnectionError(ProxmoxError, exceptions.services.generics.Retryabl
 class ProxmoxAuthError(ProxmoxError, exceptions.services.generics.FatalError):
     pass
 
+
 class ProxmoxDoesNotExists(ProxmoxError):
     pass
 
+
 class ProxmoxNotFound(ProxmoxDoesNotExists, exceptions.services.generics.NotFoundError):
     pass
 

src/uds/services/Proxmox/service.py

@@ -190,7 +190,8 @@ def initialize(self, values: 'types.core.ValuesType') -> None:
                 )
             except StopIteration:
                 raise exceptions.ui.ValidationError(_('Selected storage not found on Proxmox'))
-            if not storage.supports_linked_clone() and self.use_full_clone.value:
+            if not storage.supports_linked_clone() and not self.use_full_clone.value:
+                # if storage does not support linked clones and user wants to use linked clones, raise an error
                 raise exceptions.ui.ValidationError(_('Linked clones are not allowed on the storage'))
 
             # if int(self.memory.value) < 128:
@@ -339,3 +340,30 @@ def is_deleted(self, vmid: str) -> bool:
             return False
         except Exception:
             return True
+
+    def snapshot_creation(self, userservice_instance: DynamicUserService) -> None:
+        userservice_instance = typing.cast(ProxmoxUserserviceLinked, userservice_instance)
+        vmid = int(userservice_instance._vmid)
+        logger.debug('Using snapshots')
+        # If no snapshot exists for this vm, try to create one for it on background
+        # Lauch an snapshot. We will not wait for it to finish, but instead let it run "as is"
+        try:
+            if not self.provider().api.get_current_vm_snapshot(vmid):
+                logger.debug('No current snapshot')
+                self.provider().api.create_snapshot(
+                    vmid,
+                    name='UDS_Snapshot',
+                )
+        except Exception as e:
+            self.do_log(types.log.LogLevel.WARNING, 'Could not create SNAPSHOT for this VM. ({})'.format(e))
+
+    def snapshot_recovery(self, userservice_instance: DynamicUserService) -> None:
+        userservice_instance = typing.cast(ProxmoxUserserviceLinked, userservice_instance)
+        vmid = int(userservice_instance._vmid)
+        try:
+            # try to revert to snapshot
+            snapshot = self.provider().api.get_current_vm_snapshot(vmid)
+            if snapshot:
+                userservice_instance._store_task(self.provider().api.restore_snapshot(vmid, name=snapshot.name))
+        except Exception as e:
+            self.do_log(types.log.LogLevel.WARNING, 'Could not restore SNAPSHOT for this VM. ({})'.format(e))

tests/services/proxmox/fixtures.py

@@ -524,6 +524,7 @@ def inner(*args: typing.Any, **kwargs: typing.Any) -> T:
     'ha': HA_GROUPS[0],
     'try_soft_shutdown': False,
     'machine': VMINFO_LIST[0].id,
+    'use_full_clone': False,
     'datastore': STORAGES[0].storage,
     'gpu': VGPUS[0].type,
     'basename': 'base',
@@ -551,7 +552,7 @@ def create_client_mock() -> mock.Mock:
 @contextlib.contextmanager
 def patched_provider(
     **kwargs: typing.Any,
-) -> typing.Generator[provider.ProxmoxProvider, None, None]:
+) -> collections.abc.Generator[provider.ProxmoxProvider, None, None]:
     client = create_client_mock()
     provider = create_provider(**kwargs)
     provider._cached_api = client
@@ -567,7 +568,7 @@ def create_provider(**kwargs: typing.Any) -> provider.ProxmoxProvider:
 
     uuid_ = str(uuid.uuid4())
     return provider.ProxmoxProvider(
-        environment=environment.Environment.private_environment(uuid), values=values, uuid=uuid_
+        environment=environment.Environment.private_environment(uuid_), values=values, uuid=uuid_
     )
 
 
@@ -656,10 +657,11 @@ def create_userservice_linked(
     Create a linked user service
     """
     uuid_ = str(uuid.uuid4())
+    service = service or create_service_linked()
     return deployment_linked.ProxmoxUserserviceLinked(
         environment=environment.Environment.private_environment(uuid_),
-        service=service or create_service_linked(),
-        publication=publication or create_publication(),
+        service=service,
+        publication=publication or create_publication(service=service),
         uuid=uuid_,
     )
 

tests/services/proxmox/test_client.py

@@ -35,6 +35,7 @@
 import typing
 import logging
 import contextlib
+import collections.abc
 
 from uds.core import types as core_types
 
@@ -133,15 +134,15 @@ def _wait_for_task(self, exec_result: prox_types.ExecResult, timeout: int = 16)
     @contextlib.contextmanager
     def _create_test_vm(
         self,
-        vmid: typing.Optional[int] = None,
+        vmid: int | None = None,
         as_linked_clone: bool = False,
-        target_node: typing.Optional[str] = None,
-        target_storage: typing.Optional[str] = None,
-        target_pool: typing.Optional[str] = None,
-        must_have_vgpus: typing.Optional[bool] = None,
-    ) -> typing.Iterator[prox_types.VMInfo]:
+        target_node: str | None = None,
+        target_storage: str | None = None,
+        target_pool: str | None = None,
+        must_have_vgpus: bool | None = None,
+    ) -> collections.abc.Iterator[prox_types.VMInfo]:
         new_vmid = self._get_new_vmid()
-        res: typing.Optional[prox_types.VmCreationResult] = None
+        res: prox_types.VmCreationResult | None = None
         try:
             res = self.pclient.clone_vm(
                 vmid=vmid or self.test_vm.id,

tests/services/proxmox/test_serialization_deployment.py

@@ -30,6 +30,7 @@
 """
 Author: Adolfo Gómez, dkmaster at dkmon dot com
 """
+import collections.abc
 import pickle
 import typing
 
@@ -81,7 +82,7 @@
 
 TEST_QUEUE_NEW: typing.Final[list[types.services.Operation]] = [i.to_operation() for i in TEST_QUEUE]
 
-SERIALIZED_DEPLOYMENT_DATA: typing.Final[typing.Mapping[str, bytes]] = {
+SERIALIZED_DEPLOYMENT_DATA: typing.Final[collections.abc.Mapping[str, bytes]] = {
     'v1': b'v1\x01name\x01ip\x01mac\x01task\x01vmid\x01reason\x01' + pickle.dumps(TEST_QUEUE, protocol=0),
 }
 
@@ -100,7 +101,9 @@ def check(self, version: str, instance: Deployment) -> None:
 
     def test_marshaling(self) -> None:
         def _create_instance(unmarshal_data: 'bytes|None' = None) -> Deployment:
-            instance = fixtures.create_userservice_linked()
+            with fixtures.patched_provider() as provider:
+                service = fixtures.create_service_linked(provider=provider)
+            instance = fixtures.create_userservice_linked(service=service)
             if unmarshal_data:
                 instance.unmarshal(unmarshal_data)
             return instance
@@ -126,7 +129,9 @@ def _create_instance(unmarshal_data: 'bytes|None' = None) -> Deployment:
 
     def test_marshaling_queue(self) -> None:
         def _create_instance(unmarshal_data: 'bytes|None' = None) -> Deployment:
-            instance = fixtures.create_userservice_linked()
+            with fixtures.patched_provider() as provider:
+                service = fixtures.create_service_linked(provider=provider)
+            instance = fixtures.create_userservice_linked(service=service)
             instance.env.storage.save_pickled('queue', TEST_QUEUE)
             if unmarshal_data:
                 instance.unmarshal(unmarshal_data)

tests/services/proxmox/test_service_linked.py

@@ -42,7 +42,8 @@
 
 class TestProxmovLinkedService(UDSTestCase):
     def test_service_data(self) -> None:
-        service = fixtures.create_service_linked()
+        with fixtures.patched_provider() as provider:
+            service = fixtures.create_service_linked(provider=provider)
 
         self.assertEqual(service.pool.value, fixtures.SERVICE_LINKED_VALUES_DICT['pool'])
         self.assertEqual(service.ha.value, fixtures.SERVICE_LINKED_VALUES_DICT['ha'])

tests/services/proxmox/test_userservice_linked.py

@@ -296,7 +296,8 @@ def test_userservice_cancel(self) -> None:
                     api.stop_vm.assert_called()
 
     def test_userservice_basics(self) -> None:
-        with fixtures.patched_provider():
-            userservice = fixtures.create_userservice_linked()
+        with fixtures.patched_provider() as provider:
+            service = fixtures.create_service_linked(provider=provider)
+            userservice = fixtures.create_userservice_linked(service=service)
             userservice.set_ip('1.2.3.4')
             self.assertEqual(userservice.get_ip(), '1.2.3.4')