examples/python/comm-testany/comm-testany.py

   1 # Copyright (c) 2010-2023. The SimGrid Team. All rights reserved.
   2 #
   3 # This program is free software; you can redistribute it and/or modify it
   4 # under the terms of the license (GNU LGPL) which comes with this package.
   5
   6 from argparse import ArgumentParser
   7 from typing import List
   8 import sys
   9
  10 from simgrid import Engine, Actor, Comm, Mailbox, this_actor
  11
  12
  13 def create_parser() -> ArgumentParser:
  14     parser = ArgumentParser()
  15     parser.add_argument(
  16         '--platform',
  17         type=str,
  18         required=True,
  19         help='path to the platform description'
  20     )
  21     return parser
  22
  23
  24 def rank0():
  25     rank0_mailbox: Mailbox = Mailbox.by_name("rank0")
  26     this_actor.info("Post my asynchronous receives")
  27     comm1 = rank0_mailbox.get_async()
  28     comm2 = rank0_mailbox.get_async()
  29     comm3 = rank0_mailbox.get_async()
  30     pending_comms: List[Comm] = [comm1, comm2, comm3]
  31
  32     this_actor.info("Send some data to rank-1")
  33     rank1_mailbox: Mailbox = Mailbox.by_name("rank1")
  34     for i in range(3):
  35         rank1_mailbox.put(i, 1)
  36
  37     this_actor.info("Test for completed comms")
  38     while pending_comms:
  39         flag = Comm.test_any(pending_comms)
  40         if flag != -1:
  41             pending_comms.pop(flag)
  42             this_actor.info("Remove a pending comm.")
  43         else:
  44             # Nothing matches, wait for a little bit
  45             this_actor.sleep_for(0.1)
  46     this_actor.info("Last comm is complete")
  47
  48
  49 def rank1():
  50     rank0_mailbox: Mailbox = Mailbox.by_name("rank0")
  51     rank1_mailbox: Mailbox = Mailbox.by_name("rank1")
  52     for i in range(3):
  53         data: int = rank1_mailbox.get()
  54         this_actor.info(f"Received {data}")
  55         msg_content = f"Message {i}"
  56         this_actor.info(f"Send '{msg_content}'")
  57         rank0_mailbox.put(msg_content, int(1e6))
  58
  59
  60 def main():
  61     settings = create_parser().parse_known_args()[0]
  62     e = Engine(sys.argv)
  63     e.load_platform(settings.platform)
  64
  65     Actor.create("rank0", e.host_by_name("Tremblay"), rank0)
  66     Actor.create("rank1", e.host_by_name("Fafard"), rank1)
  67
  68     e.run()
  69
  70
  71 if __name__ == "__main__":
  72     main()