diff --git a/metaflow/mflog/save_logs_periodically.py b/metaflow/mflog/save_logs_periodically.py index 359963cca7f..51d3444d85f 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,31 @@ def _file_size(path): "current_size=%d delta=%d elapsed_seconds=%.3f" % (path, previous, current, current - previous, elapsed), ) + + upload_start_time = time.time() + returncode = None + exception = None try: - self._call_save_logs() - except: + returncode = self._call_save_logs() + except Exception as e: + exception = e + + upload_elapsed = time.time() - upload_start_time + total_bytes = sum(new_sizes) + attrs = { + "elapsed_seconds": "%.3f" % upload_elapsed, + "total_bytes": str(total_bytes), + "files_changed": str(len([s for s, ps in zip(new_sizes, previous_sizes) if s != ps])), + } + 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__) + + with traced("save_logs_periodically.upload", attrs=attrs): + pass + + if exception is not None: pass time.sleep(update_delay(time.time() - start_time)) diff --git a/test/unit/test_save_logs_periodically.py b/test/unit/test_save_logs_periodically.py index d299b2e77f1..7ed0c53d9a3 100644 --- a/test/unit/test_save_logs_periodically.py +++ b/test/unit/test_save_logs_periodically.py @@ -271,3 +271,66 @@ 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 tracing is opt-in and doesn't affect normal upload when tracing is disabled.""" + 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)) + monkeypatch.setenv("DISABLE_TRACING", "1") + 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) + call_mock = mocker.patch( + "metaflow.mflog.save_logs_periodically.subprocess.call", + return_value=0, + ) + + sidecar._update_loop() + + call_mock.assert_called_once() + + +def test_tracing_records_upload_attributes_when_enabled(monkeypatch, mocker, tmp_path): + """Verify tracing captures upload attributes when tracing is enabled.""" + 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)) + monkeypatch.delenv("DISABLE_TRACING", raising=False) + monkeypatch.setenv("OTEL_ENDPOINT", "http://localhost:4318") + 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) + call_mock = mocker.patch( + "metaflow.mflog.save_logs_periodically.subprocess.call", + return_value=0, + ) + traced_mock = mocker.patch("metaflow.mflog.save_logs_periodically.traced") + + sidecar._update_loop() + + call_mock.assert_called_once() + traced_mock.assert_called_once() + call_args = traced_mock.call_args + assert call_args[0][0] == "save_logs_periodically.upload" + attrs = call_args[1]["attrs"] + assert "elapsed_seconds" in attrs + assert "total_bytes" in attrs + assert "files_changed" in attrs + assert "returncode" in attrs + assert "success" in attrs