Adding step name to console logs

docs/book/how-to/control-logging/disable-step-names-in-logs.md

@@ -0,0 +1,47 @@
+---
+description: How to disable step names from being displayed in logs.
+---
+
+# Disable step names in logs
+
+By default, ZenML displays the current step name as a prefix in logs during pipeline execution to help you identify which step is currently running. For example:
+
+```
+[data_loader] Loading data from source...
+[data_loader] Data loaded successfully.
+[model_trainer] Training model with parameters...
+```
+
+It's important to note that these step name prefixes are only added to the console output and are not stored in the actual log files that ZenML saves. The stored logs will always be clean without the step name prefixes.
+
+Disabling step names in logs is particularly useful when you have multiple pipeline steps running simultaneously, as the step prefixes can sometimes make the logs harder to read when steps are interleaved. In such cases, you might prefer clean logs without prefixes.
+
+If you wish to disable this feature and have logs without step name prefixes, you can do so by setting the following environment variable:
+
+```bash
+ZENML_DISABLE_STEP_NAMES_IN_LOGS=true
+```
+
+Note that setting this on the [client environment](../pipeline-development/configure-python-environments/README.md#client-environment-or-the-runner-environment) (e.g. your local machine which runs the pipeline) will only affect locally executed pipelines. If you wish to disable step names in logs for remote pipeline runs, you can set the `ZENML_DISABLE_STEP_NAMES_IN_LOGS` environment variable in your pipeline runs environment as follows:
+
+```python
+from zenml import pipeline
+from zenml.config import DockerSettings
+
+docker_settings = DockerSettings(environment={"ZENML_DISABLE_STEP_NAMES_IN_LOGS": "true"})
+
+# Either add it to the decorator
+@pipeline(settings={"docker": docker_settings})
+def my_pipeline() -> None:
+    my_step()
+
+# Or configure the pipelines options
+my_pipeline = my_pipeline.with_options(
+    settings={"docker": docker_settings}
+)
+```
+
+When this option is enabled, logs will be displayed without step name prefixes, making them cleaner but potentially harder to associate with specific steps in your pipeline.
+
+<!-- For scarf -->
+<figure><img alt="ZenML Scarf" referrerpolicy="no-referrer-when-downgrade" src="https://static.scarf.sh/a.png?x-pxid=f0b4f458-0a54-4fcd-aa95-d5ee424815bc" /></figure>

docs/book/toc.md

@@ -170,6 +170,7 @@
   * [Set logging format](how-to/control-logging/set-logging-format.md)
   * [Disable rich traceback output](how-to/control-logging/disable-rich-traceback.md)
   * [Disable colorful logging](how-to/control-logging/disable-colorful-logging.md)
+  * [Disable step names in logs](how-to/control-logging/disable-step-names-in-logs.md)
 * [Popular integrations](how-to/popular-integrations/README.md)
   * [Run on AWS](how-to/popular-integrations/aws-guide.md)
   * [Run on GCP](how-to/popular-integrations/gcp-guide.md)

src/zenml/constants.py

@@ -169,6 +169,7 @@ def handle_int_env_var(var: str, default: int = 0) -> int:
 ENV_ZENML_ENFORCE_TYPE_ANNOTATIONS = "ZENML_ENFORCE_TYPE_ANNOTATIONS"
 ENV_ZENML_ENABLE_IMPLICIT_AUTH_METHODS = "ZENML_ENABLE_IMPLICIT_AUTH_METHODS"
 ENV_ZENML_DISABLE_STEP_LOGS_STORAGE = "ZENML_DISABLE_STEP_LOGS_STORAGE"
+ENV_ZENML_DISABLE_STEP_NAMES_IN_LOGS = "ZENML_DISABLE_STEP_NAMES_IN_LOGS"
 ENV_ZENML_IGNORE_FAILURE_HOOK = "ZENML_IGNORE_FAILURE_HOOK"
 ENV_ZENML_CUSTOM_SOURCE_ROOT = "ZENML_CUSTOM_SOURCE_ROOT"
 ENV_ZENML_WHEEL_PACKAGE_NAME = "ZENML_WHEEL_PACKAGE_NAME"

src/zenml/logging/step_logging.py

@@ -22,11 +22,16 @@
 from typing import Any, Callable, List, Optional, Type, Union
 from uuid import UUID, uuid4
 
+from zenml import get_step_context
 from zenml.artifact_stores import BaseArtifactStore
 from zenml.artifacts.utils import (
     _load_artifact_store,
     _load_file_from_artifact_store,
 )
+from zenml.constants import (
+    ENV_ZENML_DISABLE_STEP_NAMES_IN_LOGS,
+    handle_bool_env_var,
+)
 from zenml.exceptions import DoesNotExistException
 from zenml.logger import get_logger
 from zenml.logging import (
@@ -447,8 +452,8 @@ def __enter__(self) -> "StepLogsStorageContext":
         setattr(sys.stdout, "write", self._wrap_write(self.stdout_write))
         setattr(sys.stdout, "flush", self._wrap_flush(self.stdout_flush))
 
-        setattr(sys.stderr, "write", self._wrap_write(self.stdout_write))
-        setattr(sys.stderr, "flush", self._wrap_flush(self.stdout_flush))
+        setattr(sys.stderr, "write", self._wrap_write(self.stderr_write))
+        setattr(sys.stderr, "flush", self._wrap_flush(self.stderr_flush))
 
         redirected.set(True)
         return self
@@ -494,9 +499,32 @@ def _wrap_write(self, method: Callable[..., Any]) -> Callable[..., Any]:
         """
 
         def wrapped_write(*args: Any, **kwargs: Any) -> Any:
-            output = method(*args, **kwargs)
+            # Check if step names in logs are disabled via env var
+            step_names_disabled = handle_bool_env_var(
+                ENV_ZENML_DISABLE_STEP_NAMES_IN_LOGS, default=False
+            )
+
+            if step_names_disabled:
+                output = method(*args, **kwargs)
+            else:
+                # Try to get step context if not available yet
+                step_context = None
+                try:
+                    step_context = get_step_context()
+                except Exception:
+                    pass
+
+                if step_context and args[0] != "\n":
+                    message = f"[{step_context.step_name}] " + args[0]
+                else:
+                    message = args[0]
+
+                output = method(message, *args[1:], **kwargs)
+
+            # Save the original message without step name prefix to storage
             if args:
                 self.storage.write(args[0])
+
             return output
 
         return wrapped_write