snowflake-ml-python 1.31.0__tar.gz → 1.33.0__tar.gz

{snowflake_ml_python-1.31.0 → snowflake_ml_python-1.33.0}/CHANGELOG.md

@@ -1,6 +1,52 @@
 # Release History
 
-## 1.31.0
+## 1.33.0
+
+### New Features
+
+* Registry: Extended `ParamSpec` support to PyTorch models. Parameters declared in the model signature
+  can now be passed at inference time, matching the existing support for custom models.
+* Registry (PrPr): Added `model_init_once` model save option to `log_model`. When set to `True`, the
+  model is loaded once per worker process at startup, eliminating model-loading overhead. Defaults to `False`.
+* Experiment Tracking live logging is out of private preview
+  and will become generally available over the next few weeks.
+* Registry: Bumped `transformers` upper bound to `<6`, adding compatibility with Hugging Face Transformers v5.
+
+### Bug Fixes
+
+* Experiment Tracking live logging: Fixed a bug where the tracebacks for uncaught exceptions were not logged when
+  the code was running in a `with exp.start_run()` block.
+* Registry: Fixed a thread leak in `create_service(block=True)` where the log-streaming thread could outlive the
+  error handler when service deployment fails.
+
+### Behavior Changes
+
+### Deprecations
+
+## 1.32.0 (2026-03-26)
+
+### New Features
+
+* Registry: Extended `ParamSpec` support to scikit-learn models. Parameters declared in the model signature
+  can now be passed at inference time, matching the existing support for custom models.
+* Registry: `infer_signature()` now accepts a `dict` for the `params` argument (e.g.,
+  `params={"temperature": 0.7, "max_tokens": 100}`), automatically inferring `ParamSpec` objects
+  from the Python value types. The existing `Sequence[BaseParamSpec]` form is still supported.
+
+### Bug Fixes
+
+* Experiment Tracking live logging (PrPr): Fixed a bug where the tracebacks for uncaught exceptions were not logged
+  when the code was running inside Snowflake notebooks.
+
+### Behavior Changes
+
+* Registry: The models from `huggingface.TransformersPipeline` and `transformers.Pipeline` will default to safetensors
+  file format by default by using `safe_serialization=True` parameter to the `save_pretrained` call. This creates
+  safetensors files instead of PyTorch binaries or pt files, addressing security vulnerability CVE-2025-32434.
+
+### Deprecations
+
+## 1.31.0 (2026-03-19)
 
 ### New Features
 
@@ -18,7 +64,7 @@
 * Registry: Fixed a bug where invalid parameter types (e.g., passing a string to an integer parameter)
   were silently coerced instead of raising an error. Parameters are now strictly validated against their
   declared types before inference. This may break code that relied on the silent coercion behavior
-  (e.g., `model.run(max_tokens="100")` or `model.run(some_int_param=True)`).
+  (e.g., `mv.run(..., params={"max_tokens": "100"})` or `mv.run(..., params={"some_int_param": True})`).
 * ML Job: Fixed explicitly named `MLJobDefinition` objects so they can be overwritten in place and invoked
   repeatedly without job name collisions; the `generate_suffix` argument has been removed.
 * Experiment Tracking live logging (PrPr): Fixed a bug that generated too many lines of log that consist entirely of

{snowflake_ml_python-1.31.0/snowflake_ml_python.egg-info → snowflake_ml_python-1.33.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowflake-ml-python
-Version: 1.31.0
+Version: 1.33.0
 Summary: The machine learning client library that is used for interacting with Snowflake to build machine learning solutions.
 Author-email: "Snowflake, Inc" <support@snowflake.com>
 License: 
@@ -266,14 +266,14 @@ Requires-Dist: keras<4,>=2.0.0; extra == "all"
 Requires-Dist: lightgbm<5,>=4.1.0; extra == "all"
 Requires-Dist: mlflow<3,>=2.16.0; extra == "all"
 Requires-Dist: prophet<2,>=1.1.0; extra == "all"
-Requires-Dist: sentence-transformers<4,>=2.7.0; extra == "all"
+Requires-Dist: sentence-transformers<6,>=2.7.0; extra == "all"
 Requires-Dist: sentencepiece<0.2.0,>=0.1.95; extra == "all"
 Requires-Dist: streamlit<2,>=1.30.0; extra == "all"
 Requires-Dist: tensorflow<3,>=2.17.0; extra == "all"
 Requires-Dist: tokenizers<1,>=0.15.1; extra == "all"
 Requires-Dist: torch<3,>=2.0.1; extra == "all"
 Requires-Dist: torchdata<1,>=0.4; extra == "all"
-Requires-Dist: transformers!=4.51.3,<5,>=4.39.3; extra == "all"
+Requires-Dist: transformers!=4.51.3,<6,>=4.39.3; extra == "all"
 Provides-Extra: altair
 Requires-Dist: altair<6,>=5; extra == "altair"
 Provides-Extra: catboost
@@ -286,12 +286,12 @@ Provides-Extra: lightgbm
 Requires-Dist: lightgbm<5,>=4.1.0; extra == "lightgbm"
 Provides-Extra: llm
 Requires-Dist: mlflow<3,>=2.16.0; extra == "llm"
-Requires-Dist: sentence-transformers<4,>=2.7.0; extra == "llm"
+Requires-Dist: sentence-transformers<6,>=2.7.0; extra == "llm"
 Requires-Dist: sentencepiece<0.2.0,>=0.1.95; extra == "llm"
 Requires-Dist: tokenizers<1,>=0.15.1; extra == "llm"
 Requires-Dist: torch<3,>=2.0.1; extra == "llm"
 Requires-Dist: torchdata<1,>=0.4; extra == "llm"
-Requires-Dist: transformers!=4.51.3,<5,>=4.39.3; extra == "llm"
+Requires-Dist: transformers!=4.51.3,<6,>=4.39.3; extra == "llm"
 Provides-Extra: mlflow
 Requires-Dist: mlflow<3,>=2.16.0; extra == "mlflow"
 Provides-Extra: prophet
@@ -304,11 +304,11 @@ Provides-Extra: torch
 Requires-Dist: torch<3,>=2.0.1; extra == "torch"
 Requires-Dist: torchdata<1,>=0.4; extra == "torch"
 Provides-Extra: transformers
-Requires-Dist: sentence-transformers<4,>=2.7.0; extra == "transformers"
+Requires-Dist: sentence-transformers<6,>=2.7.0; extra == "transformers"
 Requires-Dist: sentencepiece<0.2.0,>=0.1.95; extra == "transformers"
 Requires-Dist: tokenizers<1,>=0.15.1; extra == "transformers"
 Requires-Dist: torch<3,>=2.0.1; extra == "transformers"
-Requires-Dist: transformers!=4.51.3,<5,>=4.39.3; extra == "transformers"
+Requires-Dist: transformers!=4.51.3,<6,>=4.39.3; extra == "transformers"
 Dynamic: license-file
 
 # Snowflake ML Python
@@ -416,7 +416,53 @@ NOTE: Version 1.7.0 is used as example here. Please choose the the latest versio
 
 # Release History
 
-## 1.31.0
+## 1.33.0
+
+### New Features
+
+* Registry: Extended `ParamSpec` support to PyTorch models. Parameters declared in the model signature
+  can now be passed at inference time, matching the existing support for custom models.
+* Registry (PrPr): Added `model_init_once` model save option to `log_model`. When set to `True`, the
+  model is loaded once per worker process at startup, eliminating model-loading overhead. Defaults to `False`.
+* Experiment Tracking live logging is out of private preview
+  and will become generally available over the next few weeks.
+* Registry: Bumped `transformers` upper bound to `<6`, adding compatibility with Hugging Face Transformers v5.
+
+### Bug Fixes
+
+* Experiment Tracking live logging: Fixed a bug where the tracebacks for uncaught exceptions were not logged when
+  the code was running in a `with exp.start_run()` block.
+* Registry: Fixed a thread leak in `create_service(block=True)` where the log-streaming thread could outlive the
+  error handler when service deployment fails.
+
+### Behavior Changes
+
+### Deprecations
+
+## 1.32.0 (2026-03-26)
+
+### New Features
+
+* Registry: Extended `ParamSpec` support to scikit-learn models. Parameters declared in the model signature
+  can now be passed at inference time, matching the existing support for custom models.
+* Registry: `infer_signature()` now accepts a `dict` for the `params` argument (e.g.,
+  `params={"temperature": 0.7, "max_tokens": 100}`), automatically inferring `ParamSpec` objects
+  from the Python value types. The existing `Sequence[BaseParamSpec]` form is still supported.
+
+### Bug Fixes
+
+* Experiment Tracking live logging (PrPr): Fixed a bug where the tracebacks for uncaught exceptions were not logged
+  when the code was running inside Snowflake notebooks.
+
+### Behavior Changes
+
+* Registry: The models from `huggingface.TransformersPipeline` and `transformers.Pipeline` will default to safetensors
+  file format by default by using `safe_serialization=True` parameter to the `save_pretrained` call. This creates
+  safetensors files instead of PyTorch binaries or pt files, addressing security vulnerability CVE-2025-32434.
+
+### Deprecations
+
+## 1.31.0 (2026-03-19)
 
 ### New Features
 
@@ -434,7 +480,7 @@ NOTE: Version 1.7.0 is used as example here. Please choose the the latest versio
 * Registry: Fixed a bug where invalid parameter types (e.g., passing a string to an integer parameter)
   were silently coerced instead of raising an error. Parameters are now strictly validated against their
   declared types before inference. This may break code that relied on the silent coercion behavior
-  (e.g., `model.run(max_tokens="100")` or `model.run(some_int_param=True)`).
+  (e.g., `mv.run(..., params={"max_tokens": "100"})` or `mv.run(..., params={"some_int_param": True})`).
 * ML Job: Fixed explicitly named `MLJobDefinition` objects so they can be overwritten in place and invoked
   repeatedly without job name collisions; the `generate_suffix` argument has been removed.
 * Experiment Tracking live logging (PrPr): Fixed a bug that generated too many lines of log that consist entirely of

{snowflake_ml_python-1.31.0 → snowflake_ml_python-1.33.0}/pyproject.toml

@@ -24,18 +24,18 @@ Issues = "https://github.com/snowflakedb/snowflake-ml-python/issues"
 Changelog = "https://github.com/snowflakedb/snowflake-ml-python/blob/master/CHANGELOG.md"
 
 [project.optional-dependencies]
-all = [ "altair>=5,<6", "catboost>=1.2.0, <2", "keras>=2.0.0,<4", "lightgbm>=4.1.0, <5", "mlflow>=2.16.0, <3", "prophet>=1.1.0, <2", "sentence-transformers>=2.7.0,<4", "sentencepiece>=0.1.95,<0.2.0", "streamlit>=1.30.0,<2", "tensorflow>=2.17.0,<3", "tokenizers>=0.15.1,<1", "torch>=2.0.1,<3", "torchdata>=0.4,<1", "transformers>=4.39.3,<5,!=4.51.3",]
+all = [ "altair>=5,<6", "catboost>=1.2.0, <2", "keras>=2.0.0,<4", "lightgbm>=4.1.0, <5", "mlflow>=2.16.0, <3", "prophet>=1.1.0, <2", "sentence-transformers>=2.7.0,<6", "sentencepiece>=0.1.95,<0.2.0", "streamlit>=1.30.0,<2", "tensorflow>=2.17.0,<3", "tokenizers>=0.15.1,<1", "torch>=2.0.1,<3", "torchdata>=0.4,<1", "transformers>=4.39.3,!=4.51.3,<6",]
 altair = [ "altair>=5,<6",]
 catboost = [ "catboost>=1.2.0, <2",]
 keras = [ "keras>=2.0.0,<4", "tensorflow>=2.17.0,<3", "torch>=2.0.1,<3",]
 lightgbm = [ "lightgbm>=4.1.0, <5",]
-llm = [ "mlflow>=2.16.0, <3", "sentence-transformers>=2.7.0,<4", "sentencepiece>=0.1.95,<0.2.0", "tokenizers>=0.15.1,<1", "torch>=2.0.1,<3", "torchdata>=0.4,<1", "transformers>=4.39.3,<5,!=4.51.3",]
+llm = [ "mlflow>=2.16.0, <3", "sentence-transformers>=2.7.0,<6", "sentencepiece>=0.1.95,<0.2.0", "tokenizers>=0.15.1,<1", "torch>=2.0.1,<3", "torchdata>=0.4,<1", "transformers>=4.39.3,!=4.51.3,<6",]
 mlflow = [ "mlflow>=2.16.0, <3",]
 prophet = [ "prophet>=1.1.0, <2",]
 streamlit = [ "streamlit>=1.30.0,<2",]
 tensorflow = [ "tensorflow>=2.17.0,<3",]
 torch = [ "torch>=2.0.1,<3", "torchdata>=0.4,<1",]
-transformers = [ "sentence-transformers>=2.7.0,<4", "sentencepiece>=0.1.95,<0.2.0", "tokenizers>=0.15.1,<1", "torch>=2.0.1,<3", "transformers>=4.39.3,<5,!=4.51.3",]
+transformers = [ "sentence-transformers>=2.7.0,<6", "sentencepiece>=0.1.95,<0.2.0", "tokenizers>=0.15.1,<1", "torch>=2.0.1,<3", "transformers>=4.39.3,!=4.51.3,<6",]
 
 [tool.setuptools.package-data]
 "*" = [ "*",]

{snowflake_ml_python-1.31.0 → snowflake_ml_python-1.33.0}/snowflake/ml/_internal/env_utils.py

@@ -591,18 +591,104 @@ def save_conda_env_file(
         yaml.safe_dump(env, stream=f, default_flow_style=False)
 
 
-def save_requirements_file(path: pathlib.Path, pip_deps: list[requirements.Requirement]) -> None:
+def _get_pytorch_cuda_suffix(cuda_version: str) -> str:
+    """Get PyTorch CUDA wheel suffix (e.g., '124') for the specified CUDA version.
+
+    Args:
+        cuda_version: CUDA version string.
+
+    Returns:
+        The CUDA suffix for PyTorch wheels.
+    """
+    parsed = version.parse(cuda_version)
+    return f"{parsed.major}{parsed.minor}"
+
+
+def _get_pytorch_cuda_index_url(cuda_version: str) -> str:
+    """Get PyTorch CUDA index URL for the specified CUDA version.
+
+    Args:
+        cuda_version: CUDA version string.
+
+    Returns:
+        The --extra-index-url argument for pip/uv.
+    """
+    cuda_suffix = _get_pytorch_cuda_suffix(cuda_version)
+    return f"--extra-index-url https://download.pytorch.org/whl/cu{cuda_suffix}"
+
+
+def _pin_torch_to_cuda_variant(
+    pip_deps: list[requirements.Requirement], cuda_version: str
+) -> list[requirements.Requirement]:
+    """Pin torch requirement to the CUDA variant (e.g., torch==2.5.0+cu124).
+
+    PyTorch CUDA wheels use local version identifiers like +cu124 to distinguish
+    CUDA-enabled builds from CPU-only builds.
+
+    Args:
+        pip_deps: Validated pip requirements.
+        cuda_version: CUDA version string (e.g., "12.4", "12.4.1").
+
+    Returns:
+        A new list of requirements with ``torch==X.Y.Z`` rewritten to ``torch==X.Y.Z+cu...``
+        when the requirement is exactly pinned and has no ``+cu`` suffix yet.
+    """
+    cuda_suffix = _get_pytorch_cuda_suffix(cuda_version)
+    result = []
+
+    for req in pip_deps:
+        if req.name != "torch":
+            result.append(req)
+            continue
+
+        # Extract pinned version (e.g., ==2.5.0)
+        pinned_version = next((s.version for s in req.specifier if s.operator == "=="), None)
+
+        # Skip if not pinned or already has +cu suffix
+        if not pinned_version or "+cu" in pinned_version:
+            result.append(req)
+            continue
+
+        # Create new requirement with CUDA local version
+        new_req = requirements.Requirement(f"torch=={pinned_version}+cu{cuda_suffix}")
+        new_req.extras = req.extras
+        new_req.marker = req.marker
+        result.append(new_req)
+
+    return result
+
+
+def _has_torch_dependency(pip_deps: list[requirements.Requirement]) -> bool:
+    """Check if torch is in the pip dependencies."""
+    return any(dep.name == "torch" for dep in pip_deps)
+
+
+def save_requirements_file(
+    path: pathlib.Path,
+    pip_deps: list[requirements.Requirement],
+    cuda_version: Optional[str] = None,
+) -> None:
     """Generate Python requirements.txt file in the given directory path.
 
     Args:
         path: Path to the requirements.txt file.
-        pip_deps: List of dependencies string after validated.
+        pip_deps: List of dependencies after validation.
+        cuda_version: Optional CUDA version for GPU models. If provided and torch is in
+            dependencies, adds PyTorch CUDA index URL and pins torch to CUDA variant.
     """
     assert path.suffix in [".txt"], "PIP requirement file should have extension of txt."
     path.parent.mkdir(parents=True, exist_ok=True)
-    requirements = "\n".join(map(str, pip_deps))
+
+    lines: list[str] = []
+
+    if cuda_version and _has_torch_dependency(pip_deps):
+        lines.append(_get_pytorch_cuda_index_url(cuda_version))
+        pip_deps = _pin_torch_to_cuda_variant(pip_deps, cuda_version)
+
+    lines.extend(map(str, pip_deps))
+
     with open(path, "w", encoding="utf-8") as out:
-        out.write(requirements)
+        out.write("\n".join(lines))
 
 
 def load_conda_env_file(
@@ -702,6 +788,9 @@ def load_requirements_file(path: pathlib.Path) -> list[requirements.Requirement]
     with open(path, encoding="utf-8") as f:
         reqs = f.readlines()
 
+    # Filter out pip options
+    reqs = [line for line in reqs if line.strip() and not line.strip().startswith(("-", "#"))]
+
     return validate_pip_requirement_string_list(reqs)
 
 

{snowflake_ml_python-1.31.0 → snowflake_ml_python-1.33.0}/snowflake/ml/experiment/_client/experiment_tracking_sql_client.py

@@ -203,6 +203,36 @@ class ExperimentTrackingSQLClient(_base._BaseSQLClient):
             self._session, f"SHOW RUNS {like_clause} IN EXPERIMENT {experiment_fqn}"
         ).validate()
 
+    @telemetry.send_api_usage_telemetry(project=telemetry.TelemetryProject.EXPERIMENT_TRACKING.value)
+    def show_run_metrics_in_experiment(
+        self,
+        *,
+        experiment_name: sql_identifier.SqlIdentifier,
+        run_name: Optional[sql_identifier.SqlIdentifier] = None,
+        like: Optional[str] = None,
+    ) -> list[row.Row]:
+        experiment_fqn = self.fully_qualified_object_name(self._database_name, self._schema_name, experiment_name)
+        run_name_clause = f"RUN {run_name}" if run_name else ""
+        like_clause = f"LIKE '{like}'" if like else ""
+        return query_result_checker.SqlResultValidator(
+            self._session, f"SHOW RUN METRICS {like_clause} IN EXPERIMENT {experiment_fqn} {run_name_clause}"
+        ).validate()
+
+    @telemetry.send_api_usage_telemetry(project=telemetry.TelemetryProject.EXPERIMENT_TRACKING.value)
+    def show_run_parameters_in_experiment(
+        self,
+        *,
+        experiment_name: sql_identifier.SqlIdentifier,
+        run_name: Optional[sql_identifier.SqlIdentifier] = None,
+        like: Optional[str] = None,
+    ) -> list[row.Row]:
+        experiment_fqn = self.fully_qualified_object_name(self._database_name, self._schema_name, experiment_name)
+        run_name_clause = f"RUN {run_name}" if run_name else ""
+        like_clause = f"LIKE '{like}'" if like else ""
+        return query_result_checker.SqlResultValidator(
+            self._session, f"SHOW RUN PARAMETERS {like_clause} IN EXPERIMENT {experiment_fqn} {run_name_clause}"
+        ).validate()
+
     def _build_snow_uri(
         self, experiment_name: sql_identifier.SqlIdentifier, run_name: sql_identifier.SqlIdentifier, artifact_path: str
     ) -> str:

{snowflake_ml_python-1.31.0 → snowflake_ml_python-1.33.0}/snowflake/ml/experiment/_entities/run.py

@@ -38,10 +38,12 @@ class Run:
         self,
         exc_type: Optional[type[BaseException]],
         exc_value: Optional[BaseException],
-        traceback: Optional[types.TracebackType],
+        tb: Optional[types.TracebackType],
     ) -> None:
-        self._patcher.__exit__(exc_type, exc_value, traceback)
+        self._patcher.__exit__(exc_type, exc_value, tb)
         if self._experiment_tracking._run is self:
+            if exc_type is not None:
+                self._experiment_tracking._try_log_exception(exc_type, exc_value, tb)
             self._experiment_tracking.end_run()
 
     def _get_experiment_info(self) -> experiment_info.ExperimentInfo:

{snowflake_ml_python-1.31.0 → snowflake_ml_python-1.33.0}/snowflake/ml/experiment/_logging/experiment_logging_context.py

@@ -1,5 +1,6 @@
 import contextlib
 from dataclasses import dataclass
+from typing import Any, Callable, Optional
 
 from snowflake.ml._internal.utils import tee
 from snowflake.ml.experiment._logging.experiment_logger import ExperimentLogger
@@ -11,3 +12,5 @@ class ExperimentLoggingContext:
     stderr_logger: ExperimentLogger
     stdout_ctx: contextlib.redirect_stdout[tee.OutputTee]
     stderr_ctx: contextlib.redirect_stderr[tee.OutputTee]
+    ipython_instance: Optional[Any]
+    original_showtraceback: Optional[Callable[..., None]]

{snowflake_ml_python-1.31.0 → snowflake_ml_python-1.33.0}/snowflake/ml/experiment/experiment_tracking.py

@@ -2,11 +2,13 @@ import contextlib
 import functools
 import json
 import sys
+import traceback
 import warnings
+from collections.abc import Callable
+from types import TracebackType
 from typing import Any, Optional, Union
 from urllib.parse import quote
 
-import snowflake.snowpark._internal.utils as snowpark_utils
 from snowflake import snowpark
 from snowflake.ml import model as ml_model, registry
 from snowflake.ml._internal import env_utils
@@ -535,7 +537,6 @@ class ExperimentTracking:
                 target_path=local_dir,
             )
 
-    @snowpark_utils.private_preview(version="1.30.0")
     def set_live_logging_status(self, enabled: bool) -> None:
         """
         Enable or disable live logging. When enabled, stdout and stderr are captured and logged to persistent storage.
@@ -617,6 +618,72 @@ class ExperimentTracking:
         self._run = None
         self._unpatch_stdout_and_stderr()
 
+    def _try_patch_ipython_showtraceback(
+        self, stderr_logger: experiment_logging.ExperimentLogger
+    ) -> tuple[Any, Optional[Callable[..., None]]]:
+        try:
+            from IPython import get_ipython
+
+            ip = get_ipython()
+            if ip:
+                original_showtraceback = ip._showtraceback
+
+                def _patched_showtraceback(
+                    etype: type[BaseException],
+                    evalue: Optional[BaseException],
+                    stb: list[str],
+                ) -> None:
+                    try:
+                        text = ip.InteractiveTB.stb2text(stb)
+                        stderr_logger.write(text)
+                        stderr_logger.flush()
+                    except Exception:
+                        pass
+                    if original_showtraceback is not None:  # for mypy
+                        original_showtraceback(etype, evalue, stb)
+
+                ip._showtraceback = _patched_showtraceback
+
+                return ip, original_showtraceback
+
+        except Exception:
+            pass  # Never let IPython patching break logging setup
+
+        return None, None
+
+    def _try_log_exception(
+        self,
+        exc_type: type[BaseException],
+        exc_value: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        """Append a traceback to the stderr logger before unpatching.
+
+        On exception, `Run.__exit__` ends the run and restores `sys.stderr` before the
+        interpreter prints the traceback for the exception, so the default traceback path
+        never reaches the stderr logger.
+
+        Args:
+            exc_type: The type of the exception.
+            exc_value: The value of the exception.
+            exc_tb: The traceback of the exception.
+        """
+        if self._logging_context is None:
+            return
+        ip = self._logging_context.ipython_instance
+        if ip:
+            try:
+                stb = ip.InteractiveTB.structured_traceback(exc_type, exc_value, exc_tb)
+                text = ip.InteractiveTB.stb2text(stb)
+                self._logging_context.stderr_logger.write(text)
+                self._logging_context.stderr_logger.flush()
+                return
+            except Exception:
+                pass
+        err_msg = "".join(traceback.format_exception(exc_type, exc_value, exc_tb))
+        self._logging_context.stderr_logger.write(err_msg)
+        self._logging_context.stderr_logger.flush()
+
     def _patch_stdout_and_stderr(self) -> None:
         if not self._live_logging_enabled:
             return
@@ -626,6 +693,8 @@ class ExperimentTracking:
 
         assert self._experiment is not None and self._run is not None  # for mypy
         stdout_logger, stderr_logger, stdout_ctx, stderr_ctx = None, None, None, None  # for exception handling
+        ip: Any = None
+        original_showtraceback: Optional[Callable[..., None]] = None
         try:
             exp_id = self._sql_client.get_experiment_id(self._experiment.name)
             run_id = self._sql_client.get_run_id(
@@ -636,15 +705,23 @@ class ExperimentTracking:
             stderr_logger = experiment_logging.ExperimentLogger(exp_id=exp_id, run_id=run_id, stream="STDERR")
             stdout_ctx = contextlib.redirect_stdout(tee.OutputTee(sys.stdout, stdout_logger))
             stderr_ctx = contextlib.redirect_stderr(tee.OutputTee(sys.stderr, stderr_logger))
+            stdout_ctx.__enter__()
+            stderr_ctx.__enter__()
+
+            ip, original_showtraceback = self._try_patch_ipython_showtraceback(stderr_logger)
+
             self._logging_context = experiment_logging.ExperimentLoggingContext(
                 stdout_logger=stdout_logger,
                 stderr_logger=stderr_logger,
                 stdout_ctx=stdout_ctx,
                 stderr_ctx=stderr_ctx,
+                ipython_instance=ip,
+                original_showtraceback=original_showtraceback,
             )
-            stdout_ctx.__enter__()
-            stderr_ctx.__enter__()
+
         except Exception as e:
+            if ip:
+                ip._showtraceback = original_showtraceback
             if stdout_ctx:
                 stdout_ctx.__exit__(None, None, None)
             if stderr_ctx:
@@ -661,6 +738,8 @@ class ExperimentTracking:
 
     def _unpatch_stdout_and_stderr(self) -> None:
         if self._logging_context:
+            if self._logging_context.ipython_instance:
+                self._logging_context.ipython_instance._showtraceback = self._logging_context.original_showtraceback
             self._logging_context.stdout_ctx.__exit__(None, None, None)
             self._logging_context.stderr_ctx.__exit__(None, None, None)
             self._logging_context.stdout_logger.close()