amesa-inference-dev 0.20.5.dev13__tar.gz

amesa_inference_dev-0.20.5.dev13/PKG-INFO

@@ -0,0 +1,119 @@
+Metadata-Version: 2.4
+Name: amesa-inference-dev
+Version: 0.20.5.dev13
+Summary: Agent inference package using ONNX models without Ray or PyTorch dependencies
+Author-email: Hunter Park <hunter@amesa.com>
+Requires-Python: >=3.10, <3.13
+Description-Content-Type: text/markdown
+Requires-Dist: amesa-api-dev>=0.20.5.dev13
+Requires-Dist: amesa-core-dev>=0.20.5.dev13
+Requires-Dist: numpy<2.0.0
+Requires-Dist: onnxruntime
+Requires-Dist: onnx
+Requires-Dist: ruff
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: requires-python
+
+# Amesa Inference
+
+A lightweight inference package for running Amesa agents using ONNX models without Ray or PyTorch dependencies.
+
+## Overview
+
+`composabl_inference` provides a standalone inference engine for running trained Amesa agents. It uses ONNX Runtime for model inference, making it suitable for deployment scenarios where you want to avoid heavy dependencies like Ray and PyTorch.
+
+## Features
+
+- **ONNX-based inference**: Uses ONNX Runtime for efficient model inference
+- **No Ray or PyTorch dependencies**: Lightweight package suitable for production deployment
+- **Network management**: Supports both local and remote objects (skills, perceptors, controllers)
+- **Compatible API**: Similar interface to `Trainer.package()` for easy migration
+
+## Installation
+
+```bash
+pip install amesa-inference
+```
+
+## Usage
+
+### Basic Inference
+
+```python
+from composabl_inference import InferenceEngine
+from composabl_core import Agent
+
+# Create inference engine (only license needed for license validation)
+engine = InferenceEngine(license="your-license-key")
+
+# Load agent
+agent = Agent.load("path/to/agent")
+await engine.load_agent(agent)
+
+# Package agent for inference (similar to Trainer.package())
+await engine.package()
+
+# Run inference
+observation = {...}  # Your observation from the simulator
+action = engine.execute(observation)
+```
+
+### With Remote Objects
+
+The inference engine supports remote skills, perceptors, and controllers, just like the Trainer:
+
+```python
+from composabl_inference import InferenceEngine
+
+# Optional: provide custom config for NetworkMgr (e.g., for remote targets)
+config = {
+    "target": {
+        "local": {
+            "address": "localhost:1337",
+        },
+    },
+}
+
+engine = InferenceEngine(license="your-license-key", config=config)
+await engine.load_agent("path/to/agent")
+await engine.package()
+
+# The skill processor will automatically handle remote objects
+action = engine.execute(observation)
+```
+
+### Cleanup
+
+```python
+# Clean up resources
+await engine.close()
+```
+
+## Architecture
+
+### Components
+
+1. **InferenceEngine**: Main entry point for inference operations
+2. **NetworkMgr**: Manages network connections (non-Ray version)
+3. **ONNXInferenceEngine**: Handles ONNX model loading and inference
+4. **ONNXSkillProcessor**: Processes skills using ONNX models instead of PyTorch
+
+### Differences from Trainer
+
+- Uses ONNX Runtime instead of PyTorch for model inference
+- NetworkMgr is not a Ray actor (runs in the same process)
+- No Ray initialization required
+- Lighter weight, suitable for production deployment
+
+## Requirements
+
+- Python >= 3.10
+- composabl-core
+- composabl-api
+- onnxruntime
+- numpy
+
+## License
+
+Proprietary and confidential - Copyright (C) Amesa, Inc

amesa_inference_dev-0.20.5.dev13/README.md

@@ -0,0 +1,102 @@
+# Amesa Inference
+
+A lightweight inference package for running Amesa agents using ONNX models without Ray or PyTorch dependencies.
+
+## Overview
+
+`composabl_inference` provides a standalone inference engine for running trained Amesa agents. It uses ONNX Runtime for model inference, making it suitable for deployment scenarios where you want to avoid heavy dependencies like Ray and PyTorch.
+
+## Features
+
+- **ONNX-based inference**: Uses ONNX Runtime for efficient model inference
+- **No Ray or PyTorch dependencies**: Lightweight package suitable for production deployment
+- **Network management**: Supports both local and remote objects (skills, perceptors, controllers)
+- **Compatible API**: Similar interface to `Trainer.package()` for easy migration
+
+## Installation
+
+```bash
+pip install amesa-inference
+```
+
+## Usage
+
+### Basic Inference
+
+```python
+from composabl_inference import InferenceEngine
+from composabl_core import Agent
+
+# Create inference engine (only license needed for license validation)
+engine = InferenceEngine(license="your-license-key")
+
+# Load agent
+agent = Agent.load("path/to/agent")
+await engine.load_agent(agent)
+
+# Package agent for inference (similar to Trainer.package())
+await engine.package()
+
+# Run inference
+observation = {...}  # Your observation from the simulator
+action = engine.execute(observation)
+```
+
+### With Remote Objects
+
+The inference engine supports remote skills, perceptors, and controllers, just like the Trainer:
+
+```python
+from composabl_inference import InferenceEngine
+
+# Optional: provide custom config for NetworkMgr (e.g., for remote targets)
+config = {
+    "target": {
+        "local": {
+            "address": "localhost:1337",
+        },
+    },
+}
+
+engine = InferenceEngine(license="your-license-key", config=config)
+await engine.load_agent("path/to/agent")
+await engine.package()
+
+# The skill processor will automatically handle remote objects
+action = engine.execute(observation)
+```
+
+### Cleanup
+
+```python
+# Clean up resources
+await engine.close()
+```
+
+## Architecture
+
+### Components
+
+1. **InferenceEngine**: Main entry point for inference operations
+2. **NetworkMgr**: Manages network connections (non-Ray version)
+3. **ONNXInferenceEngine**: Handles ONNX model loading and inference
+4. **ONNXSkillProcessor**: Processes skills using ONNX models instead of PyTorch
+
+### Differences from Trainer
+
+- Uses ONNX Runtime instead of PyTorch for model inference
+- NetworkMgr is not a Ray actor (runs in the same process)
+- No Ray initialization required
+- Lighter weight, suitable for production deployment
+
+## Requirements
+
+- Python >= 3.10
+- composabl-core
+- composabl-api
+- onnxruntime
+- numpy
+
+## License
+
+Proprietary and confidential - Copyright (C) Amesa, Inc

amesa_inference_dev-0.20.5.dev13/amesa_inference/__init__.py

@@ -0,0 +1,20 @@
+# Copyright (C) Amesa, Inc - All Rights Reserved
+# Unauthorized copying of this file, via any medium is strictly prohibited
+# Proprietary and confidential
+
+from amesa_core.networking.network_mgr import NetworkMgr
+
+from amesa_inference.inference_engine import InferenceEngine
+from amesa_inference.onnx_skill_processor import ONNXSkillProcessor
+from amesa_inference.onnx_selector_processor import ONNXSelectorProcessor
+from amesa_inference.onnx_skill_processor_factory import create_onnx_skill_processor
+from amesa_inference.skill_processor_base import BaseSkillProcessor
+
+__all__ = [
+    "InferenceEngine",
+    "NetworkMgr",
+    "BaseSkillProcessor",
+    "ONNXSkillProcessor",
+    "ONNXSelectorProcessor",
+    "create_onnx_skill_processor",
+]

amesa_inference_dev-0.20.5.dev13/amesa_inference/inference_engine.py

@@ -0,0 +1,189 @@
+# Copyright (C) Amesa, Inc - All Rights Reserved
+# Unauthorized copying of this file, via any medium is strictly prohibited
+# Proprietary and confidential
+
+from typing import Any, Dict, Optional, Union
+
+from amesa_core.settings import TrainerMode
+from amesa_core.utils import license_util
+import amesa_core.utils.logger as logger_util
+from amesa_core.agent.agent import Agent
+from amesa_core.agent.skill.skill import Skill
+from amesa_core.config.trainer_config import TrainerConfig
+from amesa_core.networking.config.skill_processor_context import (
+    SkillProcessorContext,
+)
+from amesa_core.networking.network_mgr import NetworkMgr
+from amesa_inference.inference_network_mgr import InferenceNetworkMgr
+from amesa_core.settings import settings
+
+from amesa_inference.skill_processor_base import BaseSkillProcessor
+
+logger = logger_util.get_logger(__name__)
+
+
+class InferenceEngine:
+    """
+    Inference engine for running agent inference with and without remote objects.
+    Uses ONNX models instead of PyTorch/Ray for inference.
+
+    Supports all skill types including:
+    - Basic skills (ONNXSkillProcessor)
+    - Selector skills (ONNXSelectorProcessor) - automatically orchestrates child skills
+    """
+
+    def __init__(
+        self,
+        license: Optional[str] = None,
+        config: Optional[Dict] = None,
+        network_mgr: Optional[NetworkMgr] = None,
+    ):
+        """
+        Initialize the inference engine.
+
+        Args:
+            license: License key for Amesa (required for license validation)
+            config: Optional configuration dictionary for NetworkMgr (defaults to local target)
+            network_mgr: Optional pre-configured NetworkMgr instance. If not provided,
+                        creates an InferenceNetworkMgr.
+        """
+        # Set license in settings for license validation
+        if license is not None:
+            settings.AMESA_LICENSE = license
+
+        license_type, license_module, license_data = license_util.validate(
+            settings.AMESA_LICENSE,
+            settings.AMESA_ENV == TrainerMode.STAGING,
+            settings.URL_AMESA_LICENSE_SERVER,
+        )
+        logger.info(f"license_validated {license_type}")
+
+        # Create minimal config for NetworkMgr (defaults to local target)
+        if config is None:
+            config = {
+                "target": {
+                    "local": {
+                        "address": "localhost:1337",
+                    }
+                }
+            }
+
+        self.config = TrainerConfig(**config)
+
+        # Use provided network_mgr or create default InferenceNetworkMgr
+        if network_mgr is not None:
+            self.network_mgr = network_mgr
+        else:
+            self.network_mgr = InferenceNetworkMgr(self.config)
+
+        self.agent: Optional[Agent] = None
+        self.skill_processor: Optional[BaseSkillProcessor] = None
+
+    async def load_agent(self, agent: Union[Agent, str]):
+        """
+        Load an agent for inference.
+
+        Args:
+            agent: Agent object or path to agent JSON file
+        """
+        if isinstance(agent, str):
+            self.agent = Agent.load(agent)
+        else:
+            self.agent = agent
+
+        if not self.agent.is_initialized:
+            await self.agent.init()
+
+        logger.info(f"Loaded agent: {self.agent.id}")
+
+    async def package(
+        self, agent: Optional[Agent] = None, skill: Optional[Union[str, Skill]] = None
+    ) -> None:
+        """
+        Package an agent for inference, similar to Trainer.package().
+        Creates an ONNX-based skill processor that can be used for inference.
+
+        Args:
+            agent: Agent object (if None, uses the loaded agent)
+            skill: Specific skill to package (if None, uses top skill)
+        """
+        if agent is None:
+            agent = self.agent
+
+        if agent is None:
+            raise ValueError("No agent provided and no agent loaded")
+
+        if not agent.is_initialized:
+            await agent.init()
+
+        # Determine which skill to use
+        if skill is None:
+            node = agent.get_top_skill()
+        else:
+            if isinstance(skill, str):
+                node = agent.get_node_by_name(skill)
+            else:
+                node = skill
+
+        # Create skill processor context
+        context = SkillProcessorContext(
+            agent=agent,
+            skill=node,
+            network_mgr=self.network_mgr,
+            is_training=False,
+            is_validating=False,
+            for_skill_group=False,
+        )
+
+        # Use factory to create appropriate skill processor based on skill type
+        # This handles selectors, basic skills, and other skill types automatically
+        from amesa_inference.onnx_skill_processor_factory import (
+            create_onnx_skill_processor,
+        )
+
+        self.skill_processor = await create_onnx_skill_processor(context)
+
+        logger.info(
+            f"Packaged agent for inference with skill: {node.get_name()} "
+            f"(processor type: {type(self.skill_processor).__name__})"
+        )
+
+    async def execute(
+        self,
+        obs: Any,
+        sim_action_mask: Optional[Any] = None,
+        explore: bool = False,
+        previous_action: Optional[Any] = None,
+    ) -> Any:
+        """
+        Execute inference on an observation.
+
+        Args:
+            obs: Observation from the simulator
+            sim_action_mask: Optional action mask from simulator
+            explore: Whether to explore (not used in inference, kept for compatibility)
+            previous_action: Previous action (for sequential skills)
+
+        Returns:
+            Action to take
+        """
+        if self.skill_processor is None:
+            raise RuntimeError(
+                "No skill processor available. Call package() first."
+            )
+
+        return await self.skill_processor._execute(
+            obs,
+            sim_action_mask=sim_action_mask,
+            explore=explore,
+            previous_action=previous_action,
+        )
+
+    async def close(self):
+        """Close the inference engine and cleanup resources."""
+        if self.network_mgr is not None:
+            self.network_mgr.stop_all()
+            self.network_mgr.stop_watchdogs()
+
+        logger.info("Inference engine closed")
+

amesa_inference_dev-0.20.5.dev13/amesa_inference/inference_network_mgr.py

@@ -0,0 +1,44 @@
+# Copyright (C) Amesa, Inc - All Rights Reserved
+# Unauthorized copying of this file, via any medium is strictly prohibited
+# Proprietary and confidential
+
+"""
+InferenceNetworkMgr - NetworkMgr subclass for async-first usage.
+
+This module provides a NetworkMgr designed for use with InferenceEngine,
+which operates in async contexts. Unlike the base NetworkMgr which uses
+async_to_sync() (which may case event loop conflicts in async contexts), this
+class returns coroutines that callers await directly.
+"""
+
+import inspect
+from typing import Callable
+
+from amesa_core.networking.network_mgr import NetworkMgr
+
+
+class InferenceNetworkMgr(NetworkMgr):
+    """
+    NetworkMgr subclass designed for async contexts.
+
+    This class overrides _call_method() to return coroutines directly,
+    which callers await. This avoids the event loop conflict entirely.
+
+    Usage:
+        mgr = InferenceNetworkMgr(config)
+        result = await mgr.call_remote_perceptor_mgr("method", ...)
+    """
+
+    def _call_method(self, attr: Callable, *args, **kwargs):
+        """
+        Override base _call_method to return coroutines for async methods.
+
+        Callers must await the result for async methods.
+        """
+        if inspect.iscoroutinefunction(attr):
+            # Return coroutine directly - caller will await
+            return attr(*args, **kwargs)
+        else:
+            # Sync method - call directly
+            return attr(*args, **kwargs)
+

amesa_inference_dev-0.20.5.dev13/amesa_inference/onnx_helper.py

@@ -0,0 +1,152 @@
+# Copyright (C) Amesa, Inc - All Rights Reserved
+# Unauthorized copying of this file, via any medium is strictly prohibited
+# Proprietary and confidential
+
+"""
+ONNX helper utilities for amesa_inference.
+
+This module contains utilities for working with ONNX models during inference,
+including ensuring models have the correct input structure.
+"""
+
+import os
+from typing import Optional
+
+import amesa_core.utils.logger as logger_util
+import onnx
+
+logger = logger_util.get_logger(__name__)
+
+
+def ensure_onnx_model_has_action_mask(
+    onnx_path: str, action_space=None, skill=None
+) -> bool:
+    """
+    Ensure an ONNX model has an action_mask input. If it doesn't, attempt to add it.
+
+    This function is used after downloading ONNX models to ensure they have the
+    action_mask input required for inference, especially for selectors.
+
+    Args:
+        onnx_path: Path to the ONNX model file
+        action_space: Optional action space to determine action_mask shape
+        skill: Optional skill object to get action space from
+
+    Returns:
+        True if the model was modified, False if it already had action_mask
+    """
+    try:
+        onnx_model = onnx.load(onnx_path)
+
+        # Check if action_mask input already exists
+        input_names = [input.name for input in onnx_model.graph.input]
+        has_action_mask = any("action_mask" in name.lower() for name in input_names)
+
+        if has_action_mask:
+            logger.debug(f"ONNX model at {onnx_path} already has action_mask input")
+            return False
+
+        # Determine action_mask shape
+        action_mask_shape = None
+        if action_space is not None:
+            try:
+                action_mask_space = action_space.get_action_mask_space()
+                if hasattr(action_mask_space, 'shape'):
+                    action_mask_shape = action_mask_space.shape
+                elif hasattr(action_space, 'n'):
+                    action_mask_shape = (action_space.n,)
+                elif hasattr(action_space, 'shape'):
+                    action_mask_shape = action_space.shape
+            except Exception as e:
+                logger.warning(f"Could not get action_mask shape from action_space: {e}")
+
+        # If we have a skill, try to get action space from it
+        if action_mask_shape is None and skill is not None:
+            try:
+                skill_action_space = skill.get_action_space()
+                if skill_action_space is not None:
+                    action_mask_space = skill_action_space.get_action_mask_space()
+                    if hasattr(action_mask_space, 'shape'):
+                        action_mask_shape = action_mask_space.shape
+                    elif hasattr(skill_action_space, 'n'):
+                        action_mask_shape = (skill_action_space.n,)
+            except Exception as e:
+                logger.warning(f"Could not get action_mask shape from skill: {e}")
+
+        # Try to infer from action_dist_inputs output shape
+        if action_mask_shape is None:
+            try:
+                output_names = [output.name for output in onnx_model.graph.output]
+                if "action_dist_inputs" in output_names:
+                    # Find the output and get its shape
+                    for output in onnx_model.graph.output:
+                        if output.name == "action_dist_inputs":
+                            shape = [
+                                dim.dim_value if dim.dim_value > 0 else 1
+                                for dim in output.type.tensor_type.shape.dim
+                            ]
+                            if len(shape) >= 2:
+                                # Shape is [batch, action_size], we want action_size
+                                action_mask_shape = (shape[-1],)
+                            elif len(shape) == 1:
+                                action_mask_shape = (shape[0],)
+                            break
+            except Exception as e:
+                logger.warning(f"Could not infer action_mask shape from model outputs: {e}")
+
+        # Fallback to shape (1,)
+        if action_mask_shape is None:
+            logger.warning("Could not determine action_mask shape, using default (1,)")
+            action_mask_shape = (1,)
+
+        # Get observation input to determine batch dimension
+        obs_input = None
+        for input in onnx_model.graph.input:
+            if "observation" in input.name.lower() or input.name == "observation":
+                obs_input = input
+                break
+
+        if obs_input is None and len(onnx_model.graph.input) > 0:
+            obs_input = onnx_model.graph.input[0]
+
+        batch_dim = None
+        if obs_input is not None:
+            try:
+                # Get batch dimension from observation input (usually first dimension)
+                obs_shape = obs_input.type.tensor_type.shape.dim
+                if len(obs_shape) > 0:
+                    batch_dim = obs_shape[0]
+            except Exception:
+                pass
+
+        # Create action_mask input
+        # Use dynamic batch dimension if available, otherwise use 1
+        if batch_dim is not None:
+            action_mask_dims = [batch_dim]
+        else:
+            action_mask_dims = [onnx.TensorShapeProto.Dimension(dim_value=1)]
+
+        for dim_size in action_mask_shape:
+            action_mask_dims.append(onnx.TensorShapeProto.Dimension(dim_value=dim_size))
+
+        action_mask_type = onnx.TypeProto()
+        action_mask_type.tensor_type.elem_type = onnx.TensorProto.DOUBLE
+        action_mask_type.tensor_type.shape.CopyFrom(onnx.TensorShapeProto(dim=action_mask_dims))
+
+        action_mask_value_info = onnx.ValueInfoProto()
+        action_mask_value_info.name = "action_mask"
+        action_mask_value_info.type.CopyFrom(action_mask_type)
+
+        # Add action_mask input to the graph
+        onnx_model.graph.input.append(action_mask_value_info)
+
+        # Validate and save
+        onnx.checker.check_model(onnx_model)
+        onnx.save(onnx_model, onnx_path)
+
+        logger.info(f"Added action_mask input with shape {action_mask_shape} to ONNX model at {onnx_path}")
+        return True
+
+    except Exception as e:
+        logger.error(f"Failed to ensure action_mask input in ONNX model at {onnx_path}: {e}")
+        raise