Automatically migrate existing temp files and auto-resume running jobs when temp directory setting changes in web UI
This commit is contained in:
+90
-4
@@ -1097,21 +1097,107 @@ def get_settings():
|
||||
settings = load_settings()
|
||||
return settings
|
||||
|
||||
def migrate_temp_directory_task(old_dir: str, new_dir: str, active_job_ids: list):
|
||||
import time
|
||||
import shutil
|
||||
|
||||
# 1. Sleep to allow pipeline threads to clean up and exit
|
||||
time.sleep(2.0)
|
||||
|
||||
old_real = os.path.realpath(old_dir) if old_dir else ""
|
||||
new_real = os.path.realpath(new_dir) if new_dir else ""
|
||||
|
||||
if old_real and new_real and old_real != new_real:
|
||||
# 2. Migrate files
|
||||
print(f"Migrating temp files from {old_real} to {new_real}...")
|
||||
if os.path.exists(old_real):
|
||||
for item in os.listdir(old_real):
|
||||
src = os.path.join(old_real, item)
|
||||
dst = os.path.join(new_real, item)
|
||||
try:
|
||||
if os.path.exists(dst):
|
||||
if os.path.isdir(dst):
|
||||
shutil.rmtree(dst)
|
||||
else:
|
||||
os.remove(dst)
|
||||
shutil.move(src, dst)
|
||||
print(f"Moved {src} to {dst}")
|
||||
except Exception as e:
|
||||
print(f"Error moving {src} to {dst}: {e}")
|
||||
|
||||
# 3. Update job temp_dir attributes in database
|
||||
for job_id, job in jobs_db.items():
|
||||
job_temp = getattr(job, "temp_dir", None)
|
||||
if not job_temp or os.path.realpath(job_temp) == old_real:
|
||||
job.temp_dir = new_real
|
||||
save_jobs_db()
|
||||
|
||||
# 4. Resume the paused jobs
|
||||
for job_id in active_job_ids:
|
||||
job = jobs_db.get(job_id)
|
||||
if job:
|
||||
print(f"Auto-resuming job {job_id} at new temp location...")
|
||||
job.status = "queued"
|
||||
job.error = None
|
||||
job.eta = "Queued for resume..."
|
||||
if hasattr(job, "_is_paused"):
|
||||
job._is_paused = False
|
||||
job_queue.put(job_id)
|
||||
save_jobs_db()
|
||||
broadcast_progress(job_id, {
|
||||
"status": "queued",
|
||||
"progress": job.progress,
|
||||
"current_frame": job.current_frame,
|
||||
"total_frames": job.total_frames,
|
||||
"eta": "Queued for resume..."
|
||||
})
|
||||
|
||||
@app.post("/api/settings")
|
||||
def update_settings(req: UpdateSettingsRequest):
|
||||
def update_settings(req: UpdateSettingsRequest, background_tasks: BackgroundTasks):
|
||||
"""Update persistent app settings"""
|
||||
settings = load_settings()
|
||||
# Validate the directory path if provided
|
||||
old_dir = settings.get("default_temp_dir", "").strip()
|
||||
if not old_dir:
|
||||
old_dir = upscaler.TEMP_DIR
|
||||
|
||||
new_dir = req.default_temp_dir.strip()
|
||||
if new_dir:
|
||||
# Ensure the parent directory exists or can be created
|
||||
try:
|
||||
os.makedirs(new_dir, exist_ok=True)
|
||||
except OSError as e:
|
||||
raise HTTPException(status_code=400, detail=f"Invalid directory path: {e}")
|
||||
|
||||
# Save settings
|
||||
settings["default_temp_dir"] = new_dir
|
||||
save_settings(settings)
|
||||
return {"status": "saved", "settings": settings}
|
||||
|
||||
# Identify and pause active jobs
|
||||
active_jobs = []
|
||||
old_real = os.path.realpath(old_dir)
|
||||
new_real = os.path.realpath(new_dir) if new_dir else os.path.realpath(upscaler.TEMP_DIR)
|
||||
|
||||
if old_real != new_real:
|
||||
for job_id, job in list(jobs_db.items()):
|
||||
if job.status in ["queued", "pending", "analyzing", "extracting", "upscaling", "restoring_faces", "interpolating", "assembling"]:
|
||||
active_jobs.append(job_id)
|
||||
try:
|
||||
job.pause()
|
||||
job_queue.remove(job_id)
|
||||
broadcast_progress(job_id, {
|
||||
"status": "paused",
|
||||
"progress": job.progress,
|
||||
"current_frame": job.current_frame,
|
||||
"total_frames": job.total_frames,
|
||||
"eta": "Paused"
|
||||
})
|
||||
except Exception as e:
|
||||
print(f"Error pausing job {job_id} for temp migration: {e}")
|
||||
save_jobs_db()
|
||||
|
||||
# Start background migration and resumption task
|
||||
background_tasks.add_task(migrate_temp_directory_task, old_dir, new_dir, active_jobs)
|
||||
|
||||
return {"status": "saved", "settings": settings, "migrating_jobs": active_jobs}
|
||||
|
||||
|
||||
class CreateDirRequest(BaseModel):
|
||||
|
||||
Reference in New Issue
Block a user