plugin system updates

Author: mikeknep

_tmp_notes.md

@@ -0,0 +1,111 @@
+# Plugin system updates
+
+## Requirements
+
+1. Plugins MUST support defining both a configuration object (a Pydantic model) and some `engine`-related implementation object (`ConfigurableTask`, `ColumnGenerator`, etc.).
+1. The UX for making plugins discoverable MUST be simple. We should only require users define a single `Plugin` object that gets referenced in a single entry point.
+1. The plugin system MUST NOT introduce a dependency chain that makes any `config` module depend on any `engine` module.
+  a. Breaks "slim install" support, because `engine` code may include third-party deps that a `config`-only slim install will not include.
+  b. Introduces a high risk of circular imports, because `engine` code depends on `config` modules.
+1. A client using a slim-install of the library SHOULD be able to use plugins.
+
+
+## Current state
+
+The current plugin system violates REQ 3 (and by extension REQ 4):
+
+```
+config.column_types -> data_designer.plugin_manager -> data_designer.plugins.plugin -> data_designer.engine.configurable_task
+```
+(`->` means "imports" aka "depends on")
+
+
+## Blessed engine modules?
+
+One idea that was floated is to refactor existing `engine` code so that base classes exist in some "blessed" module(s) that we would ensure do not create circular imports with `config` modules,
+but this seems...
+- hard to enforce/guarantee
+- potentially restricting (what if a plugin author wants to use objects from other parts of `engine`)
+- conceptually not ideal (it's just simpler to say "`engine` can import/depend on `config`, but not vice-versa" full stop instead of carving out exceptions)
+- potentially complicated with respect to however we restructure packages to support slim installs
+
+
+## Idea proposed in this branch
+
+Make the `Plugin` object "lazy" by defining the config and and task types as fully-qualified strings rather than objects.
+
+By using strings in the `Plugin` object fields, **if** the plugin is structured with multiple files (e.g. `config.py` and `task.py`)*,
+then the core library's `config` code that uses plugins (to extend discriminated union types) can load the plugin and resolve **only**
+the config class type; it would not need to resolve/load/import the plugin's task-related module where `engine` base classes are imported and subclassed.
+
+> *This multi-file setup wouldn't be **required** out of the box; see "Plugin development lifecycle" below.
+
+Example:
+```python
+# src/my_plugin/config.py
+from data_designer.config.column_types import SingleColumnConfig
+
+class MyPluginConfig(SingleColumnConfig):
+    foo: str
+
+
+
+# src/my_plugin/generator.py
+from data_designer.engine.column_generators.generators.base import ColumnGenerator
+from my_plugin.config import MyPluginConfig
+
+class MyPluginGenerator(ColumnGenerator[MyPluginConfig]):
+    pass
+
+
+
+# src/my_plugin/plugin.py
+from data_designer.plugins.plugin import Plugin, PluginType
+
+plugin = Plugin(
+    config_cls="my_plugin.config.MyPluginConfig",
+    task_cls="my_plugin.generator.MyPluginGenerator",
+    plugin_type=PluginType.COLUMN_GENERATOR,
+)
+```
+
+
+### Strings instead of concrete types?
+
+Yeah, a little sad, but seems a reasonable compromise given the benefits this unlocks.
+
+To mitigate against dumb stuff like typos, I suggest we ship a test helper function that we'd encourage plugin authors use in their unit tests:
+```python
+# my_plugin/tests/test_plugin.py
+from data_designer.plugins.test import is_valid_plugin
+from my_plugin.plugin import plugin
+
+
+def test_plugin_validity():
+    assert is_valid_plugin(plugin)
+```
+(Similar to `pd.testing.assert_frame_equal`.)
+
+To start, that test helper would ensure two things:
+1. The string class names resolve to concrete types that do exist
+2. The resolved concrete types are subclasses of the expected base classes
+
+In the future, we could extend the helper to validate other things that are more complex than just Pydantic field type validations.
+
+Remember: we can't implement this validation as a Pydantic validator because it would break the laziness.
+We **can** at least validate that the module exists (and this branch does so), but only the test helper
+can go further and actually fully resolve the two fields.
+
+
+### Plugin development lifecycle
+
+A plugin author _could_ continue defining everything in one Python file and things would still work in the library.
+The limitation would be that a plugin defined that way would not support slim installs, and so clients like NMP would not be able to use it.
+**This might be perfectly fine for many plugins**, especially in the early going.
+A reasonable "plugin development lifecycle" might be:
+1. Develop everything in one file and get it working with the library
+2. Refactor the plugin to support slim installs (if ever desired)
+
+Plugin authors would only need to do step 2 if/when we want to make the plugin available in NMP.
+That step 2 refactor would involve breaking the plugin implementation up into multiple files _and_ (if necessary) making sure any heavyweight,
+task-only third party dependencies are included under an `engine` extra.

src/data_designer/plugins/errors.py

@@ -4,6 +4,9 @@
 from data_designer.errors import DataDesignerError
 
 
+class PluginLoadError(DataDesignerError): ...
+
+
 class PluginRegistrationError(DataDesignerError): ...
 
 

src/data_designer/plugins/plugin.py

@@ -1,14 +1,22 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
+import importlib
+import importlib.util
 from enum import Enum
-from typing import Literal, get_origin
+from functools import cached_property
+from typing import TYPE_CHECKING, Annotated, Literal, get_origin
 
-from pydantic import BaseModel, model_validator
+from pydantic import AfterValidator, BaseModel, model_validator
 from typing_extensions import Self
 
-from data_designer.config.base import ConfigBase
-from data_designer.engine.configurable_task import ConfigurableTask
+from data_designer.plugins.errors import PluginLoadError
+
+if TYPE_CHECKING:
+    from data_designer.config.base import ConfigBase
+    from data_designer.engine.configurable_task import ConfigurableTask
 
 
 class PluginType(str, Enum):
@@ -26,9 +34,30 @@ def display_name(self) -> str:
         return self.value.replace("-", " ")
 
 
+def _get_module_and_object_names(fully_qualified_object: str) -> tuple[str, str]:
+    try:
+        module_name, object_name = fully_qualified_object.rsplit(".", 1)
+    except ValueError:
+        # If fully_qualified_object does not have any periods, the rsplit call will return
+        # a list of length 1 and the variable assignment above will raise ValueError
+        raise PluginLoadError("Expected a fully-qualified object name, e.g. 'my_plugin.config.MyConfig'")
+
+    return module_name, object_name
+
+
+def _is_valid_module(value: str) -> str:
+    module_name, _ = _get_module_and_object_names(value)
+    try:
+        importlib.util.find_spec(module_name)
+    except:
+        raise PluginLoadError(f"Could not find module {module_name!r}.")
+
+    return value
+
+
 class Plugin(BaseModel):
-    task_cls: type[ConfigurableTask]
-    config_cls: type[ConfigBase]
+    task_class_name: Annotated[str, AfterValidator(_is_valid_module)]
+    config_class_name: Annotated[str, AfterValidator(_is_valid_module)]
     plugin_type: PluginType
     emoji: str = "🔌"
 
@@ -50,20 +79,37 @@ def discriminator_field(self) -> str:
 
     @model_validator(mode="after")
     def validate_discriminator_field(self) -> Self:
-        cfg = self.config_cls.__name__
+        _, cfg = _get_module_and_object_names(self.config_class_name)
         field = self.plugin_type.discriminator_field
         if field not in self.config_cls.model_fields:
-            raise ValueError(f"Discriminator field '{field}' not found in config class {cfg}")
+            raise ValueError(f"Discriminator field {field!r} not found in config class {cfg!r}")
         field_info = self.config_cls.model_fields[field]
         if get_origin(field_info.annotation) is not Literal:
-            raise ValueError(f"Field '{field}' of {cfg} must be a Literal type, not {field_info.annotation}.")
+            raise ValueError(f"Field {field!r} of {cfg!r} must be a Literal type, not {field_info.annotation!r}.")
         if not isinstance(field_info.default, str):
-            raise ValueError(f"The default of '{field}' must be a string, not {type(field_info.default)}.")
+            raise ValueError(f"The default of {field!r} must be a string, not {type(field_info.default)!r}.")
         enum_key = field_info.default.replace("-", "_").upper()
         if not enum_key.isidentifier():
             raise ValueError(
-                f"The default value '{field_info.default}' for discriminator field '{field}' "
-                f"cannot be converted to a valid enum key. The converted key '{enum_key}' "
+                f"The default value {field_info.default!r} for discriminator field {field!r} "
+                f"cannot be converted to a valid enum key. The converted key {enum_key!r} "
                 f"must be a valid Python identifier."
             )
         return self
+
+    @cached_property
+    def config_cls(self) -> type[ConfigBase]:
+        return self._load(self.config_class_name)
+
+    @cached_property
+    def task_cls(self) -> type[ConfigurableTask]:
+        return self._load(self.task_class_name)
+
+    @staticmethod
+    def _load(fully_qualified_object: str) -> type:
+        module_name, object_name = _get_module_and_object_names(fully_qualified_object)
+        module = importlib.import_module(module_name)
+        try:
+            return getattr(module, object_name)
+        except AttributeError:
+            raise PluginLoadError(f"Could not find class {object_name!r} in module {module_name!r}")

src/data_designer/plugins/testing/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from data_designer.plugins.testing.utils import is_valid_plugin
+
+__all__ = [
+    is_valid_plugin.__name__,
+]

src/data_designer/plugins/testing/examples.py

@@ -0,0 +1,73 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Literal
+
+from data_designer.config.base import ConfigBase
+from data_designer.config.column_configs import SingleColumnConfig
+from data_designer.engine.configurable_task import ConfigurableTask, ConfigurableTaskMetadata
+
+MODULE_NAME = __name__
+
+
+class ValidTestConfig(SingleColumnConfig):
+    """Valid config for testing plugin creation."""
+
+    column_type: Literal["test-generator"] = "test-generator"
+    name: str
+
+
+class ValidTestTask(ConfigurableTask[ValidTestConfig]):
+    """Valid task for testing plugin creation."""
+
+    @staticmethod
+    def metadata() -> ConfigurableTaskMetadata:
+        return ConfigurableTaskMetadata(
+            name="test_generator",
+            description="Test generator",
+            required_resources=None,
+        )
+
+
+class ConfigWithoutDiscriminator(ConfigBase):
+    some_field: str
+
+
+class ConfigWithStringField(ConfigBase):
+    column_type: str = "test-generator"
+
+
+class ConfigWithNonStringDefault(ConfigBase):
+    column_type: Literal["test-generator"] = 123  # type: ignore
+
+
+class ConfigWithInvalidKey(ConfigBase):
+    column_type: Literal["invalid-key-!@#"] = "invalid-key-!@#"
+
+
+class StubPluginConfigA(SingleColumnConfig):
+    column_type: Literal["test-plugin-a"] = "test-plugin-a"
+
+
+class StubPluginConfigB(SingleColumnConfig):
+    column_type: Literal["test-plugin-b"] = "test-plugin-b"
+
+
+class StubPluginTaskA(ConfigurableTask[StubPluginConfigA]):
+    @staticmethod
+    def metadata() -> ConfigurableTaskMetadata:
+        return ConfigurableTaskMetadata(
+            name="test_plugin_a",
+            description="Test plugin A",
+            required_resources=None,
+        )
+
+
+class StubPluginTaskB(ConfigurableTask[StubPluginConfigB]):
+    @staticmethod
+    def metadata() -> ConfigurableTaskMetadata:
+        return ConfigurableTaskMetadata(
+            name="test_plugin_b",
+            description="Test plugin B",
+            required_resources=None,
+        )

src/data_designer/plugins/testing/utils.py

@@ -0,0 +1,15 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from data_designer.config.base import ConfigBase
+from data_designer.engine.configurable_task import ConfigurableTask
+from data_designer.plugins.plugin import Plugin
+
+
+def is_valid_plugin(plugin: Plugin) -> bool:
+    if not isinstance(plugin.config_cls, ConfigBase):
+        return False
+    if not isinstance(plugin.task_cls, ConfigurableTask):
+        return False
+
+    return True