python - wandb:获取所有工件集合以及这些工件的所有别名的列表

Question

wandb 文档似乎没有解释如何做到这一点 - 但它应该是我想象的一个相当常见的用例？

我像这样实现了大部分(但不是完全)我想要的，但看起来有点笨拙？我希望在 ArtifactCollection 实例上有一个 self.aliases 属性？

更新:它似乎确实是 >= v0.13.10 的 sdk 的一部分: https://github.com/wandb/wandb/blob/v0.13.10/wandb/apis/public.py#L4050-L4053

ENTITY = os.environ.get("WANDB_ENTITY")
API_KEY = os.environ.get("WANDB_API_KEY")

def get_model_artifacts(key=None):
    wandb.login(key=key if key is not None else API_KEY)
    api = wandb.Api(overrides={"entity": ENTITY})
    model_names = [
        i
        for i in api.artifact_type(
            type_name="models", project="train"
        ).collections()
    ]
    for model in model_names:
        artifact = api.artifact("train/" + model.name + ":latest")
        model._attrs.update(artifact._attrs)
        model._attrs["metadata"] = json.loads(model._attrs["metadata"])
        model.aliases = [x["alias"] for x in model._attrs["aliases"]]
    return model_names

我想如果需要的话我可能会考虑编写一个自定义的 graph-ql 查询，或者只是使用这个笨重的方法。

我错过了什么吗？有没有更干净的方法来做到这一点？

这个笨重的方法缺少的一件事是任何旧的别名 - 它只显示最新的模型，然后显示该模型的任何别名(比如说“最新”和“v4”等) - 不确定这会/应该如何被显示，但我希望也能够获得旧的别名(即指向旧版本的工件的别名)。不过，这不太重要。

编辑 - 经过几个小时查看他们的 sdk 代码后，我得到了这个(仍然对它的笨重程度不满意):

ENTITY = os.environ.get("WANDB_ENTITY")
API_KEY = os.environ.get("WANDB_API_KEY")

def get_model_artifacts(key=None):
    wandb.login(key=key if key is not None else API_KEY)
    api = wandb.Api(overrides={"entity": ENTITY})
    model_artifacts = [
        a
        for a in api.artifact_type(
            type_name="models", project="train"
        ).collections()
    ]

    def get_alias_tuple(artifact_version):
        version = None
        aliases = []
        for a in artifact_version._attrs["aliases"]:
            if re.match(r"^v\d+$", a["alias"]):
                version = a["alias"]
            else:
                aliases.append(a["alias"])
        return version, aliases

    for model in model_artifacts:
        # artifact = api.artifact("train/" + model.name + ":latest")
        # model._attrs.update(artifact._attrs)
        # model._attrs["metadata"] = json.loads(model._attrs["metadata"])
        versions = model.versions()
        version_dict = dict(get_alias_tuple(version) for version in versions)
        model.version_dict = version_dict
        model.aliases = [
            x for key, val in model.version_dict.items() for x in [key] + val
        ]
    return model_artifacts

Answer 1

我是安尼鲁德。我是 W&B 的一名工程师，帮助构建工件。您的解决方案非常接近，但是通过在获取工件时使用 latest 别名，我们只会考虑该工件的别名，而不是所有版本。您可以通过循环版本来解决这个问题:

api = wandb.Api()
collections = [
    coll for coll in api.artifact_type(type_name=TYPE, project=PROJECT).collections()
]


aliases = set()
for coll in collections:
    for artifact in coll.versions():
        aliases.update(artifact.aliases)

print(collections)
print(aliases)

目前，有关集合的文档还很少，但我们正在公共(public) API 中对其进行完善，并将很快发布一些相关文档。这些 API 尚未完全准备好发布 - 因此对于粗糙的边缘表示歉意。

如果您对工件有任何其他疑问，请随时直接与我联系。总是很乐意提供帮助。