fix(studio/colab): fix OutStream startup crash and tidy the notebook cards (#7404)
* fix(studio/colab): survive ipykernel OutStream close() during startup Unsloth Studio crashed at server startup on Colab with: Unsloth Studio failed to start: 'OutStream' object has no attribute 'watch_fd_thread' Root cause: - Colab's ipykernel OutStream is created with watchfd=False, so it never gains a watch_fd_thread. The OutStream.close() in the affected ipykernel versions joins that thread unconditionally and raises AttributeError (ipython/ipykernel#867). - _setup_server_disk_logging() replaces sys.stdout/sys.stderr with a tee. That changes the console object identity, so Colab's absl logging handler (which captured the original OutStream and whose close() deliberately skips sys.stdout/sys.stderr) no longer treats it as the live console. - run_server builds uvicorn.Config(...), whose configure_logging runs logging.config.dictConfig -> logging.shutdown, closing every existing handler. The absl handler then calls close() on the orphaned OutStream and the AttributeError propagates out of uvicorn.Config and aborts startup. Fix: - Before installing the tee, harden the displaced console streams' close() so only the ipykernel#867 AttributeError is swallowed; a healthy close() runs unchanged and any other error still propagates. The buggy close() raises before it nulls pub_thread, so the stream stays fully usable. - Give _TeeStream its own close() that flushes the log copy and forwards close() to the wrapped console stream best-effort, so a handler that captured the tee cannot crash startup either. Add regression tests reproducing the exact path (an absl-style handler closing a watchfd=False OutStream stand-in during logging.shutdown) and asserting the tee/console path survives and keeps logging. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Show the Colab login password in the shareable link card * Tighten Colab card comments for PR #7404 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make the Colab tunnel URL clickable and emphasise the password * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Narrow the console close() hardening to the watch_fd_thread AttributeError * Put the Colab password on its own line so selection excludes the label * Keep the Colab password as plain selectable text --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
2d026a1184
commit
c4b777263d
4 changed files with 512 additions and 22 deletions
|
|
@ -218,11 +218,10 @@ def _colab_login_html(username: str, password: str) -> str:
|
|||
Unsloth Studio Login (Colab)
|
||||
</h2>
|
||||
<p style="color: #333333; margin: 0 0 12px 0; font-size: 14px; font-weight: bold;">
|
||||
Log in to Studio with the Cloudflare link above using these credentials. This cell
|
||||
is visible only in your notebook session.
|
||||
Log in as <code>{username}</code> with this password. This cell is visible only in
|
||||
your notebook session.
|
||||
</p>
|
||||
<p style="color: #333333; margin: 0; font-size: 14px; font-family: monospace; font-weight: bold;">
|
||||
Username: <code>{username}</code><br>
|
||||
Password: <code>{password}</code>
|
||||
</p>
|
||||
</div>
|
||||
|
|
@ -441,8 +440,29 @@ def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def _shareable_link_html(cloudflare_url: str) -> str:
|
||||
"""Branded card for the shareable Cloudflare link, styled like the show_link banner."""
|
||||
def _shareable_link_html(
|
||||
cloudflare_url: str,
|
||||
password: "str | None" = None,
|
||||
username: "str | None" = None,
|
||||
) -> str:
|
||||
"""Branded card for the shareable Cloudflare link, styled like the show_link banner.
|
||||
|
||||
*password* renders under the link so the credential sits in the card with the button
|
||||
it unlocks. The username is always the default admin, so it reads inline.
|
||||
"""
|
||||
login_block = ""
|
||||
if password:
|
||||
login_block = f"""
|
||||
<p style="color: #000000; margin: 16px 0 0 0; font-size: 20px; font-weight: 800;">
|
||||
Password
|
||||
</p>
|
||||
<p style="margin: 6px 0 0 0;"><code style="display: inline-block; font-size: 24px;
|
||||
font-weight: 800; text-decoration: underline; background: #f3f3f3;
|
||||
padding: 4px 10px; border-radius: 6px;">{password}</code></p>
|
||||
<p style="color: #666666; margin: 6px 0 0 0; font-size: 12px;">
|
||||
Log in as <code>{username}</code> with this password. Shown only in your
|
||||
notebook session, and never included in the shared link.
|
||||
</p>"""
|
||||
return f"""
|
||||
<div style="display: inline-block; padding: 20px; background: #ffffff; border: 2px solid #000000;
|
||||
border-radius: 12px; margin: 10px 0; font-family: system-ui, -apple-system, sans-serif;">
|
||||
|
|
@ -460,11 +480,12 @@ def _shareable_link_html(cloudflare_url: str) -> str:
|
|||
Open Unsloth Studio
|
||||
</a>
|
||||
<p style="color: #333333; margin: 12px 0 0 0; font-size: 14px; font-weight: bold;">
|
||||
This Cloudflare HTTPS link works from any device — share it with anyone. The Colab view below only works in this tab.
|
||||
This Cloudflare HTTPS link works from any device, so you can share it with anyone.
|
||||
</p>
|
||||
<p style="color: #333333; margin: 16px 0 0 0; font-size: 13px; font-family: monospace; font-weight: bold;">
|
||||
🔗 {cloudflare_url}
|
||||
</p>
|
||||
🔗 <a href="{cloudflare_url}" onclick="var w=window.open(this.href,'_blank');if(!w){{return true;}}return false;"
|
||||
style="color: #000000; text-decoration: underline; cursor: pointer;">{cloudflare_url}</a>
|
||||
</p>{login_block}
|
||||
</div>
|
||||
"""
|
||||
|
||||
|
|
@ -555,28 +576,37 @@ def _show_and_embed(
|
|||
cloudflare_url = cloudflare_url,
|
||||
)
|
||||
|
||||
# Fold the credentials into the link card rather than a second card below it.
|
||||
credentials_shown = False
|
||||
if cloudflare_url:
|
||||
try:
|
||||
from IPython.display import HTML, display
|
||||
display(HTML(_shareable_link_html(cloudflare_url)))
|
||||
|
||||
username, password = colab_login if colab_login else (None, None)
|
||||
display(HTML(_shareable_link_html(cloudflare_url, password, username)))
|
||||
credentials_shown = bool(colab_login)
|
||||
except Exception as e:
|
||||
logger.info(f"Could not render Cloudflare link card ({e}).")
|
||||
|
||||
if colab_login:
|
||||
if colab_login and not credentials_shown:
|
||||
try:
|
||||
_show_colab_login_credentials(*colab_login)
|
||||
except Exception as e:
|
||||
logger.info(f"Could not render Colab login card ({e}).")
|
||||
|
||||
try:
|
||||
show_link(
|
||||
port,
|
||||
_url = url,
|
||||
has_cloudflare_link = bool(cloudflare_url),
|
||||
cloudflare_requested = cloudflare_requested,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.info(f"Could not render Unsloth link card ({e}).")
|
||||
# With a tunnel up the embed below is skipped, so the ready card would only restate
|
||||
# the link card and print a proxy URL that 404s outside this tab.
|
||||
skip_ready_card = _is_colab_runtime() and bool(cloudflare_url)
|
||||
if not skip_ready_card:
|
||||
try:
|
||||
show_link(
|
||||
port,
|
||||
_url = url,
|
||||
has_cloudflare_link = bool(cloudflare_url),
|
||||
cloudflare_requested = cloudflare_requested,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.info(f"Could not render Unsloth link card ({e}).")
|
||||
|
||||
# On Colab with a working tunnel, skip the in-cell proxy embed (often blank).
|
||||
if _is_colab_runtime() and cloudflare_url:
|
||||
|
|
|
|||
|
|
@ -991,10 +991,88 @@ class _TeeStream:
|
|||
except Exception:
|
||||
pass
|
||||
|
||||
def close(self):
|
||||
# We do NOT own the console stream (it is the terminal / Jupyter kernel
|
||||
# stream we wrapped), so closing the tee must never take the server down.
|
||||
# Flush the log copy, then forward close() to the wrapped stream
|
||||
# best-effort: on Colab that stream is an ipykernel OutStream whose
|
||||
# close() can raise (see _harden_console_close / ipython/ipykernel#867).
|
||||
try:
|
||||
self._log_fh.flush()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
self._stream.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self._stream, name)
|
||||
|
||||
|
||||
_WATCH_FD_THREAD_ATTR = "watch_fd_thread"
|
||||
|
||||
|
||||
def _is_missing_watch_fd_thread(exc):
|
||||
"""True only for ipython/ipykernel#867's missing-``watch_fd_thread`` error.
|
||||
|
||||
``AttributeError.name`` exists from Python 3.10; the message carries the
|
||||
attribute name on every version (possibly with a "Did you mean" tail), so
|
||||
check both and let every other AttributeError through.
|
||||
"""
|
||||
if getattr(exc, "name", None) == _WATCH_FD_THREAD_ATTR:
|
||||
return True
|
||||
return _WATCH_FD_THREAD_ATTR in str(exc)
|
||||
|
||||
|
||||
def _harden_console_close(stream):
|
||||
"""Stop a displaced console stream's close() from aborting Studio startup.
|
||||
|
||||
``_setup_server_disk_logging`` replaces ``sys.stdout``/``sys.stderr`` with a
|
||||
tee. That changes the object identity of the console stream, so a third-party
|
||||
logging handler that captured the ORIGINAL stream (notably Colab's ``absl``
|
||||
logging handler, whose ``close()`` skips ``sys.stdout``/``sys.stderr`` but not
|
||||
a stream that is no longer either) treats it as an ordinary stream and calls
|
||||
``close()`` on it during logging teardown -- ``uvicorn.Config()`` ->
|
||||
``logging.config.dictConfig()`` -> ``logging.shutdown()``.
|
||||
|
||||
A Jupyter/Colab ``ipykernel`` ``OutStream`` created with ``watchfd=False``
|
||||
(the Colab default, and every in-process kernel) never gains a
|
||||
``watch_fd_thread``, yet the ``OutStream.close()`` shipped in the affected
|
||||
ipykernel versions joins that thread unconditionally and raises
|
||||
``AttributeError: 'OutStream' object has no attribute 'watch_fd_thread'``
|
||||
(ipython/ipykernel#867). That AttributeError propagates out of
|
||||
``uvicorn.Config(...)`` and aborts startup ("Unsloth Studio failed to start").
|
||||
|
||||
Wrap the stream's ``close()`` in a transparent pass-through that swallows
|
||||
ONLY that specific teardown AttributeError. A healthy close() (a real console
|
||||
stream, or an OutStream with fd-watching on) runs to completion exactly as
|
||||
before and any other error still propagates, so nothing changes off Colab. A
|
||||
stream whose ``close`` cannot be reassigned keeps its original close().
|
||||
"""
|
||||
try:
|
||||
_orig_close = stream.close
|
||||
except Exception:
|
||||
return
|
||||
|
||||
def _safe_close(*args, **kwargs):
|
||||
try:
|
||||
return _orig_close(*args, **kwargs)
|
||||
except AttributeError as exc:
|
||||
if not _is_missing_watch_fd_thread(exc):
|
||||
# A real teardown failure; never hide it.
|
||||
raise
|
||||
# ipython/ipykernel#867: watchfd=False OutStream.close() joins a
|
||||
# thread that was never created. Nothing to clean up; keep going.
|
||||
return None
|
||||
|
||||
try:
|
||||
stream.close = _safe_close
|
||||
except (AttributeError, TypeError):
|
||||
# A stream that forbids setting instance attributes; leave it as-is.
|
||||
pass
|
||||
|
||||
|
||||
def _setup_server_disk_logging():
|
||||
"""Tee stdout/stderr to ~/.unsloth/studio/logs/server/ and aim
|
||||
faulthandler at the same file so hard crashes (access violations /
|
||||
|
|
@ -1037,6 +1115,11 @@ def _setup_server_disk_logging():
|
|||
# the stderr the server already captures.
|
||||
os.environ.setdefault("PYTHONFAULTHANDLER", "1")
|
||||
|
||||
# Replacing the console streams orphans them from third-party "is this the
|
||||
# live console?" checks, so guard their close() first (ipython/ipykernel#867).
|
||||
_harden_console_close(sys.stdout)
|
||||
_harden_console_close(sys.stderr)
|
||||
|
||||
sys.stdout = _TeeStream(sys.stdout, log_fh)
|
||||
sys.stderr = _TeeStream(sys.stderr, log_fh)
|
||||
|
||||
|
|
|
|||
|
|
@ -336,17 +336,71 @@ def test_colab_login_html_includes_credentials():
|
|||
html = colab._colab_login_html("unsloth", "alpha-beta-gamma-delta")
|
||||
assert "unsloth" in html
|
||||
assert "alpha-beta-gamma-delta" in html
|
||||
# The username is fixed, so it reads inline rather than as its own field.
|
||||
assert "Username:" not in html
|
||||
|
||||
|
||||
def test_show_and_embed_renders_cloudflare_before_colab_login(monkeypatch):
|
||||
def test_shareable_link_html_embeds_password_under_the_link():
|
||||
"""The credential belongs in the same card as the button it unlocks."""
|
||||
html = colab._shareable_link_html("https://share.trycloudflare.com", "secret-pass", "unsloth")
|
||||
assert "share.trycloudflare.com" in html
|
||||
assert "secret-pass" in html
|
||||
# Username is stated inline, not as its own labelled field.
|
||||
assert "Username:" not in html
|
||||
assert "unsloth" in html
|
||||
# The password must sit after the link, not above it.
|
||||
assert html.index("share.trycloudflare.com") < html.index("secret-pass")
|
||||
|
||||
|
||||
def test_shareable_link_html_renders_the_url_as_a_link():
|
||||
"""The printed URL is an anchor, using the popup-safe open the button uses."""
|
||||
html = colab._shareable_link_html("https://share.trycloudflare.com")
|
||||
assert '<a href="https://share.trycloudflare.com"' in html
|
||||
assert ">https://share.trycloudflare.com</a>" in html
|
||||
assert html.count("window.open(this.href,'_blank')") == 2
|
||||
|
||||
|
||||
def test_shareable_link_html_emphasises_the_password():
|
||||
"""The password is the one thing to copy, so it is enlarged and underlined."""
|
||||
html = colab._shareable_link_html("https://share.trycloudflare.com", "secret-pass", "unsloth")
|
||||
pw_tag = html[html.index("Password") : html.index("secret-pass")]
|
||||
assert "font-size: 24px" in pw_tag
|
||||
assert "text-decoration: underline" in pw_tag
|
||||
|
||||
|
||||
def test_shareable_link_html_password_has_no_adjacent_whitespace():
|
||||
"""Whitespace beside the password is selected with it on a double click."""
|
||||
html = colab._shareable_link_html("https://share.trycloudflare.com", "secret-pass", "unsloth")
|
||||
before, after = html.split("secret-pass", 1)
|
||||
assert before.endswith(">")
|
||||
assert after.startswith("<")
|
||||
# Label on its own line, so nothing shares the password's text node.
|
||||
assert "Password:" not in html
|
||||
# Plain selectable text: user-select overrides break double click to select.
|
||||
assert "user-select" not in html
|
||||
|
||||
|
||||
def test_shareable_link_html_omits_login_block_without_password():
|
||||
html = colab._shareable_link_html("https://share.trycloudflare.com")
|
||||
assert "Password" not in html
|
||||
|
||||
|
||||
def test_show_and_embed_folds_login_into_the_cloudflare_card(monkeypatch):
|
||||
"""One card, not two: the tunnel card carries the password itself."""
|
||||
displayed: list[str] = []
|
||||
ipython_display = SimpleNamespace(
|
||||
HTML = lambda html: SimpleNamespace(html = html),
|
||||
display = lambda html: displayed.append(html.html),
|
||||
)
|
||||
login_cards: list[tuple] = []
|
||||
|
||||
monkeypatch.setattr(colab, "get_colab_url", lambda port: "https://8888-test.prod.colab.dev/")
|
||||
monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True)
|
||||
monkeypatch.setattr(
|
||||
colab,
|
||||
"_show_colab_login_credentials",
|
||||
lambda *args: login_cards.append(args),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
colab,
|
||||
"show_link",
|
||||
|
|
@ -360,9 +414,74 @@ def test_show_and_embed_renders_cloudflare_before_colab_login(monkeypatch):
|
|||
colab_login = ("unsloth", "secret-pass"),
|
||||
)
|
||||
|
||||
assert len(displayed) == 2
|
||||
assert len(displayed) == 1
|
||||
assert "share.trycloudflare.com" in displayed[0]
|
||||
assert "secret-pass" in displayed[1]
|
||||
assert "secret-pass" in displayed[0]
|
||||
assert login_cards == []
|
||||
|
||||
|
||||
def test_show_and_embed_keeps_separate_login_card_without_tunnel(monkeypatch):
|
||||
"""No tunnel card to fold into, so the standalone login card still renders."""
|
||||
login_cards: list[tuple] = []
|
||||
|
||||
monkeypatch.setattr(colab, "get_colab_url", lambda port: "https://8888-test.prod.colab.dev/")
|
||||
monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True)
|
||||
monkeypatch.setattr(
|
||||
colab,
|
||||
"_show_colab_login_credentials",
|
||||
lambda *args: login_cards.append(args),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
colab,
|
||||
"show_link",
|
||||
lambda port, *, _url = None, has_cloudflare_link = False, cloudflare_requested = False: None,
|
||||
)
|
||||
monkeypatch.setattr(colab, "_embed_kernel_port_iframe", lambda port: True)
|
||||
colab._show_and_embed(8888, colab_login = ("unsloth", "secret-pass"))
|
||||
|
||||
assert login_cards == [("unsloth", "secret-pass")]
|
||||
|
||||
|
||||
def test_show_and_embed_skips_ready_card_when_tunnel_is_up(monkeypatch):
|
||||
"""The ready card only restates the tunnel card and prints a proxy URL that 404s."""
|
||||
calls: list[str] = []
|
||||
|
||||
monkeypatch.setattr(colab, "get_colab_url", lambda port: "https://8888-test.prod.colab.dev/")
|
||||
monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True)
|
||||
monkeypatch.setattr(
|
||||
colab,
|
||||
"show_link",
|
||||
lambda port,
|
||||
*,
|
||||
_url = None,
|
||||
has_cloudflare_link = False,
|
||||
cloudflare_requested = False: calls.append("show_link"),
|
||||
)
|
||||
monkeypatch.setattr(colab, "_embed_kernel_port_iframe", lambda port: True)
|
||||
colab._show_and_embed(8888, cloudflare_url = "https://share.trycloudflare.com")
|
||||
|
||||
assert calls == []
|
||||
|
||||
|
||||
def test_show_and_embed_keeps_ready_card_without_tunnel(monkeypatch):
|
||||
"""Without a tunnel the ready card is the only guidance, so it must stay."""
|
||||
calls: list[str] = []
|
||||
|
||||
monkeypatch.setattr(colab, "get_colab_url", lambda port: "https://8888-test.prod.colab.dev/")
|
||||
monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True)
|
||||
monkeypatch.setattr(
|
||||
colab,
|
||||
"show_link",
|
||||
lambda port,
|
||||
*,
|
||||
_url = None,
|
||||
has_cloudflare_link = False,
|
||||
cloudflare_requested = False: calls.append("show_link"),
|
||||
)
|
||||
monkeypatch.setattr(colab, "_embed_kernel_port_iframe", lambda port: True)
|
||||
colab._show_and_embed(8888)
|
||||
|
||||
assert calls == ["show_link"]
|
||||
|
||||
|
||||
def test_show_and_embed_skips_iframe_on_colab_when_cloudflare_ready(monkeypatch):
|
||||
|
|
|
|||
258
studio/backend/tests/test_server_disk_logging_outstream.py
Normal file
258
studio/backend/tests/test_server_disk_logging_outstream.py
Normal file
|
|
@ -0,0 +1,258 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Regression tests for the Colab "OutStream has no attribute 'watch_fd_thread'"
|
||||
startup crash.
|
||||
|
||||
Field report (Colab): Unsloth Studio dies at server startup with
|
||||
``❌ Unsloth Studio failed to start: 'OutStream' object has no attribute
|
||||
'watch_fd_thread'``.
|
||||
|
||||
Root cause chain:
|
||||
* Colab's ipykernel ``OutStream`` is created with ``watchfd=False``, so it
|
||||
never gains a ``watch_fd_thread``; the ``OutStream.close()`` shipped in the
|
||||
affected ipykernel versions joins that thread unconditionally and raises
|
||||
``AttributeError`` (ipython/ipykernel#867).
|
||||
* ``run._setup_server_disk_logging()`` replaces ``sys.stdout``/``sys.stderr``
|
||||
with a ``_TeeStream``. That changes the console object identity, so Colab's
|
||||
``absl`` logging handler -- which captured the ORIGINAL OutStream and whose
|
||||
``close()`` deliberately skips ``sys.stdout``/``sys.stderr`` -- no longer
|
||||
recognizes it as the live console.
|
||||
* ``run_server`` builds ``uvicorn.Config(...)``, whose ``configure_logging`` ->
|
||||
``logging.config.dictConfig`` -> ``logging.shutdown`` closes every existing
|
||||
handler. The absl handler then calls ``OutStream.close()`` on the orphaned
|
||||
stream, and the AttributeError aborts startup.
|
||||
|
||||
These tests reproduce the mechanism with a stand-in OutStream (Colab-identical
|
||||
constructs are not importable off Colab) and assert the tee/console path used at
|
||||
startup survives it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import logging
|
||||
import sys
|
||||
import weakref
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
|
||||
if _BACKEND_DIR not in sys.path:
|
||||
sys.path.insert(0, _BACKEND_DIR)
|
||||
|
||||
import run as run_mod # noqa: E402
|
||||
|
||||
|
||||
class _ColabOutStream(io.TextIOBase):
|
||||
"""Stand-in for Colab's ipykernel OutStream built with ``watchfd=False``:
|
||||
no ``watch_fd_thread`` and an unguarded ``close()`` that joins it
|
||||
(ipython/ipykernel#867)."""
|
||||
|
||||
def __init__(self, name: str, sink: io.StringIO):
|
||||
self.name = name
|
||||
self._sink = sink
|
||||
|
||||
def write(self, s):
|
||||
return self._sink.write(s)
|
||||
|
||||
def flush(self):
|
||||
pass
|
||||
|
||||
def writable(self):
|
||||
return True
|
||||
|
||||
def isatty(self):
|
||||
return False
|
||||
|
||||
def close(self):
|
||||
# Never set because watchfd=False -> AttributeError, exactly as Colab.
|
||||
self.watch_fd_thread.join()
|
||||
|
||||
def __del__(self):
|
||||
# io.TextIOBase.__del__ would call our buggy close() at GC (the harmless
|
||||
# "Exception ignored" tail seen in Colab); silence it so the test is clean.
|
||||
pass
|
||||
|
||||
|
||||
class _WatchingOutStream(_ColabOutStream):
|
||||
"""OutStream with fd-watching ON: ``watch_fd_thread`` exists, close() is
|
||||
well behaved and must keep working unchanged."""
|
||||
|
||||
def __init__(self, name: str, sink: io.StringIO):
|
||||
super().__init__(name, sink)
|
||||
self.close_ran = False
|
||||
self.watch_fd_thread = type("_T", (), {"join": lambda self: None})()
|
||||
|
||||
def close(self):
|
||||
self.watch_fd_thread.join()
|
||||
self.close_ran = True
|
||||
|
||||
|
||||
class _AbslLikeHandler(logging.StreamHandler):
|
||||
"""Mirror of ``absl.logging.PythonHandler.close()``: close the captured
|
||||
stream unless it is (still) one of the user-managed console streams."""
|
||||
|
||||
def close(self):
|
||||
try:
|
||||
user_managed = (sys.stderr, sys.stdout, sys.__stderr__, sys.__stdout__)
|
||||
if self.stream not in user_managed and (
|
||||
not hasattr(self.stream, "isatty") or not self.stream.isatty()
|
||||
):
|
||||
self.stream.close()
|
||||
except ValueError:
|
||||
pass
|
||||
super().close()
|
||||
|
||||
|
||||
class TestHardenConsoleClose:
|
||||
def test_neutralizes_watchfd_false_close(self):
|
||||
stream = _ColabOutStream("stdout", io.StringIO())
|
||||
with pytest.raises(AttributeError):
|
||||
stream.close() # baseline: the ipykernel #867 bug is real
|
||||
|
||||
stream = _ColabOutStream("stdout", io.StringIO())
|
||||
run_mod._harden_console_close(stream)
|
||||
assert stream.close() is None # swallowed, no crash
|
||||
|
||||
def test_healthy_close_still_runs_fully(self):
|
||||
stream = _WatchingOutStream("stdout", io.StringIO())
|
||||
run_mod._harden_console_close(stream)
|
||||
stream.close()
|
||||
assert stream.close_ran is True
|
||||
|
||||
def test_only_attributeerror_is_swallowed(self):
|
||||
class _Boom:
|
||||
def close(self):
|
||||
raise ValueError("real teardown failure")
|
||||
|
||||
stream = _Boom()
|
||||
run_mod._harden_console_close(stream)
|
||||
with pytest.raises(ValueError):
|
||||
stream.close()
|
||||
|
||||
def test_unrelated_attributeerror_still_propagates(self):
|
||||
# Only #867 is neutralized; a genuine missing attribute during teardown
|
||||
# must still surface instead of looking like a clean close.
|
||||
class _Console:
|
||||
def close(self):
|
||||
return self.not_a_real_attribute
|
||||
|
||||
stream = _Console()
|
||||
run_mod._harden_console_close(stream)
|
||||
with pytest.raises(AttributeError, match = "not_a_real_attribute"):
|
||||
stream.close()
|
||||
|
||||
def test_swallowed_across_attributeerror_message_shapes(self):
|
||||
# Python 3.12 appends a "Did you mean" tail; the match must survive it,
|
||||
# and pre-3.10 AttributeErrors carry no ``name``, only the message.
|
||||
class _Suggesting:
|
||||
def close(self):
|
||||
raise AttributeError(
|
||||
"'OutStream' object has no attribute 'watch_fd_thread'. "
|
||||
"Did you mean: '_watch_pipe_fd'?"
|
||||
)
|
||||
|
||||
stream = _Suggesting()
|
||||
run_mod._harden_console_close(stream)
|
||||
assert stream.close() is None
|
||||
|
||||
def test_unsettable_close_is_left_alone(self):
|
||||
# A stream whose close cannot be reassigned must not raise from hardening.
|
||||
class _Frozen:
|
||||
__slots__ = ()
|
||||
|
||||
def close(self):
|
||||
return "ok"
|
||||
|
||||
stream = _Frozen()
|
||||
run_mod._harden_console_close(stream) # must not raise
|
||||
assert stream.close() == "ok"
|
||||
|
||||
|
||||
class TestTeeStreamClose:
|
||||
def test_tee_close_over_buggy_stream_never_raises(self):
|
||||
console = _ColabOutStream("stdout", io.StringIO())
|
||||
log = io.StringIO()
|
||||
tee = run_mod._TeeStream(console, log)
|
||||
tee.write("before-close")
|
||||
tee.close() # must not raise despite the wrapped stream's broken close
|
||||
assert log.getvalue() == "before-close"
|
||||
|
||||
def test_tee_close_flushes_log(self):
|
||||
class _FlushCounting(io.StringIO):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.flushes = 0
|
||||
|
||||
def flush(self):
|
||||
self.flushes += 1
|
||||
super().flush()
|
||||
|
||||
console, log = io.StringIO(), _FlushCounting()
|
||||
tee = run_mod._TeeStream(console, log)
|
||||
tee.write("x")
|
||||
tee.close()
|
||||
assert log.flushes >= 1
|
||||
|
||||
|
||||
class TestColabStartupRegression:
|
||||
"""End-to-end: the exact trigger -- an absl-style handler closing the
|
||||
orphaned OutStream during the ``logging.shutdown`` that uvicorn's
|
||||
``uvicorn.Config`` -> ``dictConfig`` runs -- must not crash Studio, and the
|
||||
tee must keep logging afterwards.
|
||||
|
||||
``logging.shutdown`` is driven over a LOCAL weakref list (identical code path
|
||||
to ``logging.config._clearExistingHandlers``) so the global logging state and
|
||||
pytest's own capture are untouched.
|
||||
"""
|
||||
|
||||
def _make_console_and_handlers(self, monkeypatch):
|
||||
out_sink, err_sink = io.StringIO(), io.StringIO()
|
||||
out_stream = _ColabOutStream("stdout", out_sink)
|
||||
err_stream = _ColabOutStream("stderr", err_sink)
|
||||
monkeypatch.setattr(sys, "stdout", out_stream)
|
||||
monkeypatch.setattr(sys, "stderr", err_stream)
|
||||
# absl-like handlers capture the ORIGINAL OutStreams (as in Colab).
|
||||
handlers = [_AbslLikeHandler(sys.stdout), _AbslLikeHandler(sys.stderr)]
|
||||
return out_sink, err_sink, out_stream, err_stream, handlers
|
||||
|
||||
def test_baseline_reproduces_crash_without_fix(self, monkeypatch):
|
||||
# Prove the test exercises the real path: swapping the console identity
|
||||
# (what the tee does) makes the absl-like close hit #867.
|
||||
_, _, out_stream, err_stream, handlers = self._make_console_and_handlers(monkeypatch)
|
||||
try:
|
||||
monkeypatch.setattr(sys, "stdout", io.StringIO())
|
||||
monkeypatch.setattr(sys, "stderr", io.StringIO())
|
||||
with pytest.raises(AttributeError, match = "watch_fd_thread"):
|
||||
logging.shutdown([weakref.ref(h) for h in handlers])
|
||||
finally:
|
||||
# Neutralize so a lingering handler can't crash global teardown.
|
||||
run_mod._harden_console_close(out_stream)
|
||||
run_mod._harden_console_close(err_stream)
|
||||
for h in handlers:
|
||||
try:
|
||||
h.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def test_startup_survives_with_harden_and_tee(self, monkeypatch):
|
||||
out_sink, _, out_stream, err_stream, handlers = self._make_console_and_handlers(monkeypatch)
|
||||
|
||||
# Exactly what _setup_server_disk_logging does before serving:
|
||||
run_mod._harden_console_close(sys.stdout)
|
||||
run_mod._harden_console_close(sys.stderr)
|
||||
log_fh = io.StringIO()
|
||||
monkeypatch.setattr(sys, "stdout", run_mod._TeeStream(sys.stdout, log_fh))
|
||||
monkeypatch.setattr(sys, "stderr", run_mod._TeeStream(sys.stderr, log_fh))
|
||||
|
||||
# The close-storm uvicorn triggers via dictConfig -> logging.shutdown,
|
||||
# closing the absl-like handlers over the (now orphaned) OutStreams.
|
||||
logging.shutdown([weakref.ref(h) for h in handlers]) # must NOT raise
|
||||
|
||||
# The tee still tees to both console and disk afterwards.
|
||||
print("post-startup-line")
|
||||
sys.stdout.flush()
|
||||
assert "post-startup-line" in out_sink.getvalue()
|
||||
assert "post-startup-line" in log_fh.getvalue()
|
||||
Loading…
Add table
Add a link
Reference in a new issue