Test your AI Agents with the all-new Agent to Agent Testing Platform.Learn More

How to use _restore_vm method in lisa

Best Python code snippet using lisa_python

transformers.py

Source:transformers.py

...105        assert isinstance(node, RemoteNode)106        self._prepare_virtual_machine(node)107        virtual_machine = get_vm(platform, node)108        vhd_location = self._export_vhd(platform, virtual_machine)109        self._restore_vm(platform, virtual_machine, node)110        return {self.__url_name: vhd_location}111    def _prepare_virtual_machine(self, node: RemoteNode) -> None:112        runbook: VhdTransformerSchema = self.runbook113        if not runbook.public_address:114            runbook.public_address = node.public_address115        # prepare vm for exporting116        wa = node.tools[Waagent]117        node.execute("export HISTSIZE=0", shell=True)118        wa.deprovision()119        # stop the vm120        startstop = node.features[StartStop]121        startstop.stop()122    def _export_vhd(self, platform: AzurePlatform, virtual_machine: Any) -> str:123        runbook: VhdTransformerSchema = self.runbook124        compute_client = get_compute_client(platform)125        # generate sas url from os disk, so it can be copied.126        self._log.debug("generating sas url...")127        location = virtual_machine.location128        os_disk_name = virtual_machine.storage_profile.os_disk.name129        operation = compute_client.disks.begin_grant_access(130            resource_group_name=runbook.resource_group_name,131            disk_name=os_disk_name,132            grant_access_data=GrantAccessData(access="Read", duration_in_seconds=86400),133        )134        wait_operation(operation)135        sas_url = operation.result().access_sas136        assert sas_url, "cannot get sas_url from os disk"137        self._log.debug("getting or creating storage account and container...")138        # get vhd container139        if not runbook.storage_account_name:140            runbook.storage_account_name = get_storage_account_name(141                subscription_id=platform.subscription_id, location=location, type="t"142            )143        check_or_create_storage_account(144            credential=platform.credential,145            subscription_id=platform.subscription_id,146            account_name=runbook.storage_account_name,147            resource_group_name=runbook.shared_resource_group_name,148            location=location,149            log=self._log,150        )151        container_client = get_or_create_storage_container(152            credential=platform.credential,153            subscription_id=platform.subscription_id,154            account_name=runbook.storage_account_name,155            container_name=runbook.container_name,156            resource_group_name=runbook.shared_resource_group_name,157        )158        path = _generate_vhd_path(container_client, runbook.file_name_part)159        vhd_path = f"{container_client.url}/{path}"160        blob_client = container_client.get_blob_client(path)161        blob_client.start_copy_from_url(sas_url, metadata=None, incremental_copy=False)162        wait_copy_blob(blob_client, vhd_path, self._log)163        return vhd_path164    def _restore_vm(165        self, platform: AzurePlatform, virtual_machine: Any, node: Node166    ) -> None:167        runbook: VhdTransformerSchema = self.runbook168        self._log.debug("restoring vm...")169        # release the vhd export lock, so it can be started back170        compute_client = get_compute_client(platform)171        os_disk_name = virtual_machine.storage_profile.os_disk.name172        operation = compute_client.disks.begin_revoke_access(173            resource_group_name=runbook.resource_group_name,174            disk_name=os_disk_name,175        )176        wait_operation(operation)177        if runbook.restore:178            start_stop = node.features[StartStop]...

openmm.py

Source:openmm.py

1# SPDX-License-Identifier: MIT2# Copyright (c) 2020-2021: PySAGES contributors3# See LICENSE.md and CONTRIBUTORS.md at https://github.com/SSAGESLabs/PySAGES4from functools import partial5from typing import Callable6from jax import jit, numpy as np7from jax.dlpack import from_dlpack as asarray8from jax.lax import cond9from openmm_dlext import ContextView, DeviceType, Force10from pysages.backends.core import ContextWrapper11from pysages.backends.snapshot import (12    Box,13    HelperMethods,14    Snapshot,15    SnapshotMethods,16    build_data_querier,17    restore as _restore,18    restore_vm as _restore_vm,19)20from pysages.methods import SamplingMethod21from pysages.utils import copy, try_import22import importlib23import jax24import openmm_dlext as dlext25openmm = try_import("openmm", "simtk.openmm")26unit = try_import("openmm.unit", "simtk.unit")27class Sampler:28    def __init__(self, method_bundle, bias, callback: Callable, restore):29        snapshot, initialize, method_update = method_bundle30        self.snapshot = snapshot31        self.state = initialize()32        self._update = method_update33        self._restore = restore34        self.bias = bias35        self.callback = callback36    def update(self, timestep=0):37        self.state = self._update(self.snapshot, self.state)38        self.bias(self.snapshot, self.state)39        if self.callback:40            self.callback(self.snapshot, self.state, timestep)41    def restore(self, prev_snapshot):42        self._restore(self.snapshot, prev_snapshot)43    def take_snapshot(self):44        return copy(self.snapshot)45def is_on_gpu(view: ContextView):46    return view.device_type() == DeviceType.GPU47def take_snapshot(wrapped_context):48    #49    context = wrapped_context.context.context  # extra indirection for OpenMM50    context_view = wrapped_context.view51    #52    positions = asarray(dlext.positions(context_view))53    forces = asarray(dlext.forces(context_view))54    ids = asarray(dlext.atom_ids(context_view))55    #56    velocities = asarray(dlext.velocities(context_view))57    if is_on_gpu(context_view):58        vel_mass = velocities59    else:60        inverse_masses = asarray(dlext.inverse_masses(context_view))61        vel_mass = (velocities, inverse_masses.reshape((-1, 1)))62    #63    box_vectors = context.getSystem().getDefaultPeriodicBoxVectors()64    a = box_vectors[0].value_in_unit(unit.nanometer)65    b = box_vectors[1].value_in_unit(unit.nanometer)66    c = box_vectors[2].value_in_unit(unit.nanometer)67    H = ((a[0], b[0], c[0]), (a[1], b[1], c[1]), (a[2], b[2], c[2]))68    origin = (0.0, 0.0, 0.0)69    dt = context.getIntegrator().getStepSize() / unit.picosecond70    # OpenMM doesn't have images71    return Snapshot(positions, vel_mass, forces, ids, None, Box(H, origin), dt)72def identity(x):73    return x74def safe_divide(v, invm):75    return cond(invm[0] == 0, lambda x: v, lambda x: np.divide(v, x), invm)76def build_snapshot_methods(context, sampling_method):77    if is_on_gpu(context):78        def unpack(vel_mass):79            return vel_mass[:, :3], vel_mass[:, 3:]80        def indices(snapshot):81            return snapshot.ids.argsort()82        def masses(snapshot):83            return snapshot.vel_mass[:, 3:]84    else:85        unpack = identity86        def indices(snapshot):87            return snapshot.ids88        def masses(snapshot):89            _, invM = snapshot.vel_mass90            return jax.vmap(safe_divide)(1, invM)91    def positions(snapshot):92        return snapshot.positions93    def momenta(snapshot):94        V, invM = unpack(snapshot.vel_mass)95        return jax.vmap(safe_divide)(V, invM).flatten()96    return SnapshotMethods(jit(positions), jit(indices), jit(momenta), jit(masses))97def build_helpers(context, sampling_method):98    # Depending on the device being used we need to use either cupy or numpy99    # (or numba) to generate a view of jax's DeviceArrays100    if is_on_gpu(context):101        cupy = importlib.import_module("cupy")102        view = cupy.asarray103        restore_vm = _restore_vm104        def sync_forces():105            cupy.cuda.get_current_stream().synchronize()106        @jit107        def adapt(biases):108            return np.int64(2**32 * biases.T)109    else:110        utils = importlib.import_module(".utils", package="pysages.backends")111        view = utils.view112        adapt = identity113        def sync_forces():114            pass115        def restore_vm(view, snapshot, prev_snapshot):116            # TODO: Check if we can omit modifying the masses117            # (in general the masses are unlikely to change)118            velocities = view(snapshot.vel_mass[0])119            masses = view(snapshot.vel_mass[1])120            velocities[:] = view(prev_snapshot.vel_mass[0])121            masses[:] = view(prev_snapshot.vel_mass[1])122    def bias(snapshot, state, sync_backend):123        """Adds the computed bias to the forces."""124        if state.bias is None:125            return126        biases = adapt(state.bias)127        # Forces may be computed asynchronously on the GPU, so we need to128        # synchronize them before applying the bias.129        sync_backend()130        forces = view(snapshot.forces)131        biases = view(biases.block_until_ready())132        forces += biases133        sync_forces()134    snapshot_methods = build_snapshot_methods(context, sampling_method)135    flags = sampling_method.snapshot_flags136    restore = partial(_restore, view, restore_vm=restore_vm)137    helpers = HelperMethods(build_data_querier(snapshot_methods, flags))138    return helpers, restore, bias139def check_integrator(context):140    integrator = context.getIntegrator()141    if isinstance(integrator, openmm.VariableLangevinIntegrator) or isinstance(142        integrator, openmm.VariableVerletIntegrator143    ):144        raise ValueError("Variable step size integrators are not supported")145def bind(146    wrapped_context: ContextWrapper, sampling_method: SamplingMethod, callback: Callable, **kwargs147):148    # For OpenMM we need to store a Simulation object as the context,149    simulation = wrapped_context.context150    context = simulation.context151    check_integrator(context)152    force = Force()153    force.add_to(context)  # OpenMM will handle the lifetime of the force154    wrapped_context.view = force.view(context)155    wrapped_context.run = simulation.step156    helpers, restore, bias = build_helpers(wrapped_context.view, sampling_method)157    snapshot = take_snapshot(wrapped_context)158    method_bundle = sampling_method.build(snapshot, helpers)159    sync_and_bias = partial(bias, sync_backend=wrapped_context.view.synchronize)160    sampler = Sampler(method_bundle, sync_and_bias, callback, restore)161    force.set_callback_in(context, sampler.update)...

Automation Testing Tutorials

Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.