Skip to content

Commit e32d1b6

Browse files
Feature:4005 Docker settings usage warnings
- Warn when used without containerized orchestrator in place
1 parent 1bf2d9e commit e32d1b6

File tree

3 files changed

+109
-3
lines changed

3 files changed

+109
-3
lines changed

src/zenml/config/docker_settings.py

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
from pydantic import BaseModel, ConfigDict, Field, model_validator
2020

2121
from zenml.config.base_settings import BaseSettings
22-
from zenml.logger import get_logger
22+
from zenml.logger import get_logger, simple_warning_format
2323
from zenml.utils import deprecation_utils
2424
from zenml.utils.pydantic_utils import before_validator_handler
2525

@@ -394,6 +394,38 @@ def _deprecate_replicate_local_environment_commands(
394394
)
395395
return self
396396

397+
@model_validator(mode="after")
398+
def _docker_settings_usage_warnings(self) -> "DockerSettings":
399+
"""Checks active environment and warns for potential issues.
400+
401+
Returns:
402+
The validated settings values.
403+
"""
404+
import warnings
405+
406+
from zenml.client import Client
407+
from zenml.orchestrators import ContainerizedOrchestrator
408+
409+
warning_tag = "non_containerized_orchestrator"
410+
411+
active_orchestrator = Client().active_stack.orchestrator
412+
413+
if (
414+
isinstance(active_orchestrator, ContainerizedOrchestrator)
415+
or warning_tag in _docker_settings_warnings_logged
416+
):
417+
return self
418+
419+
with simple_warning_format():
420+
warnings.warn(
421+
f"WARNING: You are specifying docker settings without a containerized orchestrator: "
422+
f"{active_orchestrator.__class__.__name__}"
423+
)
424+
425+
_docker_settings_warnings_logged.append(warning_tag)
426+
427+
return self
428+
397429
model_config = ConfigDict(
398430
# public attributes are immutable
399431
frozen=True,

src/zenml/logger.py

Lines changed: 44 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,10 @@
1818
import os
1919
import re
2020
import sys
21+
import warnings
22+
from contextlib import contextmanager
2123
from contextvars import ContextVar
22-
from typing import TYPE_CHECKING, Any, Dict
24+
from typing import TYPE_CHECKING, Any, Dict, Generator, Type
2325

2426
if TYPE_CHECKING:
2527
from zenml.logging.step_logging import ArtifactStoreHandler
@@ -357,6 +359,9 @@ def init_logging() -> None:
357359
for handler in root_logger.handlers
358360
)
359361

362+
# logging capture warnings
363+
logging.captureWarnings(True)
364+
360365
if not has_console_handler:
361366
console_handler = logging.StreamHandler(sys.stdout)
362367
console_handler.setFormatter(get_formatter())
@@ -393,3 +398,41 @@ def init_logging() -> None:
393398
for logger_name in disabled_logger_names:
394399
logging.getLogger(logger_name).setLevel(logging.WARNING)
395400
logging.getLogger(logger_name).disabled = True
401+
402+
403+
def simple_warning_formatter(
404+
message: Warning | str,
405+
category: Type[Warning],
406+
filename: str,
407+
lineno: int,
408+
line: str | None = None,
409+
) -> str:
410+
"""Simplified warning formatter that ignores context and filename.
411+
412+
Args:
413+
message : The warning instance or string to display.
414+
category : The class of the warning (e.g., `UserWarning`).
415+
filename : Name of the file where the warning occurred. (Ignored in this formatter.)
416+
lineno : Line number where the warning occurred. (Ignored in this formatter.)
417+
line : The line of source code to show. Defaults to None. (Ignored in this formatter.)
418+
419+
Returns:
420+
Warning message.
421+
"""
422+
return f"{category.__name__}: {message}\n"
423+
424+
425+
@contextmanager
426+
def simple_warning_format() -> Generator[None, None, None]:
427+
"""Warning utility: Simplify formatting for specific calls (Skip path rendering).
428+
429+
Yields:
430+
Empty generator, used as context manager.
431+
"""
432+
original = warnings.formatwarning
433+
434+
warnings.formatwarning = simple_warning_formatter
435+
try:
436+
yield
437+
finally:
438+
warnings.formatwarning = original

tests/unit/config/test_docker_settings.py

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,11 @@
1616
import pytest
1717
from pydantic import ValidationError
1818

19-
from zenml.config import DockerSettings
19+
from zenml.config.docker_settings import (
20+
DockerSettings,
21+
_docker_settings_warnings_logged,
22+
)
23+
from zenml.orchestrators import ContainerizedOrchestrator
2024

2125

2226
def test_build_skipping():
@@ -28,3 +32,30 @@ def test_build_skipping():
2832
with does_not_raise():
2933
DockerSettings(skip_build=False)
3034
DockerSettings(skip_build=True, parent_image="my_parent_image")
35+
36+
37+
def test_generated_warnings():
38+
from zenml.client import Client
39+
40+
if isinstance(
41+
Client().active_stack.orchestrator, ContainerizedOrchestrator
42+
):
43+
pytest.skip(reason="Check warning generation for local orchestrators")
44+
45+
if "non_containerized_orchestrator" in _docker_settings_warnings_logged:
46+
_docker_settings_warnings_logged.remove(
47+
"non_containerized_orchestrator"
48+
)
49+
50+
with pytest.warns(UserWarning) as warning:
51+
DockerSettings()
52+
53+
assert (
54+
"You are specifying docker settings without a containerized orchestrator"
55+
in str(warning[0].message)
56+
)
57+
58+
with pytest.warns(None):
59+
DockerSettings()
60+
61+
assert "non_containerized_orchestrator" in _docker_settings_warnings_logged

0 commit comments

Comments
 (0)