Added support for accessing timesync channels from models imported as functions

Added support for returning all variables to models via a timesync conneciton if True passed in the timesync parameter 'additional_variables' for a model
Fix bugs in model queue handling which dropped some buffered messages from the model stdout
Fix bug in use of server name in bytes vs str

Author: langmm

HISTORY.rst

@@ -13,6 +13,9 @@ History
 * Migrate from setup.py + setup.cfg to pyproject.toml
 * Migrate build system to scikit-build-core
 * Introduce disassembler compilation tool class for dissecting compiled binaries
+* Added support for accessing timesync channels from models imported as functions
+* Added support for returning all variables to models via a timesync conneciton if True passed in the timesync parameter 'additional_variables' for a model
+* Fix bugs in model queue handling which dropped some buffered messages from the model stdout
 
 1.10.2 (2023-10-12) Minor bug fixes and dependency updates
 -------------------

tests/test_runner.py

@@ -91,23 +91,26 @@ def test_import_as_function():
     r"""Test import_as_function."""
     yamlfile = ex_yamls['fakeplant']['python']
     fmodel = import_as_function(yamlfile, remove_products=True)
-    input_args = {}
-    for x in fmodel.arguments:
-        input_args[x] = 1.0
-    fmodel.model_info()
-    result = fmodel(**input_args)
-    for x in fmodel.returns:
-        assert x in result
-    result = fmodel(*list(input_args.values()))
-    for x in fmodel.returns:
-        assert x in result
-    fmodel.reload()
-    fmodel.run()
-    result = fmodel(**input_args)
-    for x in fmodel.returns:
-        assert x in result
-    fmodel.stop()
-    fmodel.stop()
+    try:
+        input_args = {}
+        for x in fmodel.arguments:
+            input_args[x] = 1.0
+        fmodel.function_info
+        assert len(fmodel.returns) == 1
+        result = fmodel(**input_args)
+        for x in fmodel.returns:
+            assert x in result
+        result = fmodel(*list(input_args.values()))
+        for x in fmodel.returns:
+            assert x in result
+        fmodel.reload()
+        fmodel.run()
+        result = fmodel(**input_args)
+        for x in fmodel.returns:
+            assert x in result
+        fmodel.stop()
+    finally:
+        fmodel.stop()
 
 
 def test_import_as_function_server():
@@ -119,18 +122,67 @@ def test_import_as_function_server():
             break
     assert yamlfile
     fmodel = import_as_function(yamlfile, remove_products=True)
-    input_args = {}
-    for x in fmodel.arguments:
-        input_args[x] = 'hello'
-    fmodel.model_info()
-    result = fmodel(**input_args)
-    for x in fmodel.returns:
-        assert x in result
-    result = fmodel(*list(input_args.values()))
-    for x in fmodel.returns:
-        assert x in result
-    fmodel.stop()
-    fmodel.stop()
+    try:
+        input_args = {}
+        for x in fmodel.arguments:
+            input_args[x] = 'hello'
+        fmodel.function_info
+        assert len(fmodel.returns) == 1
+        result = fmodel(**input_args)
+        for x in fmodel.returns:
+            assert x in result
+        result = fmodel(*list(input_args.values()))
+        for x in fmodel.returns:
+            assert x in result
+        fmodel.stop()
+    finally:
+        fmodel.stop()
+
+
+def test_import_as_function_timesync():
+    r"""Test import_as_function with timesync."""
+    from yggdrasil import units
+    contents = r"""models:
+      - name: modelA
+        language: python
+        args:
+          - ./src/timesync.py
+          - 2
+          - day
+        timesync: True
+        outputs:
+          name: output
+          default_file:
+            name: modelA_output.txt
+            in_temp: True
+            filetype: table"""
+    yamlfile = os.path.join(os.path.dirname(ex_yamls['timesync1']['python']),
+                            'test_import_timesync.yml')
+    assert not os.path.isfile(yamlfile)
+    with open(yamlfile, 'w') as fd:
+        fd.write(contents)
+    fmodel = None
+    try:
+        # TODO: Test where x/y modified in call
+        steps = [
+            ({'timesync': (units.add_units(0.0, 'day'), {})},
+             {'x': 0.0, 'y': 1.0}),
+            ({'timesync': (units.add_units(1.0, 'day'), {})},
+             {'x': 0.47552825814757677, 'y': 0.09549150281252626}),
+        ]
+        fmodel = import_as_function(yamlfile, remove_products=True,
+                                    partial_timesync=True,
+                                    partial_comms=['timesync'])
+        fmodel.function_info
+        input_args = {}
+        for input_args, output_args in steps:
+            result = fmodel(**input_args)
+            assert result == output_args
+        fmodel.stop()
+    finally:
+        if fmodel is not None:
+            fmodel.stop()
+        os.remove(yamlfile)
 
 
 @pytest.mark.language('c')
@@ -148,9 +200,10 @@ def test_import_as_function_C():
     assert not os.path.isfile(yamlfile)
     with open(yamlfile, 'w') as fd:
         fd.write(contents)
+    fmodel = None
     try:
         fmodel = import_as_function(yamlfile, remove_products=True)
-        fmodel.model_info()
+        fmodel.function_info
         input_args = {}
         for x in fmodel.arguments:
             input_args[x] = b'hello'
@@ -161,6 +214,7 @@ def test_import_as_function_C():
         for x in fmodel.returns:
             assert x in result
         fmodel.stop()
-        fmodel.stop()
     finally:
+        if fmodel is not None:
+            fmodel.stop()
         os.remove(yamlfile)

yggdrasil/.ygg_schema.yml

@@ -3474,7 +3474,9 @@ definitions:
         additionalProperties:
           items:
             type: string
-          type: array
+          type:
+          - boolean
+          - array
         default: {}
         type: object
       aggregation:
@@ -5453,7 +5455,9 @@ definitions:
         additionalProperties:
           items:
             type: string
-          type: array
+          type:
+          - boolean
+          - array
         default: {}
         type: object
       aggregation:

yggdrasil/broker.py

@@ -89,16 +89,27 @@ def drivers(self):
         r"""iterator: All drivers."""
         return itertools.chain(self.models, self.connections)
 
+    @classmethod
+    def _get_commlist(cls, comm):
+        if isinstance(comm['commtype'], list):
+            return comm['commtype']
+        elif (comm['commtype'] in ['server', 'client', 'fork']
+              and isinstance(comm.get('comm_list', None), list)):
+            return comm['comm_list']
+        elif (comm['commtype'] in ['server', 'client']
+              and isinstance(comm.get('request_commtype', None), list)):
+            return comm['request_commtype']
+        return None
+
     @classmethod
     def _param_model_comm(cls, comm):
         commlist = None
         if isinstance(comm, dict):
             name = comm['name']
             direction = comm['direction']
-            model = comm['model']
+            model = comm.get('model', None)
             name = strip_model_prefix(name, model)
-            if isinstance(comm['commtype'], list):
-                commlist = comm['commtype']
+            commlist = cls._get_commlist(comm)
         else:
             name = comm.opp_name
             direction = comm.opp_direction
@@ -107,9 +118,15 @@ def _param_model_comm(cls, comm):
             if comm._commtype == 'fork':
                 commlist = comm.comm_list
             elif comm._commtype == 'server' and direction == 'send':
-                commlist = [comm.icomm]
+                if comm.icomm._commtype == 'fork':
+                    commlist = comm.icomm.comm_list
+                else:
+                    commlist = [comm.icomm]
             elif comm._commtype == 'client' and direction == 'recv':
-                commlist = [comm.ocomm]
+                if comm.ocomm._commtype == 'fork':
+                    commlist = comm.ocomm.comm_list
+                else:
+                    commlist = [comm.ocomm]
         return (model, direction, name, commlist)
 
     @classmethod
@@ -126,7 +143,8 @@ def _check_model_comm(cls, comm, model, direction, name):
                     pass
         raise CommBase.CommError("no match")
 
-    def find_model_comm(self, model, direction, name):
+    def find_model_comm(self, model, direction, name,
+                        is_split_server=False):
         r"""Locate a matching communicator from the registered
         connections.
 
@@ -135,6 +153,8 @@ def find_model_comm(self, model, direction, name):
             direction (str): Direction that the model communicator
                 operates in.
             name (str): Channel name for the communicator.
+            is_split_server (bool, optional): If True, the comm is part
+                of a split server comm.
 
         Returns:
             dict, CommBase: Partner communicator or parameters for the
@@ -155,8 +175,13 @@ def find_model_comm(self, model, direction, name):
                         comm, model, direction, name)
                 except CommBase.CommError:
                     pass
+        if is_split_server:
+            try:
+                return self.find_model_comm(model, 'recv', model)
+            except BrokerError:
+                pass
         raise BrokerError(f"Could not locate a {direction} communicator "
-                          f"with name \"{name}\" for model \"{model}\": "
+                          f"with name \"{name}\" for model \"{model}\":\n"
                           f"{pprint.pformat(self.comms)}")
 
     @classmethod
@@ -512,6 +537,24 @@ def _send_request(cls, action, *args, **kwargs):
         # client.info(f"RESPONSE [{action}]: {response['return']}")
         return response['return']
 
+    @classmethod
+    def is_address_set(cls, comm):
+        r"""Check if an address is defined.
+
+        Args:
+            comm (dict): Communication definition.
+
+        Returns:
+            bool: True if the address is defined, False otherwise.
+
+        """
+        if 'address' in comm:
+            return True
+        commlist = cls._get_commlist(comm)
+        if commlist is None:
+            return False
+        return all('address' in x for x in commlist)
+
     @classmethod
     def update_model_comm_kwargs(cls, name, kwargs, self=None):
         r"""Update the communicator parameters for the partner model
@@ -525,15 +568,31 @@ def update_model_comm_kwargs(cls, name, kwargs, self=None):
         if self is None:
             return cls._send_request('update_model_comm_kwargs',
                                      name, kwargs)
+        partner_model = kwargs.pop('partner_model', None)
         model = kwargs.pop('model')
         direction = kwargs.pop('direction')
         comm = self.find_model_comm(model, direction, name)
         assert isinstance(comm, dict)
+        comm0 = comm
+        commlist = self._get_commlist(comm)
+        if commlist is not None:
+            assert partner_model is not None
+            partners = [x['partner_model'] for x in commlist]
+            assert sum(x == partner_model for x in partners) == 1
+            idx = partners.index(partner_model)
+            comm = commlist[idx]
+            if kwargs['commtype'] in ['server', 'client']:
+                assert comm0['commtype'] == kwargs['commtype']
+                kwargs.pop('commtype')
+                if 'request_commtype' in kwargs:
+                    kwargs['commtype'] = kwargs.pop('request_commtype')
         comm.update(kwargs)
         assert comm['direction'] == direction
-        if name in self._awaiting_comm:
+        if commlist is None:
             assert 'address' in comm
-            # self.server_comm.info(f"UNLOCKING: {name}")
+        if name in self._awaiting_comm and self.is_address_set(comm0):
+            # self.server_comm.info(f"UNLOCKING: {name} "
+            #                       f"{self._awaiting_comm}")
             self._awaiting_comm.remove(name)
 
     @classmethod
@@ -569,9 +628,9 @@ def model_comm_kwargs(cls, name, direction, self=None):
             if isinstance(model_driver.is_server, dict):
                 if ((name == model_driver.is_server['input']
                      or name == model_driver.is_server['output'])):
-                    # TODO: Verify that this works
                     is_split_server = True
-        comm = self.find_model_comm(model, direction, name)
+        comm = self.find_model_comm(model, direction, name,
+                                    is_split_server)
         if isinstance(comm, dict):
             out = copy.deepcopy(comm)
             if (('address' not in comm
@@ -591,8 +650,11 @@ def model_comm_kwargs(cls, name, direction, self=None):
         )
         if is_split_server or out['commtype'] == 'model_function':
             out['global_scope'] = model
-        if 'address' not in out and 'partner_name' in out:
-            # self.server_comm.info(f"LOCKING {out['partner_name']}")
+        if (('partner_name' in out
+             and out['partner_name'] not in self._awaiting_comm
+             and (not self.is_address_set(out)))):
+            # self.server_comm.info(f"LOCKING: {out['partner_name']} "
+            #                       f"{self._awaiting_comm}")
             self._awaiting_comm.append(out['partner_name'])
         return out
 
@@ -611,7 +673,8 @@ def model_comm(cls, name, direction, **kwargs):
 
         """
         kwargs = dict(cls.model_comm_kwargs(name, direction), **kwargs)
-        assert direction == kwargs['direction']
+        if kwargs.get('commtype', None) not in ['client', 'server']:
+            assert direction == kwargs['direction']
         name = kwargs['name']
         global_scope = kwargs.pop('global_scope', False)
         global_name = name
@@ -622,13 +685,16 @@ def model_comm(cls, name, direction, **kwargs):
             # for case where server or function comm is split between
             # two aliases
             return cls._global_scope_comms[global_name]
-        if 'address' in kwargs:
+        if cls.is_address_set(kwargs):
             out = get_comm(**kwargs)
         else:
             partner_name = kwargs['partner_name']
             out = new_comm(**kwargs)
             cls.update_model_comm_kwargs(
-                partner_name, out.model_comm_kwargs)
+                partner_name,
+                dict(out.model_comm_kwargs,
+                     partner_model=kwargs['model']),
+            )
         if global_scope:
             cls._global_scope_comms[global_name] = out
         return out