Skip to content

Add from_entity parameter to _experimental_fork_checkpoint #649

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.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

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