Async RPC support #540

mwojcik · 2024-09-11T19:24:12+08:00

mwojcik commented

2024-09-11 19:24:12 +08:00

Starting out with support for asynchronous RPCs. Work in progress as I am open to feedback from more experienced contributors.

RPC codegen path seems straightforward enough, but I did wonder how to pass the flags to rpc_codegen_callback. So I am proposing to extend FunSignature to include options, which for now only include Async - they could also include FastMath (like in old kernel) and SubkernelDestination(u8) (will for subkernels). That will mean also that codegen callbacks don't have to change their signature to include flags.

Additionally, decorators with arguments have been missing (compiler would say the called function does not exist) - took me a while to find them again - see nac3artiq/src/lib.rs - the body.retain function. The ExprKind::Name that the function was looking for is nested inside an ExprKind::Call. I am not sure if this is the correct course of action.

Additionally, I am not sure where to find __artiq_rpc_flags__ defined in the Python part that could be safely passed to the FunSignature. Otherwise I could still find the flag within the Stmt::FunctionDef.

Starting out with support for asynchronous RPCs. Work in progress as I am open to feedback from more experienced contributors. RPC codegen path seems straightforward enough, but I did wonder how to pass the flags to ``rpc_codegen_callback``. So I am proposing to extend ``FunSignature`` to include options, which for now only include ``Async`` - they could also include ``FastMath`` (like in old kernel) and ``SubkernelDestination(u8)`` (will for subkernels). That will mean also that codegen callbacks don't have to change their signature to include flags. Additionally, decorators with arguments have been missing (compiler would say the called function does not exist) - took me a while to find them again - see nac3artiq/src/lib.rs - the body.retain function. The ``ExprKind::Name`` that the function was looking for is nested inside an ``ExprKind::Call``. I am not sure if this is the correct course of action. Additionally, I am not sure where to find ``__artiq_rpc_flags__`` defined in the Python part that could be safely passed to the ``FunSignature``. Otherwise I could still find the flag within the ``Stmt::FunctionDef``.

sb10q reviewed 2024-09-12 10:45:32 +08:00

nac3artiq/src/lib.rs Outdated

						
				@ -200,1 +200,4 @@

				                                        || id.to_string() == "rpc"

				                                } else if let ExprKind::Call { func, .. } = &decorator.node {

				                                    // decorators with flags (e.g. rpc async) have Call for the node;

				                                    // this is to remove the middle part

sb10q commented

2024-09-12 10:41:41 +08:00

Why remove it? Then you don't know what it is afterwards.

And you could potentially have Call for any decorator. I suggest writing a helper function like:

fn decorator_id(node: ast::Node) -> Option<StrRef> {
  if let ExprKind::Name { id, .. } = node {
     Some(id)
   else if let ExprKind::Call { func, .. } = node {
      if let ExprKind::Name { id, .. } = func.node {
         Some(id)
      } else {
          None
      }
   } else {
     None
  }
}

and using it consistently to process decorators consistently, including for @nac3 above. @nac3(), @kernel() should potentially also be legal.

Why remove it? Then you don't know what it is afterwards. And you could potentially have Call for any decorator. I suggest writing a helper function like: ```rust fn decorator_id(node: ast::Node) -> Option<StrRef> { if let ExprKind::Name { id, .. } = node { Some(id) else if let ExprKind::Call { func, .. } = node { if let ExprKind::Name { id, .. } = func.node { Some(id) } else { None } } else { None } } ``` and using it consistently to process decorators consistently, including for ``@nac3`` above. ``@nac3()``, ``@kernel()`` should potentially also be legal.

sb10q reviewed 2024-09-12 10:47:13 +08:00

nac3artiq/demo/min_artiq.py Outdated

						
				@ -119,0 +118,4 @@

				    @wraps(function)

				    def wrapped(function)

				        return function

				    wrapped.__artiq_rpc_flags__ = flags

sb10q commented

2024-09-12 10:47:13 +08:00

Actually, processing __artiq_rpc_flags__ doesn't work easily with NAC3, contrary to what I said earlier in the chat, sorry for that mistake. Probably better to process the AST after all.

Actually, processing ``__artiq_rpc_flags__`` doesn't work easily with NAC3, contrary to what I said earlier in the chat, sorry for that mistake. Probably better to process the AST after all.

👍 1

mwojcik marked this conversation as resolved

sb10q reviewed 2024-09-12 10:47:58 +08:00

nac3core/src/toplevel/builtins.rs Outdated

						
				@ -77,6 +77,7 @@ pub fn get_exn_constructor(

				        args: exn_cons_args,

				        ret: exn_type,

				        vars: VarMap::default(),

				        opts: vec![],

sb10q commented

2024-09-12 10:47:58 +08:00

Can this be handled (1) entirely in nac3artiq and (2) without passing opts all over the place?

Can this be handled (1) entirely in nac3artiq and (2) without passing ``opts`` all over the place?

mwojcik marked this conversation as resolved

mwojcik force-pushed async_rpc from b79e3312ef to 94f0c467fa

2024-09-12 14:40:53 +08:00

Compare

mwojcik force-pushed async_rpc from 733dc218d0 to 50e9ccb7c4

2024-09-12 16:33:51 +08:00

Compare

mwojcik force-pushed async_rpc from 50e9ccb7c4 to 201a89b6d3

2024-09-12 16:48:24 +08:00

Compare

mwojcik commented

2024-09-12 16:50:21 +08:00

No need for FuncOptions, I can just pass is_async (extracted from the decorator arguments) to the rpc gen function.

Similarly, other ARTIQ-only options (like for subkernels later) would be still handled within nac3artiq in similar fashion.

Seems alright now, just need to test with ARTIQ.

No need for ``FuncOptions``, I can just pass ``is_async`` (extracted from the decorator arguments) to the rpc gen function. Similarly, other ARTIQ-only options (like for subkernels later) would be still handled within ``nac3artiq`` in similar fashion. Seems alright now, just need to test with ARTIQ.

mwojcik commented

2024-09-12 17:29:28 +08:00

I did test it successfully with the kc705. No changes required in ARTIQ besides flake update (and removing the restriction with #nac3todo in few examples.

I believe there may be comments about code style though, but I'm taking off WIP now.

Testing experiment:

from artiq.language.core import *
from artiq.experiment import *
from artiq.coredevice.core import Core
from artiq.coredevice.ttl import TTLOut

@nac3
class AsyncTest(EnvExperiment):
    core: KernelInvariant[Core]
    led: KernelInvariant[TTLOut]

    @rpc
    def simple_rpc(self):
        print("totally normal RPC method")

    @rpc(flags={"async"})
    def simple_async_rpc(self):
        print("totally normal RPC method but async")
        while True:
            from time import sleep
            sleep(1)
            print("ZzzZzz")

    def build(self):
        self.setattr_device("core")
        self.setattr_device("led")
        
    @kernel
    def run(self):
        self.core.reset()
        self.simple_rpc()
        self.simple_async_rpc()
        self.core.delay(100.*ms)
        
        while True:
            self.led.pulse(100.*ms)
            self.core.delay(100.*ms)

I did test it successfully with the kc705. No changes required in ARTIQ besides flake update (and removing the restriction with ``#nac3todo`` in few examples. I believe there may be comments about code style though, but I'm taking off WIP now. Testing experiment: ```python from artiq.language.core import * from artiq.experiment import * from artiq.coredevice.core import Core from artiq.coredevice.ttl import TTLOut @nac3 class AsyncTest(EnvExperiment): core: KernelInvariant[Core] led: KernelInvariant[TTLOut] @rpc def simple_rpc(self): print("totally normal RPC method") @rpc(flags={"async"}) def simple_async_rpc(self): print("totally normal RPC method but async") while True: from time import sleep sleep(1) print("ZzzZzz") def build(self): self.setattr_device("core") self.setattr_device("led") @kernel def run(self): self.core.reset() self.simple_rpc() self.simple_async_rpc() self.core.delay(100.*ms) while True: self.led.pulse(100.*ms) self.core.delay(100.*ms) ```

mwojcik changed title from ~~[WIP] Async RPC support~~ to Async RPC support

2024-09-12 17:29:31 +08:00

mwojcik force-pushed async_rpc from 201a89b6d3 to 613159da39

2024-09-12 17:32:23 +08:00

Compare

mwojcik force-pushed async_rpc from 613159da39 to c94f407732

2024-09-12 17:34:25 +08:00

Compare

sb10q requested review from derppening 2024-09-13 11:24:56 +08:00

derppening approved these changes 2024-09-13 11:55:55 +08:00

derppening left a comment

Please run cargo fmt and cargo clippy before merging.

Please run `cargo fmt `and `cargo clippy` before merging.

nac3artiq/src/lib.rs Outdated

						
				@ -844,6 +863,38 @@ impl Nac3 {

				    }

				}

				fn decorator_id_string(decorator: &Located<ExprKind>) -> Option<String> {

derppening commented

2024-09-13 11:53:52 +08:00

Please add documentation for these two functions.

nac3core/src/toplevel/composer.rs Outdated

						
				@ -2060,0 +2063,4 @@

				                            if matches!(&func.node,

				                                ast::ExprKind::Name{ id, .. } if id == &"rpc".into())

				                            {

				                                instance_to_symbol.insert(String::new(), simple_name.to_string());

derppening commented

2024-09-13 11:55:06 +08:00

What exactly is this for?

mwojcik commented

2024-09-13 11:58:35 +08:00

Mimics the check above it - for decorators with arguments the Name is behind a Call as I mentioned before. Without it, it wouldn't be considered an external function and compiler would try to compile it together with the kernel.

👍 1