Skip to content

Add from_entity parameter to _experimental_fork_checkpoint #649

@arcticfly

Description

@arcticfly

Problem

_experimental_fork_checkpoint in ServerlessBackend constructs the source artifact path using the destination model's entity (line ~89):

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

This means you can't fork from a checkpoint in a different W&B entity. For example, forking from 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.

Proposed Fix

Add an optional from_entity parameter to _experimental_fork_checkpoint:

async def _experimental_fork_checkpoint(
    self,
    model: Model,
    from_model: str,
    from_project: str | None = None,
    from_entity: str | None = None,  # NEW
    from_s3_bucket: str | None = None,
    not_after_step: int | None = None,
    verbose: bool = False,
    prefix: str | None = None,
) -> None:

And use it when constructing the artifact path:

from_entity = from_entity or model.entity or api.default_entity

Current Workaround

Download the artifact from the source entity and re-upload it to the destination entity before calling fork. This works but doubles the artifact storage and adds an unnecessary copy step.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions