feat(serverless): support forking checkpoints from a different W&B entity

[AnnaSuSu]

What

ServerlessBackend._experimental_fork_checkpoint builds the source W&B artifact path from the destination model's entity:

from_entity = model.entity or api.default_entity
collection_path = f"{from_entity}/{from_project}/{from_model}"

So a checkpoint can only be forked within the same entity — forking e.g. willow-voice/willow_normal/kl-000-1 into wb-training/willow_normal/my-new-run fails because it looks for the artifact under wb-training. The only workaround is to download from the source entity and re-upload to the destination, doubling storage.

Closes #649.

Fix

Add an optional from_entity parameter to _experimental_fork_checkpoint and resolve the source entity as from_entity → model.entity → api.default_entity. The resolution + path building move into a small pure helper, _wandb_checkpoint_collection_path, which also raises a clear ValueError when no entity can be determined (previously this produced a "None/…" path).

await backend._experimental_fork_checkpoint(
    model,                       # destination, e.g. wb-training/...
    from_model="kl-000-1",
    from_project="willow_normal",
    from_entity="willow-voice",  # NEW: source entity
)

This re-implements the approach from #676 (validated there but voluntarily closed by its author); credit to @poofeth for the original.

Tests

New tests/unit/test_serverless_fork_checkpoint.py — no GPU/backend deps, runs on the base install:

• _wandb_checkpoint_collection_path resolution: explicit from_entity wins → falls back to model entity → falls back to default entity → raises when none.
• An async test drives _experimental_fork_checkpoint through the W&B branch with a fake wandb.Api and asserts the artifact query uses the explicit source entity (("lora", "src-entity/src-project/src-model")).

tests/unit/test_serverless_fork_checkpoint.py ..... [5 passed]
existing serverless unit tests: 12 passed (no regressions)

[AnnaSuSu]

any chance of a look here? adds support for forking checkpoints from a different W&B entity. it's mergeable and ready, glad to adjust if you'd want it done differently.

[AnnaSuSu]

rebased onto current main (88fad83) — the conflict is resolved and the PR is mergeable again.

on the conflict itself: #739 swapped the serverless backend's import wandb / wandb.Api(...) for the ..utils.wandb_sdk lazy-import shim. that landed on the line directly above the one this PR removes (from_entity = model.entity or api.default_entity, which becomes a parameter here), so the two changes were textually adjacent but not semantically at odds. the resolution takes wandb_sdk.api(...) and keeps from_entity as the new parameter — wandb_sdk.api() returns the same wandb.apis.public.Api, so .default_entity and .artifacts() behave as before.

the rebase did require one real change. the test used to stub W&B by injecting a fake module into sys.modules, which no longer works now that the shim does from wandb.apis.public import Api (it raises ModuleNotFoundError: No module named 'wandb.apis'; 'wandb' is not a package). it now patches at the seam #739 established, matching tests/unit/test_serverless_pipeline_trainer_compat.py:

monkeypatch.setattr("art.serverless.backend.wandb_sdk.api", FakeApi)

verified locally on the rebased branch:

• uv run pytest tests/unit/test_serverless_fork_checkpoint.py → 5 passed. reverting the fix reddens the end-to-end test with exactly the bug in Add from_entity parameter to _experimental_fork_checkpoint #649 (it queries dst-entity/src-project/src-model instead of src-entity/...).
• full tests/unit failure set is byte-identical to clean main: 13 pre-existing failures + 20 collection errors, all from torch / transformers / tinker / datasets not being installed in my local venv. 182 passed on main, 187 with this PR (the 5 new ones).
• uv run prek run --all-files: ruff, ruff format and the uv.lock sync check all pass. the only ty diagnostics are pre-existing unresolved-imports under src/art/megatron/** (torch/tilelang absent locally) — none in either file this PR touches.

@bradhilton the workflows on this head are sitting at action_required — fork PRs need a maintainer to click Approve and run, so no checks are displayed. could you approve the run when you get a moment?