forked from enviPath/enviPy
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
60 lines
1.9 KiB
Python
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)
|