Remove collecting task output models from Models collection during migration

This commit is contained in:
allegroai 2021-05-03 17:57:27 +03:00
parent 179661a0d4
commit 8ea8ad34e6

View File

@ -11,35 +11,13 @@ from .utils import _drop_all_indices_from_collections
def _migrate_task_models(db: Database): def _migrate_task_models(db: Database):
""" """
Collect the task output models from the models collections
Move the execution and output models to new models.input and output lists Move the execution and output models to new models.input and output lists
""" """
tasks: Collection = db["task"] tasks: Collection = db["task"]
models: Collection = db["model"]
models_field = "models" models_field = "models"
now = datetime.utcnow() now = datetime.utcnow()
pipeline = [
{"$match": {"task": {"$exists": True}}},
{"$project": {"name": 1, "task": 1}},
{"$group": {"_id": "$task", "models": {"$push": "$$ROOT"}}},
]
output_models = f"{models_field}.{TaskModelTypes.output}"
for group in models.aggregate(pipeline=pipeline, allowDiskUse=True):
task_id = group.get("_id")
task_models = group.get("models")
if task_id and models:
task_models = [
{"model": m["_id"], "name": m.get("name", m["_id"]), "updated": now}
for m in task_models
]
tasks.update_one(
{"_id": task_id, output_models: {"$in": [None, []]}},
{"$set": {output_models: task_models}},
upsert=False,
)
fields = { fields = {
TaskModelTypes.input: "execution.model", TaskModelTypes.input: "execution.model",
TaskModelTypes.output: "output.model", TaskModelTypes.output: "output.model",