Merge pull request #176 from amd/alex_invocation

Framework updates: PluginRunInvocation

Author: alexandraBara

nodescraper/cli/__init__.py

@@ -2,7 +2,7 @@
 #
 # MIT License
 #
-# Copyright (c) 2025 Advanced Micro Devices, Inc.
+# Copyright (C) 2026 Advanced Micro Devices, Inc.
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
@@ -25,5 +25,19 @@
 ###############################################################################
 
 from .cli import main as cli_entry
+from .embed import run_main_return_code
+from .invocation import (
+    PluginRunInvocation,
+    get_plugin_run_invocation,
+    plugin_run_invocation_scope,
+    run_plugin_queue_with_invocation,
+)
 
-__all__ = ["cli_entry"]
+__all__ = [
+    "cli_entry",
+    "run_main_return_code",
+    "PluginRunInvocation",
+    "get_plugin_run_invocation",
+    "plugin_run_invocation_scope",
+    "run_plugin_queue_with_invocation",
+]

nodescraper/cli/cli.py

@@ -49,6 +49,7 @@
     process_args,
 )
 from nodescraper.cli.inputargtypes import ModelArgHandler, json_arg, log_path_arg
+from nodescraper.cli.invocation import run_plugin_queue_with_invocation
 from nodescraper.configregistry import ConfigRegistry
 from nodescraper.connection.redfish import (
     RedfishConnection,
@@ -380,11 +381,17 @@ def setup_logger(
     return logger
 
 
-def main(arg_input: Optional[list[str]] = None):
+def main(
+    arg_input: Optional[list[str]] = None,
+    *,
+    host_cli_args: Optional[argparse.Namespace] = None,
+):
     """Main entry point for the CLI
 
     Args:
         arg_input (Optional[list[str]], optional): list of args to parse. Defaults to None.
+        host_cli_args: Optional namespace from an embedding host (e.g. detect-errors) for code that
+            calls get_plugin_run_invocation during the plugin queue.
     """
     if arg_input is None:
         arg_input = sys.argv[1:]
@@ -552,21 +559,23 @@ def main(arg_input: Optional[list[str]] = None):
                 "skip_sudo"
             ] = True
 
-        log_system_info(log_path, system_info, logger)
     except Exception as e:
         parser.error(str(e))
 
-    plugin_executor = PluginExecutor(
-        logger=logger,
-        plugin_configs=plugin_config_inst_list,
-        connections=parsed_args.connection_config,
-        system_info=system_info,
-        log_path=log_path,
-        plugin_registry=plugin_reg,
-    )
-
     try:
-        results = plugin_executor.run_queue()
+        results = run_plugin_queue_with_invocation(
+            plugin_reg=plugin_reg,
+            parsed_args=parsed_args,
+            plugin_config_inst_list=plugin_config_inst_list,
+            system_info=system_info,
+            log_path=log_path,
+            logger=logger,
+            timestamp=timestamp,
+            sname=sname,
+            host_cli_args=host_cli_args,
+        )
+
+        log_system_info(log_path, system_info, logger)
 
         dump_results_to_csv(results, sname, log_path, timestamp, logger)
 

nodescraper/cli/embed.py

@@ -0,0 +1,53 @@
+###############################################################################
+#
+# MIT License
+#
+# Copyright (C) 2026 Advanced Micro Devices, Inc.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+###############################################################################
+"""In-process CLI entry without adding new argparse flags."""
+
+from __future__ import annotations
+
+import argparse
+from typing import Optional
+
+__all__ = ["run_main_return_code"]
+
+
+def run_main_return_code(
+    arg_input: list[str],
+    *,
+    host_cli_args: Optional[argparse.Namespace] = None,
+) -> int:
+    """Runs the nodescraper main entrypoint and maps SystemExit to an integer return code."""
+    from nodescraper.cli.cli import main
+
+    try:
+        main(arg_input, host_cli_args=host_cli_args)
+    except SystemExit as exc:
+        code = exc.code
+        if code is None:
+            return 0
+        if isinstance(code, int):
+            return code
+        return 1
+    return 0

nodescraper/cli/invocation.py

@@ -0,0 +1,117 @@
+###############################################################################
+#
+# MIT License
+#
+# Copyright (C) 2026 Advanced Micro Devices, Inc.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+###############################################################################
+
+from __future__ import annotations
+
+import argparse
+import logging
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass
+from typing import Iterator, Optional
+
+from nodescraper.models import PluginConfig, SystemInfo
+from nodescraper.models.pluginresult import PluginResult
+from nodescraper.pluginexecutor import PluginExecutor
+from nodescraper.pluginregistry import PluginRegistry
+
+_plugin_run_invocation_ctx: ContextVar[Optional["PluginRunInvocation"]] = ContextVar(
+    "nodescraper_plugin_run_invocation", default=None
+)
+
+
+def get_plugin_run_invocation() -> Optional[PluginRunInvocation]:
+    """Return the active invocation while run_plugin_queue_with_invocation is running, if any."""
+    return _plugin_run_invocation_ctx.get()
+
+
+@contextmanager
+def plugin_run_invocation_scope(inv: PluginRunInvocation) -> Iterator[None]:
+    """Bind *inv* for nested code (connection managers, plugins) for the scope of the context."""
+    token = _plugin_run_invocation_ctx.set(inv)
+    try:
+        yield
+    finally:
+        _plugin_run_invocation_ctx.reset(token)
+
+
+@dataclass
+class PluginRunInvocation:
+    """Recorded inputs for one plugin run; optional host_cli_args for embedded hosts."""
+
+    plugin_reg: PluginRegistry
+    parsed_args: argparse.Namespace
+    plugin_config_inst_list: list[PluginConfig]
+    system_info: SystemInfo
+    log_path: Optional[str]
+    logger: logging.Logger
+    timestamp: str
+    sname: str
+    host_cli_args: Optional[argparse.Namespace] = None
+
+
+def run_plugin_queue_with_invocation(
+    *,
+    plugin_reg: PluginRegistry,
+    parsed_args: argparse.Namespace,
+    plugin_config_inst_list: list[PluginConfig],
+    system_info: SystemInfo,
+    log_path: Optional[str],
+    logger: logging.Logger,
+    timestamp: str,
+    sname: str,
+    host_cli_args: Optional[argparse.Namespace] = None,
+) -> list[PluginResult]:
+    """Constructs the plugin executor, binds invocation context, and runs the plugin queue."""
+    inv = PluginRunInvocation(
+        plugin_reg=plugin_reg,
+        parsed_args=parsed_args,
+        plugin_config_inst_list=plugin_config_inst_list,
+        system_info=system_info,
+        log_path=log_path,
+        logger=logger,
+        timestamp=timestamp,
+        sname=sname,
+        host_cli_args=host_cli_args,
+    )
+    plugin_executor = PluginExecutor(
+        logger=logger,
+        plugin_configs=plugin_config_inst_list,
+        connections=parsed_args.connection_config,
+        system_info=system_info,
+        log_path=log_path,
+        plugin_registry=plugin_reg,
+    )
+    with plugin_run_invocation_scope(inv):
+        return plugin_executor.run_queue()
+
+
+__all__ = [
+    "PluginRunInvocation",
+    "get_plugin_run_invocation",
+    "plugin_run_invocation_scope",
+    "run_plugin_queue_with_invocation",
+]

nodescraper/interfaces/datacollectortask.py

@@ -26,6 +26,7 @@
 import abc
 import inspect
 import logging
+from enum import Enum
 from functools import wraps
 from typing import Callable, ClassVar, Generic, Optional, Type, Union
 
@@ -47,6 +48,21 @@
 from .taskresulthook import TaskResultHook
 
 
+def _supported_sku_name_set(supported: Optional[set[Union[str, Enum]]]) -> Optional[set[str]]:
+    """Map ``SUPPORTED_SKUS`` to string names for comparison with ``SystemInfo.sku``."""
+    if not supported:
+        return None
+    names: set[str] = set()
+    for item in supported:
+        if isinstance(item, Enum):
+            names.add(item.name)
+        elif isinstance(item, str):
+            names.add(item)
+        else:
+            names.add(str(item))
+    return names
+
+
 def collect_decorator(
     func: Callable[..., tuple[TaskResult, Optional[TDataModel]]],
 ) -> Callable[..., tuple[TaskResult, Optional[TDataModel]]]:
@@ -111,8 +127,8 @@ class DataCollector(Task, abc.ABC, Generic[TConnection, TDataModel, TCollectArg]
 
     DATA_MODEL: Type[TDataModel]
 
-    # A set of supported SKUs for this data collector
-    SUPPORTED_SKUS: ClassVar[Optional[set[str]]] = None
+    # A set of supported SKUs for this data collector (strings or enum members; enum uses .name)
+    SUPPORTED_SKUS: ClassVar[Optional[set[Union[str, Enum]]]] = None
 
     # A set of supported Platforms for this data collector,
     SUPPORTED_PLATFORMS: ClassVar[Optional[set[str]]] = None
@@ -153,7 +169,12 @@ def __init__(
         self.system_interaction_level = system_interaction_level
         self.connection = connection
 
-        if self.SUPPORTED_SKUS and self.system_info.sku not in self.SUPPORTED_SKUS:
+        allowed_skus = _supported_sku_name_set(self.SUPPORTED_SKUS)
+        if (
+            allowed_skus is not None
+            and self.system_info.sku is not None
+            and self.system_info.sku not in allowed_skus
+        ):
             raise SystemCompatibilityError(
                 f"{self.system_info.sku} SKU is not supported for this collector"
             )

nodescraper/models/systeminfo.py

@@ -39,6 +39,8 @@ class SystemInfo(BaseModel):
     os_family: OSFamily = OSFamily.UNKNOWN
     sku: Optional[str] = None
     platform: Optional[str] = None
+    gpu_count: Optional[int] = None
+    cpu_count: Optional[int] = None
     metadata: Optional[dict] = Field(default_factory=dict)
     location: Optional[SystemLocation] = SystemLocation.LOCAL
     vendorid_ep: int = 0x1002

nodescraper/pluginexecutor.py

@@ -26,6 +26,7 @@
 from __future__ import annotations
 
 import copy
+import inspect
 import logging
 from collections import deque
 from typing import Optional, Type, Union
@@ -160,30 +161,38 @@ def run_queue(self) -> list[PluginResult]:
                     connection_manager_class: Type[ConnectionManager] = plugin_class.CONNECTION_TYPE
                     if (
                         connection_manager_class.__name__
-                        not in self.plugin_registry.connection_managers
+                        in self.plugin_registry.connection_managers
                     ):
+                        mgr_impl = self.plugin_registry.connection_managers[
+                            connection_manager_class.__name__
+                        ]
+                    elif (
+                        inspect.isclass(connection_manager_class)
+                        and issubclass(connection_manager_class, ConnectionManager)
+                        and not inspect.isabstract(connection_manager_class)
+                    ):
+                        # External packages set CONNECTION_TYPE on the plugin;
+                        # use it when not listed under nodescraper.connection_managers entry points.
+                        mgr_impl = connection_manager_class
+                    else:
                         self.logger.error(
                             "Unable to find registered connection manager class for %s that is required by",
                             connection_manager_class.__name__,
                         )
                         continue
 
-                    if connection_manager_class not in self.connection_library:
+                    if mgr_impl not in self.connection_library:
                         self.logger.info(
                             "Initializing connection manager for %s with default args",
-                            connection_manager_class.__name__,
+                            mgr_impl.__name__,
                         )
-                        self.connection_library[connection_manager_class] = (
-                            connection_manager_class(
-                                system_info=self.system_info,
-                                logger=self.logger,
-                                task_result_hooks=self.connection_result_hooks,
-                            )
+                        self.connection_library[mgr_impl] = mgr_impl(
+                            system_info=self.system_info,
+                            logger=self.logger,
+                            task_result_hooks=self.connection_result_hooks,
                         )
 
-                    init_payload["connection_manager"] = self.connection_library[
-                        connection_manager_class
-                    ]
+                    init_payload["connection_manager"] = self.connection_library[mgr_impl]
 
                 try:
                     plugin_inst = plugin_class(**init_payload)