diff --git a/metaflow/mflog/save_logs_periodically.py b/metaflow/mflog/save_logs_periodically.py index 359963cca7f..baffdbac0ee 100644 --- a/metaflow/mflog/save_logs_periodically.py +++ b/metaflow/mflog/save_logs_periodically.py @@ -5,6 +5,7 @@ from threading import Thread from metaflow.sidecar import MessageTypes +from metaflow.tracing import traced from metaflow.util import to_unicode from . import update_delay, BASH_SAVE_LOGS_ARGS, TASK_LOG_SOURCE from .mflog import decorate @@ -109,8 +110,34 @@ def _file_size(path): "current_size=%d delta=%d elapsed_seconds=%.3f" % (path, previous, current, current - previous, elapsed), ) - try: - self._call_save_logs() - except: - pass + + upload_start_time = time.time() + returncode = None + exception = None + attrs = { + "total_bytes": str(sum(new_sizes)), + "files_changed": str( + len([s for s, ps in zip(new_sizes, previous_sizes) if s != ps]) + ), + } + with traced("save_logs_periodically.upload", attrs=attrs) as span: + try: + returncode = self._call_save_logs() + except BaseException as e: + # Upload failures are intentionally non-fatal to the + # sidecar, as they were before tracing was added. + exception = e + + attrs["elapsed_seconds"] = "%.3f" % ( + time.time() - upload_start_time + ) + if returncode is not None: + attrs["returncode"] = str(returncode) + attrs["success"] = str(returncode == 0) + if exception is not None: + attrs["exception"] = str(type(exception).__name__) + + if span is not None: + for key, value in attrs.items(): + span.set_attribute(key, value) time.sleep(update_delay(time.time() - start_time)) diff --git a/metaflow/tracing/__init__.py b/metaflow/tracing/__init__.py index e1cd1128e36..f4757c0ebc3 100644 --- a/metaflow/tracing/__init__.py +++ b/metaflow/tracing/__init__.py @@ -43,7 +43,7 @@ def get_trace_id() -> str: def traced(name, attrs=None): if attrs is None: attrs = {} - yield + yield None def tracing(func): diff --git a/metaflow/tracing/tracing_modules.py b/metaflow/tracing/tracing_modules.py index e6d508c1dee..cdb0cbd10f1 100644 --- a/metaflow/tracing/tracing_modules.py +++ b/metaflow/tracing/tracing_modules.py @@ -125,7 +125,7 @@ def traced(name: str, attrs: Optional[Dict] = None): if attrs: for k, v in attrs.items(): span.set_attribute(k, v) - yield + yield span def tracing(func): diff --git a/test/unit/test_save_logs_periodically.py b/test/unit/test_save_logs_periodically.py index d299b2e77f1..07abef46d55 100644 --- a/test/unit/test_save_logs_periodically.py +++ b/test/unit/test_save_logs_periodically.py @@ -1,3 +1,4 @@ +import contextlib import os import sys import textwrap @@ -15,6 +16,56 @@ def _new_sidecar(enable_debug_logs): return sidecar +@contextlib.contextmanager +def _no_trace_context(name=None, attrs=None): + yield None + + +class _RecordingSpan(object): + def __init__(self, attrs): + self.attributes = dict(attrs or {}) + + def set_attribute(self, key, value): + self.attributes[key] = value + + +@contextlib.contextmanager +def _recording_trace(events, spans, name, attrs=None): + assert name == "save_logs_periodically.upload" + events.append("trace_start") + span = _RecordingSpan(attrs) + spans.append(span) + try: + yield span + finally: + events.append("trace_end") + + +def _run_single_upload_update(monkeypatch, mocker, tmp_path, traced_context, upload): + stdout = tmp_path / "stdout" + stderr = tmp_path / "stderr" + stdout.write_bytes(b"out\n") + stderr.write_bytes(b"err\n") + monkeypatch.setenv("MFLOG_STDOUT", str(stdout)) + monkeypatch.setenv("MFLOG_STDERR", str(stderr)) + sidecar = _new_sidecar(False) + sidecar.is_alive = True + mocker.patch( + "metaflow.mflog.save_logs_periodically.time.sleep", + side_effect=lambda _: setattr(sidecar, "is_alive", False), + ) + mocker.patch("metaflow.mflog.save_logs_periodically.time.time", return_value=100) + traced_mock = mocker.patch( + "metaflow.mflog.save_logs_periodically.traced", + side_effect=traced_context, + ) + call_mock = mocker.patch.object(sidecar, "_call_save_logs", side_effect=upload) + + sidecar._update_loop() + + return call_mock, traced_mock + + def _read_uploader_messages(path): if not path.exists(): return [] @@ -271,3 +322,85 @@ def test_call_save_logs_confirms_absence_of_logs_when_child_crashes( assert returncode == -9 assert _read_uploader_messages(uploader_log) == [] process.communicate.assert_called_once_with() + + +def test_tracing_does_not_interfere_when_disabled(monkeypatch, mocker, tmp_path): + """Verify a disabled tracing context does not affect a normal upload.""" + call_mock, traced_mock = _run_single_upload_update( + monkeypatch, mocker, tmp_path, _no_trace_context, lambda: 0 + ) + + call_mock.assert_called_once() + traced_mock.assert_called_once() + + +def test_upload_runs_inside_active_trace_and_records_attributes( + monkeypatch, mocker, tmp_path +): + """Verify the upload is enclosed by tracing and outcome attributes are recorded.""" + events = [] + spans = [] + + def active_traced(name, attrs=None): + return _recording_trace(events, spans, name, attrs) + + def upload(): + events.append("upload_start") + assert events == ["trace_start", "upload_start"] + events.append("upload_end") + return 0 + + call_mock, _ = _run_single_upload_update( + monkeypatch, mocker, tmp_path, active_traced, upload + ) + + call_mock.assert_called_once() + assert events == ["trace_start", "upload_start", "upload_end", "trace_end"] + assert spans[0].attributes == { + "elapsed_seconds": "0.000", + "total_bytes": "8", + "files_changed": "2", + "returncode": "0", + "success": "True", + } + + +def test_failed_upload_records_returncode_and_failure_status( + monkeypatch, mocker, tmp_path +): + """Verify a non-zero upload result is recorded without changing sidecar behavior.""" + spans = [] + + def active_traced(name, attrs=None): + return _recording_trace([], spans, name, attrs) + + call_mock, _ = _run_single_upload_update( + monkeypatch, mocker, tmp_path, active_traced, lambda: 1 + ) + + call_mock.assert_called_once() + assert spans[0].attributes["returncode"] == "1" + assert spans[0].attributes["success"] == "False" + + +def test_upload_exception_closes_trace_and_remains_non_fatal( + monkeypatch, mocker, tmp_path +): + """Verify upload exceptions remain swallowed after the trace is closed.""" + events = [] + spans = [] + + def active_traced(name, attrs=None): + return _recording_trace(events, spans, name, attrs) + + def upload(): + events.append("upload_start") + raise RuntimeError("upload failed") + + call_mock, _ = _run_single_upload_update( + monkeypatch, mocker, tmp_path, active_traced, upload + ) + + call_mock.assert_called_once() + assert events == ["trace_start", "upload_start", "trace_end"] + assert spans[0].attributes["exception"] == "RuntimeError"