feat: opt-in tracing for periodic log uploads

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))

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):

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):

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"