icrdr/AnyImage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add unit tests for service functionality and model caching

Browse Source 
- Implement tests for ben2 and da3 model caching to ensure models are reused for matching runtime keys.
- Verify service state management, including client registration, job cancellation, and client expiration.
- Test discovery file creation and removal based on service ownership.
- Ensure frame extraction functionality works without models and produces expected results.
This commit is contained in:
马南
2026-08-13 00:35:40 +08:00
parent 7d0983529e
commit 92d2cea1d1
21 changed files with 935 additions and 887 deletions
Download Patch File Download Diff File Expand all files Collapse all files

132
src/snapform/operators.py
View File

@@ -6,7 +6,7 @@ from pathlib import Path
import bpy
from . import runtime
from . import client
from .models import (
BEN2_MODEL,
get_downloadable_model,
@@ -95,7 +95,7 @@ class ActiveJob:
mesh_create_camera: bool = False
temporary_input_path: Path | None = None
cancelled: bool = False
uses_host: bool = False
uses_service: bool = False
started: bool = False
start_error: str = ""
start_thread: threading.Thread | None = None
@@ -114,9 +114,11 @@ def cancel_active_job():
_active_job.cancelled = True
if not getattr(_active_job, "started", True):
return
if _active_job.uses_host:
if _active_job.uses_service:
try:
runtime.cancel_host_job(_active_job.status_path.stem)
client.cancel_service_job(
client.job_id_from_status(_active_job.status_path)
)
except RuntimeError:
pass
elif _active_job.process is not None and _active_job.process.poll() is None:
@@ -150,7 +152,7 @@ class RunJob(bpy.types.Operator):
job_type: bpy.props.EnumProperty(
items=(
("INSTALL_RUNTIME", "Install Backend", ""),
("INSTALL_ENV", "Install Environment", ""),
("DOWNLOAD_MODEL", "Download Model", ""),
(RECONSTRUCTION_JOB, "Export DA3 Result", ""),
(IMAGE_DEPTH_JOB, "Convert Image to Depth Plane", ""),
@@ -219,10 +221,10 @@ class RunJob(bpy.types.Operator):
settings = context.scene.snapform_settings
try:
command, paths = self._prepare_command(context, settings)
uses_host = isinstance(command, runtime.HostJobRequest)
log_path = runtime.host_log_path() if uses_host else paths["log"]
working_directory = runtime.storage_root()
creation_flags = runtime.creation_flags()
uses_service = isinstance(command, client.JobRequest)
log_path = client.service_log_path() if uses_service else paths["log"]
working_directory = client.storage_root()
creation_flags = client.creation_flags()
except (OSError, RuntimeError, ValueError) as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
@@ -262,7 +264,7 @@ class RunJob(bpy.types.Operator):
and self.snapform_delete_input
else None
),
uses_host=uses_host,
uses_service=uses_service,
)
_active_job = job
settings.status = self._starting_message(settings)
@@ -292,11 +294,11 @@ class RunJob(bpy.types.Operator):
job.started = True
return
try:
if getattr(job, "uses_host", False):
runtime.submit_host_job(command)
if getattr(job, "uses_service", False):
client.submit_service_job(command)
else:
if job.kind == "INSTALL_RUNTIME":
runtime.stop_host()
if job.kind == "INSTALL_ENV":
client.stop_service()
job.log_file = job.log_path.open("w", encoding="utf-8")
job.process = subprocess.Popen(
command,
@@ -307,8 +309,10 @@ class RunJob(bpy.types.Operator):
)
job.started = True
if job.cancelled:
if job.uses_host:
runtime.cancel_host_job(job.status_path.stem)
if job.uses_service:
client.cancel_service_job(
client.job_id_from_status(job.status_path)
)
elif job.process is not None:
job.process.terminate()
except (OSError, RuntimeError) as error:
@@ -318,12 +322,12 @@ class RunJob(bpy.types.Operator):
job.start_error = str(error)
def _prepare_command(self, _context, settings):
paths = runtime.new_job_paths(self.job_type)
paths = client.new_job_paths(self.job_type)
if (
getattr(self, "snapform_delete_input", False)
and self.job_type in IMAGE_EMPTY_JOBS
):
archived_input = runtime.move_input_into_job(
archived_input = client.move_input_into_job(
self.snapform_input_path,
paths["result"],
)
@@ -332,18 +336,18 @@ class RunJob(bpy.types.Operator):
remove_background = bool(
getattr(self, "snapform_remove_background", False)
)
if self.job_type == "INSTALL_RUNTIME":
if self.job_type == "INSTALL_ENV":
self._require_online_access()
return runtime.install_runtime_command(paths["status"]), paths
return client.install_env_command(paths["status"]), paths
if not runtime.runtime_ready():
raise RuntimeError("Install the SnapForm Backend first")
if not client.env_installed():
raise RuntimeError("Install the SnapForm Environment first")
if self.job_type == "DOWNLOAD_MODEL":
self._require_online_access()
model = get_downloadable_model(self.model_key)
return (
runtime.download_model_command(model.key, paths["status"]),
client.download_model_command(model.key, paths["status"]),
paths,
)
@@ -372,10 +376,10 @@ class RunJob(bpy.types.Operator):
raise ValueError(
"BEN2 requires an image, video, or directory of image frames"
)
if not runtime.model_ready(BEN2_MODEL.key):
if not client.model_ready(BEN2_MODEL.key):
raise RuntimeError(f"Download {BEN2_MODEL.label} first")
return (
runtime.remove_background_request(
client.remove_background_request(
settings,
input_path,
paths["status"],
@@ -388,7 +392,7 @@ class RunJob(bpy.types.Operator):
)
if self.job_type == ORIGINAL_FRAMES_JOB:
return (
runtime.extract_frames_request(
client.extract_frames_request(
settings,
input_path,
paths["status"],
@@ -416,12 +420,12 @@ class RunJob(bpy.types.Operator):
use_ray_pose = self.snapform_use_ray_pose
process_res = self.snapform_process_res
model = get_model(model_key)
if not runtime.model_ready(model_key):
if not client.model_ready(model_key):
raise RuntimeError(f"Download {model.label} first")
if (
self.job_type in DA3_IMAGE_JOBS
and remove_background
and not runtime.model_ready(BEN2_MODEL.key)
and not client.model_ready(BEN2_MODEL.key)
):
raise RuntimeError(f"Download {BEN2_MODEL.label} first")
if export_format == "gs_ply" and not model.supports_gs:
@@ -429,7 +433,7 @@ class RunJob(bpy.types.Operator):
"PLY export requires a DA3 Giant or Nested Giant model"
)
return (
runtime.reconstruct_request(
client.reconstruct_request(
settings,
export_format,
paths["status"],
@@ -456,7 +460,7 @@ class RunJob(bpy.types.Operator):
def _task_device(self, settings):
"""Debug panels use the scene device; production tasks use Preferences."""
if self.job_type in PRODUCTION_AI_JOBS:
return runtime.production_device()
return client.production_device()
return settings.device
def _require_online_access(self):
@@ -467,7 +471,7 @@ class RunJob(bpy.types.Operator):
def _starting_message(self, settings):
labels = {
"INSTALL_RUNTIME": "Starting Backend installation",
"INSTALL_ENV": "Starting Environment installation",
"DOWNLOAD_MODEL": "Starting model download",
}
if self.job_type == RECONSTRUCTION_JOB:
@@ -516,14 +520,14 @@ class RunJob(bpy.types.Operator):
return {"CANCELLED"}
if not job.started:
return {"RUNNING_MODAL"}
if job.uses_host:
status = runtime.read_status(job.status_path) or {}
if job.uses_service:
status = client.read_status(job.status_path) or {}
state = status.get("state")
if state == "succeeded":
return self._finish(context, job, 0)
if state in {"failed", "cancelled"}:
return self._finish(context, job, 1)
if runtime.host_status().get("state") in {"ERROR", "STOPPED"}:
if client.service_status().get("state") in {"ERROR", "STOPPED"}:
return self._finish(context, job, 1)
return {"RUNNING_MODAL"}
return_code = job.process.poll()
@@ -532,7 +536,7 @@ class RunJob(bpy.types.Operator):
return self._finish(context, job, return_code)
def _update_status(self, context, job):
status = runtime.read_status(job.status_path)
status = client.read_status(job.status_path)
if status is None:
return
settings = context.scene.snapform_settings
@@ -593,28 +597,28 @@ class RunJob(bpy.types.Operator):
if job.cancelled:
settings.status = "Task cancelled"
settings.progress = 0.0
if getattr(job, "uses_host", False):
runtime.mark_host_job_complete(job.status_path.stem)
if getattr(job, "uses_service", False):
client.mark_service_job_complete(client.job_id_from_status(job.status_path))
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
status = runtime.read_status(job.status_path) or {}
status = client.read_status(job.status_path) or {}
error = status.get("error")
if return_code != 0 or error:
settings.status = "Task failed"
settings.progress = 0.0
fallback_error = (
"Backend Service task failed"
if job.uses_host
"Service task failed"
if job.uses_service
else "Task subprocess failed"
)
self.report(
{"ERROR"},
f"{error or fallback_error}; log: {job.log_path}",
)
if job.uses_host:
runtime.mark_host_job_complete(job.status_path.stem)
if job.uses_service:
client.mark_service_job_complete(client.job_id_from_status(job.status_path))
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
@@ -644,7 +648,7 @@ class RunJob(bpy.types.Operator):
settings.status = "Depth import failed"
settings.progress = 0.0
self.report({"ERROR"}, str(import_error))
runtime.mark_host_job_complete(job.status_path.stem)
client.mark_service_job_complete(client.job_id_from_status(job.status_path))
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
@@ -659,7 +663,7 @@ class RunJob(bpy.types.Operator):
settings.status = "Surface import failed"
settings.progress = 0.0
self.report({"ERROR"}, str(import_error))
runtime.mark_host_job_complete(job.status_path.stem)
client.mark_service_job_complete(client.job_id_from_status(job.status_path))
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
@@ -674,7 +678,7 @@ class RunJob(bpy.types.Operator):
settings.status = "Depth import failed"
settings.progress = 0.0
self.report({"ERROR"}, str(import_error))
runtime.mark_host_job_complete(job.status_path.stem)
client.mark_service_job_complete(client.job_id_from_status(job.status_path))
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
@@ -689,7 +693,7 @@ class RunJob(bpy.types.Operator):
settings.status = "Image update failed"
settings.progress = 0.0
self.report({"ERROR"}, str(import_error))
runtime.mark_host_job_complete(job.status_path.stem)
client.mark_service_job_complete(client.job_id_from_status(job.status_path))
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
@@ -701,8 +705,8 @@ class RunJob(bpy.types.Operator):
job.status_path.unlink(missing_ok=True)
except OSError:
pass
if getattr(job, "uses_host", False):
runtime.mark_host_job_complete(job.status_path.stem)
if getattr(job, "uses_service", False):
client.mark_service_job_complete(client.job_id_from_status(job.status_path))
close_active_job()
self._redraw_ui(context, force=True)
return {"FINISHED"}
@@ -836,17 +840,17 @@ class CancelJob(bpy.types.Operator):
return {"FINISHED"}
class ControlHost(bpy.types.Operator):
bl_idname = "snapform.control_host"
bl_label = "Control SnapForm Backend Service"
class ControlService(bpy.types.Operator):
bl_idname = "snapform.control_service"
bl_label = "Control SnapForm Service"
action: bpy.props.EnumProperty(
items=(
("START", "Start", "Start or connect to the Backend Service"),
("START", "Start", "Start or connect to the Service"),
("CLEAR", "Clear Models", "Release cached AI models"),
("RESTART", "Restart", "Restart the Backend Service"),
("STOP", "Stop", "Stop the Backend Service"),
("LOG", "Open Log", "Open the Backend Service log"),
("RESTART", "Restart", "Restart the Service"),
("STOP", "Stop", "Stop the Service"),
("LOG", "Open Log", "Open the Service log"),
),
options={"HIDDEN", "SKIP_SAVE"},
)
@@ -854,23 +858,23 @@ class ControlHost(bpy.types.Operator):
def execute(self, _context):
try:
if self.action == "START":
runtime.enable_host()
message = "Backend Service is ready"
client.enable_service()
message = "Service is ready"
elif self.action == "CLEAR":
runtime.clear_host_models()
client.clear_service_models()
message = "AI model cache is cleared"
elif self.action == "RESTART":
runtime.restart_host()
message = "Backend Service is restarting"
client.restart_service()
message = "Service is restarting"
elif self.action == "STOP":
runtime.stop_host()
message = "Backend Service is stopping"
client.stop_service()
message = "Service is stopping"
else:
log_path = runtime.host_log_path()
log_path = client.service_log_path()
target = log_path if log_path.is_file() else log_path.parent
result = bpy.ops.wm.path_open(filepath=str(target))
if "FINISHED" not in result:
raise RuntimeError("Unable to open the Backend Service log")
raise RuntimeError("Unable to open the Service log")
return {"FINISHED"}
except (OSError, RuntimeError) as error:
self.report({"ERROR"}, str(error))