PR review fixes & improvements

Author: Json-Andriopoulos

src/zenml/config/compiler.py

@@ -385,14 +385,9 @@ def _validate_docker_settings_usage(
             stack: The stack the settings are validated against.
 
         """
-        from zenml.orchestrators import (
-            ContainerizedOrchestrator,
-            LocalOrchestrator,
-        )
+        from zenml.orchestrators import LocalOrchestrator
 
-        if not docker_settings or isinstance(
-            stack.orchestrator, ContainerizedOrchestrator
-        ):
+        if not docker_settings:
             return
 
         warning_message = (
@@ -406,10 +401,6 @@ def _validate_docker_settings_usage(
             WARNING_CONTROLLER.info(
                 warning_code=WarningCodes.ZML002, message=warning_message
             )
-        else:
-            WARNING_CONTROLLER.warn(
-                warning_code=WarningCodes.ZML002, message=warning_message
-            )
 
     def _filter_and_validate_settings(
         self,

src/zenml/utils/warnings/base.py

@@ -13,26 +13,26 @@
 #  permissions and limitations under the License.
 """Warning configuration class and helper enums."""
 
-from enum import Enum
-
 from pydantic import BaseModel, Field
 
+from zenml.utils.enum_utils import StrEnum
+
 
-class WarningSeverity(str, Enum):
+class WarningSeverity(StrEnum):
     """Enum class describing the warning severity."""
 
     LOW = "low"
     MEDIUM = "medium"
     HIGH = "high"
 
 
-class WarningCategory(str, Enum):
+class WarningCategory(StrEnum):
     """Enum class describing the warning category."""
 
     USAGE = "USAGE"
 
 
-class WarningVerbosity(str, Enum):
+class WarningVerbosity(StrEnum):
     """Enum class describing the warning verbosity."""
 
     LOW = "low"

src/zenml/utils/warnings/controller.py

@@ -14,12 +14,13 @@
 """Module for centralized WarningController implementation."""
 
 import logging
-from collections import defaultdict
+from collections import Counter
 from typing import Any
 
 from zenml.enums import LoggingLevels
 from zenml.utils.singleton import SingletonMetaClass
 from zenml.utils.warnings.base import WarningConfig, WarningVerbosity
+from zenml.utils.warnings.registry import WarningCodes
 
 logger = logging.getLogger(__name__)
 
@@ -30,7 +31,7 @@ class WarningController(metaclass=SingletonMetaClass):
     def __init__(self) -> None:
         """WarningController constructor."""
         self._warning_configs: dict[str, WarningConfig] = {}
-        self._warning_statistics: dict[str, int] = defaultdict(int)
+        self._warning_statistics: dict[str, int] = Counter()
 
     def register(self, warning_configs: dict[str, WarningConfig]) -> None:
         """Register a warning config collection to the controller.
@@ -81,7 +82,7 @@ def _get_display_message(
 
     def _log(
         self,
-        warning_code: str,
+        warning_code: WarningCodes,
         message: str,
         level: LoggingLevels,
         **kwargs: dict[str, Any],
@@ -125,7 +126,9 @@ def _log(
             # Assumes warning level is the default if an invalid option is passed.
             logger.warning(display_message.format(**kwargs))
 
-    def warn(self, *, warning_code: str, message: str, **kwargs: Any) -> None:
+    def warn(
+        self, *, warning_code: WarningCodes, message: str, **kwargs: Any
+    ) -> None:
         """Method to execute warning handling logic with warning log level.
 
         Args:
@@ -135,7 +138,9 @@ def warn(self, *, warning_code: str, message: str, **kwargs: Any) -> None:
         """
         self._log(warning_code, message, LoggingLevels.WARNING, **kwargs)
 
-    def info(self, *, warning_code: str, message: str, **kwargs: Any) -> None:
+    def info(
+        self, *, warning_code: WarningCodes, message: str, **kwargs: Any
+    ) -> None:
         """Method to execute warning handling logic with info log level.
 
         Args:

src/zenml/utils/warnings/registry.py

@@ -13,8 +13,7 @@
 #  permissions and limitations under the License.
 """Module for warning configurations organization and resolution."""
 
-from enum import Enum
-
+from zenml.utils.enum_utils import StrEnum
 from zenml.utils.warnings.base import (
     WarningCategory,
     WarningConfig,
@@ -23,7 +22,7 @@
 )
 
 
-class WarningCodes(str, Enum):
+class WarningCodes(StrEnum):
     """Enum class organizing the warning codes."""
 
     ZML001 = "ZML001"

tests/unit/utils/test_warnings.py

@@ -2,8 +2,8 @@
 from pydantic import ValidationError
 
 from zenml.utils.warnings.base import (
-    WarningConfig,
     WarningCategory,
+    WarningConfig,
     WarningSeverity,
     WarningVerbosity,
 )
@@ -16,7 +16,7 @@ def controller() -> WarningController:
 
     We reuse the singleton instance but clear its internal state to ensure test isolation.
     """
-    c = WarningController()
+    c = WarningController.create()
     # hard reset of internal state
     c._warning_configs.clear()
     c._warning_statistics.clear()
@@ -109,7 +109,9 @@ def test_config_validations_reject_invalid_values():
         )
 
 
-def test_registration_merges_configs(controller: WarningController, sample_configs: dict[str, WarningConfig]):
+def test_registration_merges_configs(
+    controller: WarningController, sample_configs: dict[str, WarningConfig]
+):
     controller.register(sample_configs)
     assert controller._warning_configs["TEST001"].code == "TEST001"
     assert controller._warning_configs["TEST002"].is_throttled is True
@@ -120,7 +122,9 @@ def test_singleton_instance_is_shared(controller: WarningController):
     assert controller is another
 
 
-def test_statistics_increment_for_non_throttled(controller: WarningController, sample_configs: dict[str, WarningConfig]):
+def test_statistics_increment_for_non_throttled(
+    controller: WarningController, sample_configs: dict[str, WarningConfig]
+):
     controller.register(sample_configs)
 
     controller.warn(warning_code="TEST001", message="once")
@@ -130,7 +134,9 @@ def test_statistics_increment_for_non_throttled(controller: WarningController, s
     assert controller._warning_statistics["TEST001"] == 3
 
 
-def test_statistics_throttled_only_once(controller: WarningController, sample_configs: dict[str, WarningConfig]):
+def test_statistics_throttled_only_once(
+    controller: WarningController, sample_configs: dict[str, WarningConfig]
+):
     controller.register(sample_configs)
 
     controller.warn(warning_code="TEST002", message="first")
@@ -140,7 +146,9 @@ def test_statistics_throttled_only_once(controller: WarningController, sample_co
     assert controller._warning_statistics["TEST002"] == 1
 
 
-def test_warn_and_info_do_not_break_with_format_kwargs(controller, sample_configs):
+def test_warn_and_info_do_not_break_with_format_kwargs(
+    controller: WarningController, sample_configs: dict[str, WarningConfig]
+):
     controller.register(sample_configs)
 
     controller.warn(
@@ -159,24 +167,39 @@ def test_warn_and_info_do_not_break_with_format_kwargs(controller, sample_config
     assert controller._warning_statistics["TEST001"] == 2
 
 
-def test_unregistered_warning_no_crash_and_no_stats_increment(controller):
-    controller.warn(warning_code="UNKNOWN", message="should fallback to default")
+def test_unregistered_warning_no_crash_and_no_stats_increment(
+    controller: WarningController,
+):
+    controller.warn(
+        warning_code="UNKNOWN", message="should fallback to default"
+    )
     controller.info(warning_code="UNKNOWN", message="also fallback")
 
     assert "UNKNOWN" not in controller._warning_statistics
 
 
-def test_get_display_message_varies_by_verbosity(sample_configs):
+def test_get_display_message_varies_by_verbosity(
+    sample_configs: dict[str, WarningConfig],
+):
     from zenml.utils.warnings.controller import WarningController as WC
+
     module_name = "mod.path"
     line_number = 42
 
-    low = sample_configs["TEST001"].model_copy(update={"verbosity": WarningVerbosity.LOW})
-    med = sample_configs["TEST001"].model_copy(update={"verbosity": WarningVerbosity.MEDIUM})
-    high = sample_configs["TEST001"].model_copy(update={"verbosity": WarningVerbosity.HIGH})
+    low = sample_configs["TEST001"].model_copy(
+        update={"verbosity": WarningVerbosity.LOW}
+    )
+    med = sample_configs["TEST001"].model_copy(
+        update={"verbosity": WarningVerbosity.MEDIUM}
+    )
+    high = sample_configs["TEST001"].model_copy(
+        update={"verbosity": WarningVerbosity.HIGH}
+    )
 
     msg_low = WC._get_display_message("hello", module_name, line_number, low)
-    assert msg_low.startswith("[TEST001](WarningCategory.USAGE)") or msg_low.startswith("[TEST001](USAGE)")
+    assert msg_low.startswith(
+        "[TEST001](WarningCategory.USAGE)"
+    ) or msg_low.startswith("[TEST001](USAGE)")
     assert "hello" in msg_low
     assert "mod.path" not in msg_low and "42" not in msg_low
     assert low.description not in msg_low
@@ -191,3 +214,26 @@ def test_get_display_message_varies_by_verbosity(sample_configs):
     assert "[TEST001]" in msg_high and "hello" in msg_high
     assert "\n" in msg_high
     assert high.description.strip() in msg_high
+
+
+def test_combined_registries(
+    controller: WarningController, sample_configs: dict[str, WarningConfig]
+):
+    from zenml.utils.warnings.registry import (
+        WARNING_CONFIG_REGISTRY,
+        WarningCodes,
+    )
+
+    controller.register(sample_configs)
+    controller.register(WARNING_CONFIG_REGISTRY)
+
+    assert WarningCodes.ZML001.value in controller._warning_configs
+    assert "TEST001" in controller._warning_configs
+
+    controller.warn(
+        warning_code=WarningCodes.ZML001,
+        message="You are warned {name}",
+        name="tester",
+    )
+
+    assert controller._warning_statistics["ZML001"] >= 1