forked from enviPath/enviPy
[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
This commit is contained in:
59
epdb/management/commands/dump_enviformer.py
Normal file
59
epdb/management/commands/dump_enviformer.py
Normal file
@ -0,0 +1,59 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user