madsci.node_module 0.7.0__tar.gz → 0.7.1__tar.gz

{madsci_node_module-0.7.0 → madsci_node_module-0.7.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: madsci.node_module
-Version: 0.7.0
+Version: 0.7.1
 Summary: The Modular Autonomous Discovery for Science (MADSci) Node Module Helper Classes.
 Author-Email: Tobias Ginsburg <tginsburg@anl.gov>, "Ryan D. Lewis" <ryan.lewis@anl.gov>, Casey Stone <cstone@anl.gov>, Doga Ozgulbas <dozgulbas@anl.gov>
 License: MIT

{madsci_node_module-0.7.0 → madsci_node_module-0.7.1}/madsci/node_module/abstract_node_module.py

@@ -1,8 +1,11 @@
 """Base Node Module helper classes."""
 
+import atexit
 import contextlib
 import inspect
+import logging
 import threading
+import weakref
 from pathlib import Path
 from typing import (
     Annotated,
@@ -115,6 +118,12 @@ class AbstractNode(MadsciClientMixin):
             module_name=module_name,
         )
 
+        # Resolve stable identity from registry if enabled
+        self._resolver = None
+        self._atexit_registered = False
+        if self.config.enable_registry_resolution:
+            self._resolve_identity_from_registry()
+
         global_ownership_info.node_id = self.node_info.node_id
         self._configure_clients()
 
@@ -136,6 +145,23 @@ class AbstractNode(MadsciClientMixin):
     """Node Lifecycle and Public Methods"""
     """------------------------------------------------------------------------------------------------"""
 
+    def close(self) -> None:
+        """Release registry identity and clean up client resources.
+
+        This method is idempotent and safe to call multiple times.
+        Call this before reassigning a node variable to a new node
+        with the same name (e.g. in notebook cells) to avoid
+        ``RegistryLockError`` from lingering heartbeat threads.
+        """
+        self._release_registry_identity()
+        self._resolver = None
+        self.teardown_clients()
+
+    def __del__(self) -> None:
+        """Best-effort cleanup when the node is garbage-collected."""
+        with contextlib.suppress(Exception):
+            self.close()
+
     def start_node(self) -> None:
         """
         Called once to start the node.
@@ -170,6 +196,74 @@ class AbstractNode(MadsciClientMixin):
     def shutdown_handler(self) -> None:
         """Called to shut down the node. Should be used to clean up any resources."""
 
+    def _resolve_identity_from_registry(self) -> None:
+        """Resolve node identity from the ID Registry.
+
+        Uses IdentityResolver to look up or create a stable ID for this
+        node.  The resolved ID is written back to ``self.node_info.node_id``.
+
+        A ``RegistryLockError`` (after retry exhaustion) is fatal — the
+        node cannot start without a stable identity.  Other errors
+        (e.g. missing registry file) are non-fatal: a warning is logged
+        and the generated ID is kept.
+        """
+        from madsci.common.registry.identity_resolver import (  # noqa: PLC0415
+            IdentityResolver,
+        )
+        from madsci.common.registry.lock_manager import (  # noqa: PLC0415
+            RegistryLockError,
+        )
+
+        try:
+            self._resolver = IdentityResolver(lab_url=self.config.lab_url)
+
+            node_name = self.node_info.node_name
+
+            result = self._resolver.resolve_with_info(
+                name=node_name,
+                component_type="node",
+                metadata={"module_name": self.node_info.module_name},
+                retry_timeout=self.config.registry_lock_timeout,
+            )
+
+            self.node_info.node_id = result.id
+
+            if not self._atexit_registered:
+                ref = weakref.ref(self)
+
+                def _release_on_exit(weak_ref: weakref.ref = ref) -> None:
+                    obj = weak_ref()
+                    if obj is not None:
+                        obj._release_registry_identity()
+
+                atexit.register(_release_on_exit)
+                self._atexit_registered = True
+
+        except RegistryLockError:
+            raise
+        except Exception:
+            logging.getLogger(__name__).warning(
+                "Registry identity resolution failed; using generated ID",
+                exc_info=True,
+            )
+
+    def _release_registry_identity(self) -> None:
+        """Release the registry lock for this node's identity.
+
+        Called via ``atexit`` to allow other instances to acquire the
+        same name.
+        """
+        if self._resolver is None:
+            return
+
+        try:
+            self._resolver.release(self.node_info.node_name)
+        except Exception:
+            logging.getLogger(__name__).warning(
+                "Failed to release registry identity",
+                exc_info=True,
+            )
+
     """------------------------------------------------------------------------------------------------"""
     """Interface Methods"""
     """------------------------------------------------------------------------------------------------"""

{madsci_node_module-0.7.0 → madsci_node_module-0.7.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "madsci.node_module"
-version = "0.7.0"
+version = "0.7.1"
 description = "The Modular Autonomous Discovery for Science (MADSci) Node Module Helper Classes."
 authors = [
     { name = "Tobias Ginsburg", email = "tginsburg@anl.gov" },

{madsci_node_module-0.7.0 → madsci_node_module-0.7.1}/tests/conftest.py

@@ -117,10 +117,13 @@ def test_node_factory() -> Generator[Callable[..., TestNode], None, None]:
             Configured TestNode instance
         """
         # Create base config with identity fields
+        # Disable registry resolution by default in tests to prevent
+        # lock contention and side effects from shared registry files.
         config_params = {
             "test_required_param": 1,
             "node_name": node_name,
             "module_name": module_name,
+            "enable_registry_resolution": False,
             **config_kwargs,
         }
         if config_overrides:

{madsci_node_module-0.7.0 → madsci_node_module-0.7.1}/tests/test_argument_parsing.py

@@ -30,7 +30,7 @@ class TestArgumentParsingNode(RestNode):
     """Test node for argument parsing tests."""
 
     __test__ = False
-    config: TestConfig = TestConfig()
+    config: TestConfig = TestConfig(enable_registry_resolution=False)
     config_model = TestConfig
 
     def startup_handler(self) -> None:

madsci_node_module-0.7.1/tests/test_node_registry.py

@@ -0,0 +1,215 @@
+"""Unit tests for registry resolution in AbstractNode."""
+
+from unittest.mock import MagicMock, patch
+
+from madsci.common.types.node_types import NodeConfig
+from madsci.common.types.registry_types import RegistryResolveResult
+from madsci.common.utils import new_ulid_str
+
+from madsci_node_module.tests.test_node import TestNode, TestNodeConfig
+
+
+class TestNodeRegistrySettings:
+    """Test the registry-related fields on NodeConfig."""
+
+    def test_registry_resolution_enabled_by_default(self) -> None:
+        """Registry resolution is enabled by default in production.
+
+        Note: The test suite patches the default to False to prevent lock
+        contention. We verify explicit True still works.
+        """
+        config = NodeConfig(enable_registry_resolution=True)
+        assert config.enable_registry_resolution is True
+
+    def test_lab_url_none_by_default(self) -> None:
+        """Lab URL should be None by default."""
+        config = NodeConfig()
+        assert config.lab_url is None
+
+    def test_registry_lock_timeout_default(self) -> None:
+        """Registry lock timeout should default to 60 seconds."""
+        config = NodeConfig()
+        assert config.registry_lock_timeout == 60.0
+
+    def test_can_disable_registry_resolution(self) -> None:
+        """Should be able to disable registry resolution."""
+        config = NodeConfig(enable_registry_resolution=False)
+        assert config.enable_registry_resolution is False
+
+
+class TestNodeRegistryResolution:
+    """Test the registry resolution integration in AbstractNode."""
+
+    def test_resolver_not_called_when_disabled(self) -> None:
+        """When registry resolution is disabled, resolver should not be called."""
+        config = TestNodeConfig(
+            test_required_param=1,
+            node_name="test_node",
+            enable_registry_resolution=False,
+        )
+        node = TestNode(node_config=config)
+
+        assert node._resolver is None
+
+    def test_resolver_called_when_enabled(self) -> None:
+        """When registry resolution is enabled, resolver should update node_info.node_id."""
+        resolved_id = new_ulid_str()
+        mock_resolver = MagicMock()
+        mock_resolver.resolve_with_info.return_value = RegistryResolveResult(
+            name="test_node",
+            id=resolved_id,
+            component_type="node",
+            is_new=False,
+            source="local",
+        )
+        mock_resolver_cls = MagicMock(return_value=mock_resolver)
+
+        config = TestNodeConfig(
+            test_required_param=1,
+            node_name="test_node",
+            enable_registry_resolution=True,
+        )
+
+        with patch(
+            "madsci.common.registry.identity_resolver.IdentityResolver",
+            mock_resolver_cls,
+        ):
+            node = TestNode(node_config=config)
+
+        assert node.node_info.node_id == resolved_id
+
+    def test_resolution_failure_falls_back_to_generated_id(self) -> None:
+        """If registry resolution fails, the generated ID should be preserved."""
+        config = TestNodeConfig(
+            test_required_param=1,
+            node_name="test_node",
+            enable_registry_resolution=True,
+        )
+
+        with patch(
+            "madsci.common.registry.identity_resolver.IdentityResolver",
+            side_effect=Exception("Connection refused"),
+        ):
+            node = TestNode(node_config=config)
+
+        # The node should still have a valid node_id (the originally generated one)
+        assert node.node_info.node_id is not None
+        assert len(node.node_info.node_id) > 0
+
+    def test_stable_id_across_restarts(self) -> None:
+        """Two sequential node instances with the same name should get the same ID."""
+        stable_id = new_ulid_str()
+        mock_resolver = MagicMock()
+        mock_resolver.resolve_with_info.return_value = RegistryResolveResult(
+            name="stable_node",
+            id=stable_id,
+            component_type="node",
+            is_new=False,
+            source="local",
+        )
+        mock_resolver_cls = MagicMock(return_value=mock_resolver)
+
+        with patch(
+            "madsci.common.registry.identity_resolver.IdentityResolver",
+            mock_resolver_cls,
+        ):
+            node1 = TestNode(
+                node_config=TestNodeConfig(
+                    test_required_param=1,
+                    node_name="stable_node",
+                    enable_registry_resolution=True,
+                )
+            )
+            node1._release_registry_identity()
+
+            node2 = TestNode(
+                node_config=TestNodeConfig(
+                    test_required_param=1,
+                    node_name="stable_node",
+                    enable_registry_resolution=True,
+                )
+            )
+
+        assert node1.node_info.node_id == stable_id
+        assert node2.node_info.node_id == stable_id
+
+    def test_registry_disabled_generates_ulid(self) -> None:
+        """With enable_registry_resolution=False, a fresh ULID should be generated."""
+        node1 = TestNode(
+            node_config=TestNodeConfig(
+                test_required_param=1,
+                node_name="ulid_node",
+                enable_registry_resolution=False,
+            )
+        )
+        node2 = TestNode(
+            node_config=TestNodeConfig(
+                test_required_param=1,
+                node_name="ulid_node",
+                enable_registry_resolution=False,
+            )
+        )
+
+        # Both should have IDs, but they should be different (fresh ULIDs)
+        assert node1.node_info.node_id is not None
+        assert node2.node_info.node_id is not None
+        assert node1.node_info.node_id != node2.node_info.node_id
+
+    def test_release_identity_calls_resolver_release(self) -> None:
+        """_release_registry_identity() should call resolver.release()."""
+        config = TestNodeConfig(
+            test_required_param=1,
+            node_name="release_test_node",
+            enable_registry_resolution=False,
+        )
+        node = TestNode(node_config=config)
+        node._resolver = MagicMock()
+
+        node._release_registry_identity()
+
+        node._resolver.release.assert_called_once_with("release_test_node")
+
+    def test_release_identity_noop_without_resolver(self) -> None:
+        """_release_registry_identity() should be a no-op when no resolver is set."""
+        config = TestNodeConfig(
+            test_required_param=1,
+            node_name="noop_node",
+            enable_registry_resolution=False,
+        )
+        node = TestNode(node_config=config)
+
+        # Should not raise
+        node._release_registry_identity()
+
+    def test_node_name_used_for_resolution(self) -> None:
+        """The node name should be used for registry lookup."""
+        resolved_id = new_ulid_str()
+        mock_resolver = MagicMock()
+        mock_resolver.resolve_with_info.return_value = RegistryResolveResult(
+            name="custom_node",
+            id=resolved_id,
+            component_type="node",
+            is_new=True,
+            source="local",
+        )
+        mock_resolver_cls = MagicMock(return_value=mock_resolver)
+
+        config = TestNodeConfig(
+            test_required_param=1,
+            node_name="custom_node",
+            module_name="custom_module",
+            enable_registry_resolution=True,
+        )
+
+        with patch(
+            "madsci.common.registry.identity_resolver.IdentityResolver",
+            mock_resolver_cls,
+        ):
+            TestNode(node_config=config)
+
+        mock_resolver.resolve_with_info.assert_called_once_with(
+            name="custom_node",
+            component_type="node",
+            metadata={"module_name": "custom_module"},
+            retry_timeout=60.0,
+        )