wandb 文档似乎没有解释如何做到这一点 - 但它应该是我想象的一个相当常见的用例?
我基本上(但不是完全)实现了我想要的,但它似乎有点笨拙?我本来希望在实例上有一个self.aliases
属性?ArtifactCollection
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