olas-operate-middleware 0.13.6__tar.gz → 0.13.7__tar.gz

{olas_operate_middleware-0.13.6 → olas_operate_middleware-0.13.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: olas-operate-middleware
-Version: 0.13.6
+Version: 0.13.7
 Summary: 
 License-File: LICENSE
 Author: David Vilela

{olas_operate_middleware-0.13.6 → olas_operate_middleware-0.13.7}/operate/services/funding_manager.py

@@ -60,6 +60,7 @@ from operate.ledger.profiles import (
 from operate.operate_types import Chain, ChainAmounts, LedgerType, OnChainState
 from operate.services.protocol import EthSafeTxBuilder, StakingManager, StakingState
 from operate.services.service import NON_EXISTENT_TOKEN, Service
+from operate.utils import concurrent_execute
 from operate.utils.gnosis import drain_eoa, get_asset_balance, get_owners
 from operate.utils.gnosis import transfer as transfer_from_safe
 from operate.utils.gnosis import transfer_erc20_from_safe
@@ -344,13 +345,27 @@ class FundingManager:
 
             # os.environ["CUSTOM_CHAIN_RPC"] = ledger_config.rpc  # TODO do we need this?
 
-            # Determine bonded native amount
-            service_registry_address = CHAIN_PROFILES[chain]["service_registry"]
+            # Fetch on-chain data
             service_registry = registry_contracts.service_registry.get_instance(
                 ledger_api=ledger_api,
-                contract_address=service_registry_address,
+                contract_address=CHAIN_PROFILES[chain]["service_registry"],
+            )
+            (
+                service_info,
+                operator_balance,
+                current_staking_program,
+            ) = concurrent_execute(
+                (service_registry.functions.getService(service_id).call, ()),
+                (
+                    service_registry.functions.getOperatorBalance(
+                        master_safe, service_id
+                    ).call,
+                    (),
+                ),
+                (staking_manager.get_current_staking_program, (service_id,)),
             )
-            service_info = service_registry.functions.getService(service_id).call()
+
+            # Determine bonded native amount
             security_deposit = service_info[0]
             service_state = service_info[6]
             agent_ids = service_info[7]
@@ -362,15 +377,9 @@ class FundingManager:
             ):
                 bonded_assets[ZERO_ADDRESS] += security_deposit
 
-            operator_balance = service_registry.functions.getOperatorBalance(
-                master_safe, service_id
-            ).call()
             bonded_assets[ZERO_ADDRESS] += operator_balance
 
             # Determine bonded token amount for staking programs
-            current_staking_program = staking_manager.get_current_staking_program(
-                service_id=service_id,
-            )
             target_staking_program = user_params.staking_program_id
             staking_contract = staking_manager.get_staking_contract(
                 staking_program_id=current_staking_program or target_staking_program,
@@ -381,9 +390,7 @@ class FundingManager:
 
             staking_manager = StakingManager(Chain(chain))
             staking_params = staking_manager.get_staking_params(
-                staking_contract=staking_manager.get_staking_contract(
-                    staking_program_id=user_params.staking_program_id,
-                ),
+                staking_contract=staking_contract,
             )
 
             service_registry_token_utility_address = staking_params[
@@ -396,24 +403,62 @@ class FundingManager:
                 )
             )
 
+            (
+                *agent_instances_and_bonds,
+                token_bond,
+                security_deposits,
+                staking_state,
+            ) = concurrent_execute(
+                *(
+                    [
+                        (
+                            service_registry.functions.getInstancesForAgentId(
+                                service_id, agent_id
+                            ).call,
+                            (),
+                        )
+                        for agent_id in agent_ids
+                    ]
+                    + [
+                        (
+                            service_registry_token_utility.functions.getAgentBond(
+                                service_id, agent_id
+                            ).call,
+                            (),
+                        )
+                        for agent_id in agent_ids
+                    ]
+                    + [
+                        (
+                            service_registry_token_utility.functions.getOperatorBalance(
+                                master_safe, service_id
+                            ).call,
+                            (),
+                        ),
+                        (
+                            service_registry_token_utility.functions.mapServiceIdTokenDeposit(
+                                service_id
+                            ).call,
+                            (),
+                        ),
+                        (
+                            staking_manager.staking_state,
+                            (service_id, staking_params["staking_contract"]),
+                        ),
+                    ]
+                ),
+            )
+
             agent_bonds = 0
-            for agent_id in agent_ids:
-                num_agent_instances = service_registry.functions.getInstancesForAgentId(
-                    service_id, agent_id
-                ).call()[0]
-                agent_bond = service_registry_token_utility.functions.getAgentBond(
-                    service_id, agent_id
-                ).call()
+            for agent_instances, agent_bond in zip(
+                agent_instances_and_bonds[: len(agent_ids)],
+                agent_instances_and_bonds[len(agent_ids) :],
+            ):
+                num_agent_instances = agent_instances[0]
                 agent_bonds += num_agent_instances * agent_bond
 
             if service_state == OnChainState.TERMINATED_BONDED:
                 num_agent_instances = service_info[5]
-                token_bond = (
-                    service_registry_token_utility.functions.getOperatorBalance(
-                        master_safe,
-                        service_id,
-                    ).call()
-                )
                 agent_bonds += num_agent_instances * token_bond
 
             security_deposit = 0
@@ -422,20 +467,11 @@ class FundingManager:
                 <= service_state
                 < OnChainState.TERMINATED_BONDED
             ):
-                security_deposit = (
-                    service_registry_token_utility.functions.mapServiceIdTokenDeposit(
-                        service_id
-                    ).call()[1]
-                )
+                security_deposit = security_deposits[1]
 
             bonded_assets[staking_params["staking_token"]] += agent_bonds
             bonded_assets[staking_params["staking_token"]] += security_deposit
 
-            staking_state = staking_manager.staking_state(
-                service_id=service_id,
-                staking_contract=staking_params["staking_contract"],
-            )
-
             if staking_state in (StakingState.STAKED, StakingState.EVICTED):
                 for token, amount in staking_params[
                     "additional_staking_tokens"
@@ -564,40 +600,64 @@ class FundingManager:
 
     def _get_master_safe_balances(self, thresholds: ChainAmounts) -> ChainAmounts:
         output = ChainAmounts()
+        batch_calls_args = {}
         for chain_str, addresses in thresholds.items():
             chain = Chain(chain_str)
             master_safe = self._resolve_master_safe(chain)
-            master_safe_dict = output.setdefault(chain_str, {}).setdefault(
-                master_safe, {}
-            )
+            output.setdefault(chain_str, {}).setdefault(master_safe, {})
             for _, assets in addresses.items():
                 for asset, _ in assets.items():
-                    master_safe_dict[asset] = get_asset_balance(
-                        ledger_api=get_default_ledger_api(chain),
-                        asset_address=asset,
-                        address=master_safe,
-                        raise_on_invalid_address=False,
+                    batch_calls_args[
+                        (
+                            get_default_ledger_api(chain),
+                            asset,
+                            master_safe,
+                            False,
+                        )
+                    ] = (
+                        chain_str,
+                        master_safe,
+                        asset,
                     )
 
+        batch_calls_results = concurrent_execute(
+            *[(get_asset_balance, args) for args in batch_calls_args.keys()]
+        )
+        for args, balance in zip(batch_calls_args.keys(), batch_calls_results):
+            chain_str, master_safe, asset = batch_calls_args[args]
+            output[chain_str][master_safe][asset] = balance
+
         return output
 
     def _get_master_eoa_balances(self, thresholds: ChainAmounts) -> ChainAmounts:
         output = ChainAmounts()
+        batch_calls_args = {}
         for chain_str, addresses in thresholds.items():
             chain = Chain(chain_str)
             master_eoa = self._resolve_master_eoa(chain)
-            master_eoa_dict = output.setdefault(chain_str, {}).setdefault(
-                master_eoa, {}
-            )
+            output.setdefault(chain_str, {}).setdefault(master_eoa, {})
             for _, assets in addresses.items():
                 for asset, _ in assets.items():
-                    master_eoa_dict[asset] = get_asset_balance(
-                        ledger_api=get_default_ledger_api(chain),
-                        asset_address=asset,
-                        address=master_eoa,
-                        raise_on_invalid_address=False,
+                    batch_calls_args[
+                        (
+                            get_default_ledger_api(chain),
+                            asset,
+                            master_eoa,
+                            False,
+                        )
+                    ] = (
+                        chain_str,
+                        master_eoa,
+                        asset,
                     )
 
+        batch_calls_results = concurrent_execute(
+            *[(get_asset_balance, args) for args in batch_calls_args.keys()]
+        )
+        for args, balance in zip(batch_calls_args.keys(), batch_calls_results):
+            chain_str, master_eoa, asset = batch_calls_args[args]
+            output[chain_str][master_eoa][asset] = balance
+
         return output
 
     def fund_master_eoa(self) -> None:
@@ -623,12 +683,30 @@ class FundingManager:
             }
         )
         master_eoa_balances = self._get_master_eoa_balances(master_eoa_topups)
+        master_safe_balance = self._get_master_safe_balances(master_eoa_topups)
         master_eoa_shortfalls = self._compute_shortfalls(
             balances=master_eoa_balances,
             thresholds=master_eoa_topups * DEFAULT_EOA_THRESHOLD,
             topups=master_eoa_topups,
         )
-        self.fund_chain_amounts(master_eoa_shortfalls)
+        possible_to_fund_shortfalls = ChainAmounts(
+            {
+                chain_str: {
+                    address: {
+                        asset: min(
+                            amount,
+                            master_safe_balance.get(chain_str, {})
+                            .get(self._resolve_master_safe(Chain(chain_str)), {})
+                            .get(asset, 0),
+                        )
+                        for asset, amount in assets.items()
+                    }
+                    for address, assets in addresses.items()
+                }
+                for chain_str, addresses in master_eoa_shortfalls.items()
+            }
+        )
+        self.fund_chain_amounts(possible_to_fund_shortfalls)
 
     def funding_requirements(self, service: Service) -> t.Dict:
         """Funding requirements"""
@@ -639,9 +717,13 @@ class FundingManager:
         total_requirements: ChainAmounts
         chains = [Chain(chain_str) for chain_str in service.chain_configs.keys()]
 
-        # Protocol shortfall
-        protocol_thresholds = self._compute_protocol_asset_requirements(service)
-        protocol_balances = self._compute_protocol_bonded_assets(service)
+        (
+            protocol_thresholds,
+            protocol_balances,
+        ) = concurrent_execute(
+            (self._compute_protocol_asset_requirements, (service,)),
+            (self._compute_protocol_bonded_assets, (service,)),
+        )
         protocol_topups = protocol_thresholds
         protocol_shortfalls = self._compute_shortfalls(
             balances=protocol_balances,

{olas_operate_middleware-0.13.6 → olas_operate_middleware-0.13.7}/operate/services/protocol.py

@@ -80,6 +80,7 @@ from operate.ledger.profiles import CONTRACTS, STAKING
 from operate.operate_types import Chain as OperateChain
 from operate.operate_types import ContractAddresses
 from operate.services.service import NON_EXISTENT_TOKEN
+from operate.utils import concurrent_execute
 from operate.utils.gnosis import (
     MultiSendOperation,
     SafeOperation,
@@ -226,37 +227,20 @@ class StakingManager:
     def _get_staking_params(chain: OperateChain, staking_contract: str) -> t.Dict:
         """Get staking params"""
         ledger_api = get_default_ledger_api(chain=chain)
-        instance = StakingManager.staking_ctr.get_instance(
-            ledger_api=ledger_api,
-            contract_address=staking_contract,
+
+        second_token_func = (
+            lambda: None  # pylint: disable=unnecessary-lambda-assignment  # noqa: E731
         )
-        agent_ids = instance.functions.getAgentIds().call()
-        service_registry = instance.functions.serviceRegistry().call()
-        staking_token = instance.functions.stakingToken().call()
-        service_registry_token_utility = (
-            instance.functions.serviceRegistryTokenUtility().call()
+        second_token_amount_func = (
+            lambda: None  # pylint: disable=unnecessary-lambda-assignment  # noqa: E731
         )
-        min_staking_deposit = instance.functions.minStakingDeposit().call()
-        activity_checker = instance.functions.activityChecker().call()
-
-        output = {
-            "staking_contract": staking_contract,
-            "agent_ids": agent_ids,
-            "service_registry": service_registry,
-            "staking_token": staking_token,
-            "service_registry_token_utility": service_registry_token_utility,
-            "min_staking_deposit": min_staking_deposit,
-            "activity_checker": activity_checker,
-            "additional_staking_tokens": {},
-        }
         try:
             instance = StakingManager.dual_staking_ctr.get_instance(
                 ledger_api=ledger_api,
                 contract_address=staking_contract,
             )
-            output["additional_staking_tokens"][
-                instance.functions.secondToken().call()
-            ] = instance.functions.secondTokenAmount().call()
+            second_token_func = instance.functions.secondToken().call
+            second_token_amount_func = instance.functions.secondTokenAmount().call
         except Exception:  # pylint: disable=broad-except # nosec
             # Contract is not a dual staking contract
 
@@ -269,6 +253,45 @@ class StakingManager:
             # avoid any issues we are simply catching all exceptions.
             pass
 
+        instance = StakingManager.staking_ctr.get_instance(
+            ledger_api=ledger_api,
+            contract_address=staking_contract,
+        )
+        (
+            agent_ids,
+            service_registry,
+            staking_token,
+            service_registry_token_utility,
+            min_staking_deposit,
+            activity_checker,
+            second_token,
+            second_token_amount,
+        ) = concurrent_execute(
+            (instance.functions.getAgentIds().call, ()),
+            (instance.functions.serviceRegistry().call, ()),
+            (instance.functions.stakingToken().call, ()),
+            (instance.functions.serviceRegistryTokenUtility().call, ()),
+            (instance.functions.minStakingDeposit().call, ()),
+            (instance.functions.activityChecker().call, ()),
+            (second_token_func, ()),
+            (second_token_amount_func, ()),
+            ignore_exceptions=True,
+        )
+
+        output = {
+            "staking_contract": staking_contract,
+            "agent_ids": agent_ids,
+            "service_registry": service_registry,
+            "staking_token": staking_token,
+            "service_registry_token_utility": service_registry_token_utility,
+            "min_staking_deposit": min_staking_deposit,
+            "activity_checker": activity_checker,
+            "additional_staking_tokens": (
+                {second_token: second_token_amount}
+                if second_token and second_token_amount
+                else {}
+            ),
+        }
         return output
 
     def get_staking_params(self, staking_contract: str) -> t.Dict:

{olas_operate_middleware-0.13.6 → olas_operate_middleware-0.13.7}/operate/utils/__init__.py

@@ -19,14 +19,25 @@
 
 """Helper utilities."""
 
+import asyncio
+import inspect
+import logging
 import os
 import platform
 import shutil
 import time
 import typing as t
+from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import TimeoutError as FuturesTimeoutError
+from contextlib import contextmanager
 from pathlib import Path
 from threading import Lock
 
+from operate.constants import DEFAULT_TIMEOUT
+
+
+logger = logging.getLogger(__name__)
+
 
 class SingletonMeta(type):
     """A metaclass for creating thread-safe singleton classes."""
@@ -153,3 +164,100 @@ def unrecoverable_delete(file_path: Path, passes: int = 3) -> None:
         print(f"Permission denied to securely delete file '{file_path}'.")
     except Exception as e:  # pylint: disable=broad-except
         print(f"Error during secure deletion of '{file_path}': {e}")
+
+
+@contextmanager
+def timing_context(label: str = "Block") -> t.Generator[None, None, None]:
+    """Context manager for timing a code block."""
+    start = time.perf_counter()
+    try:
+        yield
+    finally:
+        end = time.perf_counter()
+        logger.debug(f"[{label}] Elapsed time: {end - start:.4f} seconds")
+
+
+def concurrent_execute(
+    *func_calls: t.Tuple[t.Callable, t.Tuple],
+    ignore_exceptions: bool = False,
+) -> t.List[t.Any]:
+    """Execute callables concurrently.
+
+    This is a synchronous convenience wrapper around `parallel_execute_async`.
+    If called from within an active asyncio event loop, use
+    `await parallel_execute_async(...)` instead.
+    """
+
+    async def _runner() -> t.List[t.Any]:
+        return await concurrent_execute_async(
+            *func_calls,
+            ignore_exceptions=ignore_exceptions,
+        )
+
+    try:
+        asyncio.get_running_loop()
+    except RuntimeError:
+        # No running loop in this thread.
+        return asyncio.run(_runner())
+
+    # Running inside an event loop thread.
+    # We cannot call `asyncio.run` here, so offload to a background thread.
+    # NOTE: this blocks the current thread until completion.
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        future = executor.submit(asyncio.run, _runner())
+        return future.result()
+
+
+async def concurrent_execute_async(
+    *func_calls: t.Tuple[t.Callable, t.Tuple],
+    ignore_exceptions: bool = False,
+) -> t.List[t.Any]:
+    """Execute callables concurrently using asyncio.
+
+    - Async callables are awaited directly.
+    - Sync callables are executed via `asyncio.to_thread`.
+
+    Results are returned in the same order as `funcs`/`args_list`.
+    """
+
+    async def _invoke(func: t.Callable, args: t.Tuple) -> t.Any:
+        with timing_context(f"Executing {func.__name__}"):
+            if inspect.iscoroutinefunction(func):
+                return await t.cast(t.Awaitable[t.Any], func(*args))
+            return await asyncio.to_thread(func, *args)
+
+    results: t.List[t.Any] = [None] * len(func_calls)
+
+    async def _invoke_indexed(
+        idx: int, func: t.Callable, args: t.Tuple
+    ) -> t.Tuple[int, t.Any]:
+        try:
+            return idx, await _invoke(func, args)
+        except Exception as e:  # pylint: disable=broad-except
+            return idx, e
+
+    tasks: t.List[asyncio.Task] = [
+        asyncio.create_task(_invoke_indexed(idx, func, args))
+        for idx, (func, args) in enumerate(func_calls)
+    ]
+
+    try:
+        for task in asyncio.as_completed(tasks, timeout=DEFAULT_TIMEOUT):
+            idx, outcome = await task
+            if isinstance(outcome, BaseException):
+                if ignore_exceptions:
+                    results[idx] = None
+                else:
+                    raise outcome
+            else:
+                results[idx] = outcome
+    except asyncio.TimeoutError as e:
+        raise FuturesTimeoutError() from e
+    finally:
+        # Ensure we don't leak pending tasks.
+        for task in tasks:
+            if not task.done():
+                task.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
+
+    return results

{olas_operate_middleware-0.13.6 → olas_operate_middleware-0.13.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "olas-operate-middleware"
-version = "0.13.6"
+version = "0.13.7"
 description = ""
 authors = ["David Vilela <dvilelaf@gmail.com>", "Viraj Patel <vptl185@gmail.com>"]
 readme = "README.md"
@@ -40,6 +40,8 @@ cryptography = "^46.0.3"
 [tool.poetry.group.development.dependencies]
 tomte = {version = "0.4.0", extras = ["cli"]}
 build = "1.2.2.post1"
+yappi = "^1.6.10"
+viztracer = "^1.1.1"
 
 [build-system]
 requires = ["poetry-core"]