diff --git a/swanlab/sdk/internal/core_python/transport/sender.py b/swanlab/sdk/internal/core_python/transport/sender.py index b60590793..bf5463462 100644 --- a/swanlab/sdk/internal/core_python/transport/sender.py +++ b/swanlab/sdk/internal/core_python/transport/sender.py @@ -12,7 +12,7 @@ import mimetypes import threading from collections.abc import Callable, Sequence -from pathlib import Path, PurePosixPath +from pathlib import Path, PurePosixPath, PureWindowsPath from typing import IO, TYPE_CHECKING, Literal, Optional, Union, cast import yaml @@ -20,7 +20,7 @@ from swanlab.exceptions import ApiError from swanlab.proto.swanlab.metric.column.v1.column_pb2 import ColumnClass, ColumnRecord, ColumnType from swanlab.proto.swanlab.record.v1.record_pb2 import Record -from swanlab.proto.swanlab.save.v1.save_pb2 import SaveType +from swanlab.proto.swanlab.save.v1.save_pb2 import SaveRecord, SaveType from swanlab.sdk.internal.core_python.api.save import ( complete_multipart_save, complete_save_files, @@ -261,6 +261,36 @@ def upload_log(self, records: Sequence[Record]) -> None: # ── 文件保存上传 ── + def _resolve_save_source(self, save: SaveRecord) -> Optional[Path]: + """解析 save 记录对应的可读本地文件路径。 + + ``source_path`` 为训练机绝对路径；online/local/offline 本机运行时可读。 + sync 在另一台机器、以不同挂载根读取同一 run 目录时该绝对路径不可读， + 此时回退到当前 run 目录的 ``files`` 子目录按文件名重新定位： + + - 用户保存（CUSTOM）：按记录相对名 ``save.name`` 定位（镜像位置，与 ``create_save_links`` 约定一致）； + - 内部保存（metadata/requirements/conda/config）：由 probe 直接写入 files 目录（真实文件）， + 按 ``source_path`` 的 basename 定位。注意 config 的 ``save.name`` 为 ``"config"`` 而非 ``config.yaml``， + 故此处一律用 basename 而非 name。 + + :param save: 文件保存记录 + :return: 可读的本地文件路径；若原始路径与回退路径均不可读则返回 None + """ + primary = Path(save.source_path) + if primary.is_file(): + return primary + # source_path/name 可能由异构系统写入（例如训练在 Windows，sync 在 POSIX）， + # 其分隔符为反斜杠时，POSIX 的 Path 无法正确切分。用 PureWindowsPath 仅做分隔符解析， + # 再交给本地 Path 复原，确保 basename 取正确文件名、子目录层级被正确还原。 + if save.type == SaveType.SAVE_TYPE_CUSTOM: + fallback = self._ctx.files_dir / Path(*PureWindowsPath(save.name).parts) + else: + fallback = self._ctx.files_dir / PureWindowsPath(save.source_path).name + if fallback.is_file(): + console.debug(f"Save source path not readable, recovered from run directory: {fallback}") + return fallback + return None + def upload_save(self, records: Sequence[Record]) -> None: """ 内部保存（如config、metadata等）和用户保存（文件保存）共用 SaveRecord 结构 @@ -276,27 +306,34 @@ def upload_save(self, records: Sequence[Record]) -> None: # 根据约定的 type 字段区分内部保存和用户保存，内部保存直接上传内容，用户保存走后续文件上传逻辑 if save.type == SaveType.SAVE_TYPE_CUSTOM: save_records.append(record) - elif save.type == SaveType.SAVE_TYPE_METADATA: - with safe.block(message=f"Failed to upload metadata, skipping; file kept at {save.source_path}"): - with open(save.source_path, "r", encoding="utf-8") as f: + continue + # 内部保存（metadata/requirements/conda/config）：source_path 为训练机绝对路径， + # 跨挂载根 sync 时不可读，回退到 run 目录 files 子目录按 basename 重新定位 + source_ref_path = self._resolve_save_source(save) + if source_ref_path is None: + console.warning(f"Save file not found, skipping: {save.source_path}") + continue + if save.type == SaveType.SAVE_TYPE_METADATA: + with safe.block(message=f"Failed to upload metadata, skipping; file kept at {source_ref_path}"): + with open(source_ref_path, "r", encoding="utf-8") as f: content = json.load(f) if isinstance(content, dict): upload_metadata(self._username, self._project, self._experiment_id, content=content) elif save.type == SaveType.SAVE_TYPE_REQUIREMENTS: - with safe.block(message=f"Failed to upload requirements, skipping; file kept at {save.source_path}"): - with open(save.source_path, "r", encoding="utf-8") as f: + with safe.block(message=f"Failed to upload requirements, skipping; file kept at {source_ref_path}"): + with open(source_ref_path, "r", encoding="utf-8") as f: content = f.read() if len(content) > 0: upload_requirements(self._username, self._project, self._experiment_id, content=content) elif save.type == SaveType.SAVE_TYPE_CONDA: - with safe.block(message=f"Failed to upload conda, skipping; file kept at {save.source_path}"): - with open(save.source_path, "r", encoding="utf-8") as f: + with safe.block(message=f"Failed to upload conda, skipping; file kept at {source_ref_path}"): + with open(source_ref_path, "r", encoding="utf-8") as f: content = f.read() if len(content) > 0: upload_conda(self._username, self._project, self._experiment_id, content=content) elif save.type == SaveType.SAVE_TYPE_CONFIG: - with safe.block(message=f"Failed to upload config, skipping; file kept at {save.source_path}"): - with open(save.source_path, "r", encoding="utf-8") as f: + with safe.block(message=f"Failed to upload config, skipping; file kept at {source_ref_path}"): + with open(source_ref_path, "r", encoding="utf-8") as f: content = yaml.safe_load(f) if isinstance(content, dict): upload_config(self._username, self._project, self._experiment_id, content=content) @@ -314,21 +351,24 @@ def upload_save(self, records: Sequence[Record]) -> None: if not record.HasField("save"): continue save = record.save - source = Path(save.source_path) - if not source.is_file(): + source_ref_path = self._resolve_save_source(save) + if source_ref_path is None: console.warning(f"Save file not found, skipping: {save.source_path}") continue size: Optional[int] = None - with safe.block(message=f"Failed to stat save file, skipping: {save.source_path}"): - size = source.stat().st_size + with safe.block(message=f"Failed to stat save file, skipping: {source_ref_path}"): + size = source_ref_path.stat().st_size if size is None: continue if size > config.save_size: console.warning( - f"Save file exceeds size limit ({size} > {config.save_size}), skipping: {save.source_path}" + f"Save file exceeds size limit ({size} > {config.save_size}), skipping: {source_ref_path}" ) continue - pending.append((save.source_path, size, save.name)) + # 用解析后的可读路径（可能是回退到 files 子目录的路径）作为实际读取路径 + # 保持 str() 原生路径形式，与下游 compute_md5 / upload_saves 返回的 source_path 字符串一致， + # 避免 Windows 上 .as_posix() 与 str() 形式不一致导致 source_to_path 映射查不到 + pending.append((str(source_ref_path), size, save.name)) if not pending: console.warning("No valid files to save.") return diff --git a/tests/unit/sdk/internal/core_python/transport/test_sender.py b/tests/unit/sdk/internal/core_python/transport/test_sender.py index 800d6bdd9..bf7fd8824 100644 --- a/tests/unit/sdk/internal/core_python/transport/test_sender.py +++ b/tests/unit/sdk/internal/core_python/transport/test_sender.py @@ -8,7 +8,7 @@ from swanlab.proto.swanlab.metric.column.v1.column_pb2 import ColumnType from swanlab.proto.swanlab.metric.data.v1.data_pb2 import MediaItem, MediaRecord, MediaValue from swanlab.proto.swanlab.record.v1.record_pb2 import Record -from swanlab.proto.swanlab.save.v1.save_pb2 import SaveRecord +from swanlab.proto.swanlab.save.v1.save_pb2 import SaveRecord, SaveType from swanlab.sdk.internal.core_python.context import CoreConfig, CoreContext from swanlab.sdk.internal.core_python.transport.sender import HttpRecordSender from swanlab.sdk.internal.core_python.transport.tracker import UploadTracker @@ -71,8 +71,12 @@ def _make_sender( return HttpRecordSender(ctx=ctx) -def _make_save_record(source: Path, name: str = "checkpoints/model.txt") -> Record: - return Record(save=SaveRecord(name=name, source_path=str(source), target_path=str(source))) +def _make_save_record( + source: Path, + name: str = "checkpoints/model.txt", + save_type: SaveType = SaveType.SAVE_TYPE_CUSTOM, +) -> Record: + return Record(save=SaveRecord(name=name, source_path=str(source), target_path=str(source), type=save_type)) def _make_media_record(filename: str, media_type=ColumnType.COLUMN_TYPE_IMAGE) -> Record: @@ -391,3 +395,166 @@ def test_upload_does_not_advance_records_for_retryable_api_error(tmp_path: Path) pass assert tracker.snapshot().uploaded_records == 0 + + +def _make_internal_save_record(source_path: str, save_type: SaveType, name: str = "") -> Record: + """构造内部保存（metadata/requirements/conda/config）记录，source_path 可指向不存在的训练机绝对路径。""" + return Record(save=SaveRecord(name=name, source_path=source_path, target_path=source_path, type=save_type)) + + +def test_resolve_save_source_prefers_primary_when_readable(tmp_path: Path): + """本机运行时 source_path 可读 → 直接返回 primary，不走回退、不打 debug。""" + source = tmp_path / "model.txt" + source.write_text("weights", encoding="utf-8") + rec = _make_save_record(source, name="checkpoints/model.txt") + sender = _make_sender(tmp_path) + + with patch("swanlab.sdk.internal.core_python.transport.sender.console.debug") as mock_debug: + result = sender._resolve_save_source(rec.save) + + assert result == source + mock_debug.assert_not_called() + + +def test_resolve_save_source_falls_back_to_files_dir_for_custom(tmp_path: Path): + """CUSTOM 跨挂载根：source_path 不可读 → 回退到 files_dir/save.name。""" + # source_path 指向不存在路径，模拟训练机绝对路径在 sync 机不可读 + rec = _make_save_record(Path("/nonexistent/training-host/model.txt"), name="checkpoints/model.txt") + sender = _make_sender(tmp_path) + # files 子目录下有镜像（含子目录层级） + fallback = tmp_path / "files" / "checkpoints" / "model.txt" + fallback.parent.mkdir(parents=True) + fallback.write_text("recovered", encoding="utf-8") + + with patch("swanlab.sdk.internal.core_python.transport.sender.console.debug") as mock_debug: + result = sender._resolve_save_source(rec.save) + + assert result == fallback + mock_debug.assert_called_once() + + +def test_resolve_save_source_falls_back_for_internal_metadata(tmp_path: Path): + """内部 metadata：source_path 不可读 → 按 basename 回退到 files_dir/swanlab-metadata.json。""" + rec = _make_internal_save_record("/nonexistent/host/swanlab-metadata.json", SaveType.SAVE_TYPE_METADATA) + sender = _make_sender(tmp_path) + fallback = tmp_path / "files" / "swanlab-metadata.json" + fallback.parent.mkdir(parents=True) + fallback.write_text("{}", encoding="utf-8") + + result = sender._resolve_save_source(rec.save) + + assert result == fallback + + +def test_resolve_save_source_uses_basename_not_name_for_config(tmp_path: Path): + """config 的 save.name 为 "config"，但文件名是 config.yaml → 必须用 basename 而非 name 回退。""" + rec = _make_internal_save_record("/nonexistent/host/config.yaml", SaveType.SAVE_TYPE_CONFIG, name="config") + sender = _make_sender(tmp_path) + fallback = tmp_path / "files" / "config.yaml" + fallback.parent.mkdir(parents=True) + fallback.write_text("key: value", encoding="utf-8") + # files/config（即 save.name）不应被命中 + assert not (tmp_path / "files" / "config").exists() + + result = sender._resolve_save_source(rec.save) + + assert result == fallback + + +def test_resolve_save_source_returns_none_when_both_unreadable(tmp_path: Path): + rec = _make_save_record(Path("/nonexistent/training-host/model.txt"), name="checkpoints/model.txt") + sender = _make_sender(tmp_path) + # 既无 primary 也无 files 镜像 + + assert sender._resolve_save_source(rec.save) is None + + +def test_upload_save_recovers_internal_metadata_from_run_dir(tmp_path: Path): + """内部 metadata 跨挂载根 sync：source_path 不可读，但 files/swanlab-metadata.json 有内容 → 成功上传。""" + sender = _make_sender(tmp_path) + # files 目录下放置真实 metadata（probe 直接写入） + metadata_file = tmp_path / "files" / "swanlab-metadata.json" + metadata_file.parent.mkdir(parents=True) + metadata_file.write_text('{"hostname": "gpu-01"}', encoding="utf-8") + # source_path 指向训练机绝对路径（sync 机不可读） + rec = _make_internal_save_record("/nonexistent/host/swanlab-metadata.json", SaveType.SAVE_TYPE_METADATA) + + with ( + patch("swanlab.sdk.internal.core_python.transport.sender.upload_metadata") as mock_upload_meta, + patch("swanlab.sdk.internal.core_python.transport.sender.console.debug"), + ): + sender.upload_save([rec]) + + mock_upload_meta.assert_called_once() + assert mock_upload_meta.call_args.kwargs["content"] == {"hostname": "gpu-01"} + + +def test_upload_save_recovers_custom_file_from_run_dir(tmp_path: Path): + """用户保存跨挂载根 sync：source_path 不可读，files/checkpoints/model.txt 有镜像 → 走 save 上传。""" + sender = _make_sender(tmp_path) + files_dir = tmp_path / "files" + files_dir.mkdir() + recovered = files_dir / "checkpoints" / "model.txt" + recovered.parent.mkdir(parents=True) + recovered.write_text("recovered-weights", encoding="utf-8") + # source_path 指向训练机绝对路径（含子目录），sync 机不可读 + rec = _make_save_record(Path("/nonexistent/host/checkpoints/model.txt"), name="checkpoints/model.txt") + + def _fake_upload(session, *, resources, tracker=None): + # 断言读取自回退路径，而非训练机绝对路径 + assert all(r["source_path"] == str(recovered) for r in resources) + return {r["source_path"] for r in resources} + + with ( + patch( + "swanlab.sdk.internal.core_python.transport.sender.prepare_save_files", + return_value={"urls": ["https://s3.test/model"]}, + ) as mock_prepare, + patch("swanlab.sdk.internal.core_python.transport.sender.complete_save_files") as mock_complete, + patch("swanlab.sdk.internal.core_python.transport.sender.client.session.create", return_value=MagicMock()), + patch( + "swanlab.sdk.internal.core_python.transport.sender.upload_saves", side_effect=_fake_upload + ) as mock_upload_saves, + patch("swanlab.sdk.internal.core_python.transport.sender.console.debug"), + ): + sender.upload_save([rec]) + + mock_prepare.assert_called_once() + # 远端 path 仍为记录的相对名 + assert mock_prepare.call_args.kwargs["files"][0]["path"] == "checkpoints/model.txt" + mock_upload_saves.assert_called_once() + mock_complete.assert_called_once_with( + "experiment-id", files=[{"path": "checkpoints/model.txt", "state": "UPLOADED"}] + ) + + +def test_resolve_save_source_handles_windows_separators_on_posix(tmp_path: Path): + """跨平台：训练在 Windows 记录（反斜杠），sync 在 POSIX 读取时回退正确。 + + - 内部保存 source_path='C:\\host\\config.yaml'：POSIX 的 Path 无法切分反斜杠， + 直接取 .name 会得到整串；必须用 PureWindowsPath 取 basename。 + - 用户保存 name='checkpoints\\model.pt'：files 下真实镜像是分层子目录 trees， + files_dir / name 会拼成单个含反斜杠的文件名，须用 PureWindowsPath 还原层级。 + """ + sender = _make_sender(tmp_path) + + # ── 内部保存（config）：basename 解析 ── + config_fallback = tmp_path / "files" / "config.yaml" + config_fallback.parent.mkdir(parents=True) + config_fallback.write_text("key: value", encoding="utf-8") + internal_rec = _make_internal_save_record(r"C:\host\config.yaml", SaveType.SAVE_TYPE_CONFIG, name="config") + assert sender._resolve_save_source(internal_rec.save) == config_fallback + + # ── 用户保存（CUSTOM）：层级还原 ── + custom_fallback = tmp_path / "files" / "checkpoints" / "model.pt" + custom_fallback.parent.mkdir(parents=True) + custom_fallback.write_text("weights", encoding="utf-8") + custom_rec = Record( + save=SaveRecord( + name=r"checkpoints\model.pt", + source_path=r"C:\host\checkpoints\model.pt", + target_path=r"C:\host\checkpoints\model.pt", + type=SaveType.SAVE_TYPE_CUSTOM, + ) + ) + assert sender._resolve_save_source(custom_rec.save) == custom_fallback