Merge pull request #703 from NVIDIA/am/entry-agents

Allow agents to be registered via entry points

Author: amaslenn

src/cloudai/cli/cli.py

@@ -238,9 +238,8 @@ def verify_configs(configs_dir: Path, tests_dir: Path):
 
 
 @main.command()
-@click.argument("type", type=click.Choice(["reports"]))
+@click.argument("type", type=click.Choice(["reports", "agents"], case_sensitive=False))
 @click.option("-v", "--verbose", is_flag=True, default=False, help="Verbose output.")
 def list(type: str, verbose: bool):
     """List available in Registry items."""
-    args = argparse.Namespace(type=type, verbose=verbose)
-    handle_list_registered_items(args)
+    handle_list_registered_items(type, verbose)

src/cloudai/cli/handlers.py

@@ -539,15 +539,21 @@ def load_tomls_by_type(tomls: List[Path]) -> dict[str, List[Path]]:
     return files
 
 
-def handle_list_registered_items(args: argparse.Namespace) -> int:
-    item_type = args.type
+def handle_list_registered_items(item_type: str, verbose: bool) -> int:
     registry = Registry()
-    if item_type == "reports":
-        print("Registered scenario reports:")
+    if item_type.lower() == "reports":
+        print("Available scenario reports:")
         for idx, (name, report) in enumerate(sorted(registry.scenario_reports.items()), start=1):
             str = f'{idx}. "{name}" {report.__name__}'
-            if args.verbose:
+            if verbose:
                 str += f" (config={registry.report_configs[name].model_dump_json(indent=None)})"
             print(str)
+    elif item_type.lower() == "agents":
+        print("Available agents:")
+        for idx, (name, agent) in enumerate(sorted(registry.agents_map.items()), start=1):
+            str = f'{idx}. "{name}" class={agent.__name__}'
+            if verbose:
+                str += f"{agent.__doc__}"
+            print(str)
 
     return 0

src/cloudai/registration.py

@@ -14,6 +14,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import warnings
+from importlib.metadata import entry_points
+
+
+def register_entrypoint_agents():
+    from cloudai.configurator.base_agent import BaseAgent
+    from cloudai.core import Registry
+
+    eps = entry_points(group="cloudai.agents")
+    for ep in eps:
+        cls = ep.load()
+        if issubclass(cls, BaseAgent):
+            Registry().add_agent(ep.name, cls)
+        else:
+            warnings.warn(
+                f"Skipping entrypoint: {ep.name} -> {ep.value} class={cls} (not a subclass of BaseAgent)", stacklevel=2
+            )
+
 
 def register_all():
     """Register all workloads, systems, runners, installers, and strategies."""
@@ -233,3 +251,5 @@ def register_all():
     Registry().add_reward_function("ai_dynamo_weighted_normalized", ai_dynamo_weighted_normalized_reward)
     Registry().add_reward_function("ai_dynamo_ratio_normalized", ai_dynamo_ratio_normalized_reward)
     Registry().add_reward_function("ai_dynamo_log_scale", ai_dynamo_log_scale_reward)
+
+    register_entrypoint_agents()

tests/test_registry.py

@@ -15,6 +15,8 @@
 # limitations under the License.
 
 import copy
+from typing import Any
+from unittest.mock import patch
 
 import pytest
 
@@ -31,6 +33,7 @@
 )
 from cloudai.models.scenario import ReportConfig
 from cloudai.models.workload import TestDefinition
+from cloudai.registration import register_entrypoint_agents
 
 
 class MyTestDefinition(TestDefinition):
@@ -322,3 +325,20 @@ def test_get_command_gen_strategy_not_found(self, registry: Registry):
         with pytest.raises(KeyError) as exc_info:
             registry.get_command_gen_strategy(MySystem, AnotherTestDefinition)
         assert exc_info.match("Command gen strategy for 'MySystem, AnotherTestDefinition' not found.")
+
+
+def test_entrypoint_agent_type_verified():
+    class MockEP:
+        def __init__(self, load_value: Any):
+            self._load_value = load_value
+            self.name = "name"
+            self.value = "value"
+
+        def load(self):
+            return self._load_value
+
+    with (
+        patch("cloudai.registration.entry_points", return_value=[MockEP(str)]),
+        pytest.warns(UserWarning, match="(not a subclass of BaseAgent)"),
+    ):
+        register_entrypoint_agents()