Class level variables in host level functions with async RPCs

ejp9

Hello,

I am trying to keep track of several numpy arrays during my experiment. I require the ability to be able to store and modify them throughout the exp, but timing and speed is not important. To do so I have been calling functions on the host to modify and update the arrays which are stored as class level variables. They do not seem to be updating as I expect them to on the top level. I've recreated a bare bones example of my issue here. Is there something I'm misunderstanding about how the timing is handled that would cause this?

The output of the following code are three print statements: the original assignment, a print statement with the measured values from the experiment, then again the same list as the second (not sorted).

If I print the output inside the host function, it looks correct there though.

`from artiq.experiment import *
import numpy as np

class KernelTimingTest(EnvExperiment):

def build(self, **kwargs):
    
    self.setattr_device("core")
    self.setattr_device("scheduler")

    self.test_vals = np.zeros(10, dtype=np.int32)
 
    
def prepare(self):
    pass
         
@kernel
def run(self):
    self.core.reset()         
    print(self.test_vals)

    for i in range(len(test_vals)):
        self.test_vals[i] = self.update_function()
        delay(100*ms)
    print(self.test_vals)
 
    self.sort_list(self.test_vals) # sorts by value at each point       
    self.core.break_realtime()
    
    print(self.test_vals)


    
@rpc(flags={"async"})     
def sort_list(self, vals):
    indices = np.argsort(vals)
    self.test_vals = vals[indices]

lriesebos

ah yeah, this is a nice one! it really goes into the details of how RPC and class attribute synchronization works. our code has several pieces of documentation around this exact behavior. for example https://gitlab.com/duke-artiq/dax/-/blob/master/dax/modules/safety_context.py?ref_type=heads#L97

let me try to summarize it by going through the sequence of your code:

build(): self.test_vals is initialized as all zeros
run(): we enter the kernel and now two copies of self.test_vals exist. the in-kernel one (on the core device) and the host-side one. both are equal at the start of the kernel.
- you print the array at the start. both copies are equal, so you see the expected value
- you start modifying the array, these changes are only reflected in the in-kernel copy of the array
- you print the in-kernel copy of the list, seeing the expected values
- you call self.sort_list() and you provide the in-kernel list to the host. at this moment, the host still has an all zeroes array for test_vals. you sort the incoming (i.e. copy of the in-kernel array received from the RPC argument) array and assign that sorted array to test_vals. at this moment, your in-kernel array is not sorted. your host-side array is now the sorted array.
- you print the in-kernel array, so you will see the unsorted array here.
now, when you exit the kernel, the class attributes that were modified in-kernel are synchronized with your host-side class attributes. this means that the host-side test_vals will be overwritten with the unsorted data. so if you print the value in the analysis() phase of the experiment, you should see that behavior.

so based on my knowledge, working with class attributes inside and outside kernels has this sub-optimal synchronization behavior. it is for good reasons though, but the programmer needs to keep this in mind. I think that in NAC3, this will be eliminated by differentiating in-kernel and host-side attributes explicitly. until that time, this feature can be used, but you need detailed knowledge of how it works to make sure you don't run into issues like these 🙂

ejp9

Thank you for your reply. This is how I have updated the code to reflect the needs for my experiment, which are, run an experiment, update the experiment parameters based on the results, and then iterate on this process.

class KernelTimingTest(EnvExperiment):

def build(self, **kwargs):
    
    self.setattr_device("core")
    self.setattr_device("scheduler")


    self.test_vals = np.array([1,10,5,7,6,2,8], dtype=np.int32)
 
    
def prepare(self):
    pass
         

def run(self):
    self.start_exp()  
    print(self.test_vals)
    
    self.run_exp()
    print(self.test_vals)
    
    self.sort_list()   
    print(self.test_vals)

   
def sort_list(self):
    self.test_vals = self.test_vals[np.argsort(self.test_vals)]

@kernel 
def start_exp(self):
    self.core.reset()
    
@kernel
def run_exp(self):
    self.test_vals[0] = 100

I think I understand what is happening now, and this code is working, but I feel like the self.test_vals[0] = 100 executed on the kernel shouldn't be reflected on the host based on your explanation, but now it seems to be. The output of the code is

print:[ 1 10 5 7 6 2 8]
print:[100 10 5 7 6 2 8]
print:[ 2 5 6 7 8 10 100]

Which is the correct functionality I want, but feel like it is not consistent with how I understand the values are being stored.

lriesebos

ejp9 hm, to me that looks fully consistent. so maybe I didn't explain it well.

so in your run() function

start_exp() will do nothing. it will spawn a kernel and reset the device. when that kernel ends, an in-kernel to host attribute sync will happen, but test_vals is not changed in-kernel, so also on the host it will be the same.
run_exp() will spawn a kernel, change test_vals[0] in-kernel. the kernel ends, and the in-kernel test_vals overwrites the host test_vals (i.e. the in-kernel to host attribute sync). the print happens on the host (i.e. after the sync) and you see the mutated value of test_vals.
sort_list() is a host function and it sorts the host test_vals correctly and eventually prints it. no kernels are spawned, and no synchronizations are done anymore.

so the difference with your last example is that two of these attribute synchronizations happen in this code. let me know if that improves your understanding.

ejp9

So when exactly do host-kernel attribute syncs occur?

In this example, we are running on host and call kernel method, and you are saying after this kernel method is over, the kernel will update the attributes to the host. But in the previous code, we were running a kernel method and calling host methods within that kernel, and no such syncs were made. Is this because we never actually exit the kernel method in the first case until the whole experiment is over? Because the host can't sync attributes in the kernel direction? Both/neither?? 🙂

lriesebos

so the first example was a run() function as one big kernel. the calls to the host functions are called RPC functions. no attribute syncing is done during RPC functions and the in-kernel and host attributes effectively go out of sync. only when the whole run() function finished (in the last example), the single kernel in that program finishes and attribute sync happens.

in the last example consisted of multiple kernels and had no RPC calls. so when each kernel ends, attribute sync happens.

most of the time we want to put as much as possible in a single kernel because starting a new kernel takes a lot of time (order of seconds).

ejp9

Thanks so much, I think it's clear now 🙂

ejp9

I've been thinking more about how you mentioned that is is very expensive time wise to recreate kernel instances. In my experiment, I require a run of the experiment, then a bunch of host side computations, and then iterating this with updated values based on the host side computations. I don't see how to do this without consistently calling new kernel instances at every iteration (perhaps several hundred runs of the small experiment). Is there a way to get around this while maintaining the flow of information into the kernel? Synchronized timing between the different experiment calls is not important, but adding in a few seconds every kernel instance could prove to be an issue.

Best

lriesebos

Not exactly sure what the context is you are working in, but that sounds like a "partial compilation" situation where you run the same kernel multiple times, but with different parameters. Like VQE. for that, we have leveraged synchronous RPC in the past, but I don't know how well that works when returning lists form an RPC function. Roughly something like:

def run(self):
    self.core.reset()
    data = [0] * 10
    while self.decide_if_we_will_continue_with(data):
        self.experiment_iteration_with_results_stored_in(data)
        data = self.host_side_computation(data)
    self.core.wait_until_mu(now_mu())

@rpc  # this means it is a synchronous RPC function, so a value is returned
def host_side_computation(self, input_data: TList(TInt32)) -> TList(TInt32):
    output_data = your_computations(input_data)
    return output_data

So as mentioned, I don't know if RPC works like that with lists. We do know this works for tuples as they have fixed size.

Timing information for RPC (wihtout any computation time) can be found here: https://gitlab.com/duke-artiq/dax/-/jobs/6559344260#L495 . So sync RPC takes about 2ms while spawning a kernel takes 165ms. Note that these are lower bounds and both the RPC function and the spawned kernel for these benchmarks are empty. In practice, our kernel code has enough complexity such that spawning a kernel takes in the order of seconds.

I think this is the paper that talks about this too: https://arxiv.org/abs/2308.10787