Files
enviPy-bayer/epdb/management/commands/dump_enviformer.py
jebus 8fda2577ee [Feature] Dump/Restore of enviFormer Models (#170)
Dump:
`./manage.py  dump_enviformer d544303c-a1ca-439d-b036-5e3413ce4a48 --output test.tar.gz`

Restore:
`./manage.py load_enviformer test.tar.gz 1062eb09-5ec7-4bdd-a8f2-ae0252eb4b06`

Co-authored-by: Tim Lorsbach <tim@lorsba.ch>
Reviewed-on: enviPath/enviPy#170
2025-10-22 10:39:22 +13:00

60 lines
1.9 KiB
Python

import json
import os
import tarfile
from tempfile import TemporaryDirectory
from django.conf import settings as s
from django.core.management.base import BaseCommand
from django.db import transaction
from epdb.models import EnviFormer
class Command(BaseCommand):
def add_arguments(self, parser):
parser.add_argument(
"model",
type=str,
help="Model UUID of the Model to Dump",
)
parser.add_argument("--output", type=str)
def package_dict_and_folder(self, dict_data, folder_path, output_path):
with TemporaryDirectory() as tmpdir:
dict_filename = os.path.join(tmpdir, "data.json")
with open(dict_filename, "w", encoding="utf-8") as f:
json.dump(dict_data, f, indent=2)
with tarfile.open(output_path, "w:gz") as tar:
tar.add(dict_filename, arcname="data.json")
tar.add(folder_path, arcname=os.path.basename(folder_path))
os.remove(dict_filename)
@transaction.atomic
def handle(self, *args, **options):
output = options["output"]
if os.path.exists(output):
raise ValueError(f"Output file {output} already exists")
model = EnviFormer.objects.get(uuid=options["model"])
data = {
"uuid": str(model.uuid),
"name": model.name,
"description": model.description,
"kv": model.kv,
"data_packages_uuids": [str(p.uuid) for p in model.data_packages.all()],
"eval_packages_uuids": [str(p.uuid) for p in model.data_packages.all()],
"threshold": model.threshold,
"eval_results": model.eval_results,
"multigen_eval": model.multigen_eval,
"model_status": model.model_status,
}
model_folder = os.path.join(s.MODEL_DIR, "enviformer", str(model.uuid))
self.package_dict_and_folder(data, model_folder, output)