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

Implement AI Host functionality with model caching and job management

Browse Source 
- Added ai_host.py to handle AI tasks with FastAPI integration.
- Introduced HostState and ModelCache classes for managing client connections and model loading.
- Implemented job submission and cancellation endpoints in the AI host.
- Enhanced tests for AI host functionality, including client tracking and job cancellation.
- Updated existing tests to accommodate changes in the operator and command structure.
- Refactored runtime commands to use request-based architecture instead of command-line arguments.
- Improved model loading logic to handle CUDA memory management and reuse models efficiently.
This commit is contained in:
马南
2026-08-04 07:35:40 +08:00
parent d694980da7
commit 838af0ee96
16 changed files with 2112 additions and 225 deletions
Download Patch File Download Diff File Expand all files Collapse all files

114
src/snapform/operators.py
View File

@@ -34,7 +34,7 @@ def open_job_directory(result_directory):
@dataclass
class ActiveJob:
kind: str
process: subprocess.Popen
process: subprocess.Popen | None
log_file: object
log_path: Path
status_path: Path
@@ -50,6 +50,7 @@ class ActiveJob:
remove_background: bool = False
temporary_input_path: Path | None = None
cancelled: bool = False
uses_host: bool = False
_active_job = None
@@ -65,7 +66,12 @@ def cancel_active_job():
if _active_job is None:
return
_active_job.cancelled = True
if _active_job.process.poll() is None:
if _active_job.uses_host:
try:
runtime.cancel_host_job(_active_job.status_path.stem)
except RuntimeError:
pass
elif _active_job.process.poll() is None:
_active_job.process.terminate()
@@ -74,7 +80,8 @@ def close_active_job():
if _active_job is None:
return
_active_job.log_file.close()
if _active_job.log_file is not None:
_active_job.log_file.close()
if _active_job.temporary_input_path is not None:
try:
_active_job.temporary_input_path.unlink(missing_ok=True)
@@ -171,17 +178,28 @@ class RunJob(bpy.types.Operator):
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
log_file = paths["log"].open("w", encoding="utf-8")
uses_host = isinstance(command, runtime.HostJobRequest)
log_file = None
log_path = paths["log"]
try:
process = subprocess.Popen(
command,
stdout=log_file,
stderr=subprocess.STDOUT,
cwd=runtime.storage_root(),
creationflags=runtime.creation_flags(),
)
except OSError as error:
log_file.close()
if uses_host:
runtime.submit_host_job(command)
process = None
log_path = runtime.host_log_path()
else:
if self.job_type == "INSTALL_RUNTIME":
runtime.stop_host()
log_file = paths["log"].open("w", encoding="utf-8")
process = subprocess.Popen(
command,
stdout=log_file,
stderr=subprocess.STDOUT,
cwd=runtime.storage_root(),
creationflags=runtime.creation_flags(),
)
except (OSError, RuntimeError) as error:
if log_file is not None:
log_file.close()
self.report({"ERROR"}, f"Unable to start task: {error}")
return {"CANCELLED"}
@@ -189,7 +207,7 @@ class RunJob(bpy.types.Operator):
kind=self.job_type,
process=process,
log_file=log_file,
log_path=paths["log"],
log_path=log_path,
status_path=paths["status"],
result_directory=paths["result"],
export_format=(
@@ -215,6 +233,7 @@ class RunJob(bpy.types.Operator):
and self.snapform_delete_input
else None
),
uses_host=uses_host,
)
settings.status = self._starting_message(settings)
settings.progress = 0.0
@@ -274,7 +293,7 @@ class RunJob(bpy.types.Operator):
if not runtime.model_ready(BEN2_MODEL.key):
raise RuntimeError(f"Download {BEN2_MODEL.label} first")
return (
runtime.remove_background_command(
runtime.remove_background_request(
settings,
input_path,
paths["status"],
@@ -311,7 +330,7 @@ class RunJob(bpy.types.Operator):
"PLY export requires a DA3 Giant or Nested Giant model"
)
return (
runtime.reconstruct_command(
runtime.reconstruct_request(
settings,
export_format,
paths["status"],
@@ -372,6 +391,16 @@ class RunJob(bpy.types.Operator):
return {"CANCELLED"}
self._update_status(context, job)
if job.uses_host:
status = runtime.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"}:
return self._finish(context, job, 1)
return {"RUNNING_MODAL"}
return_code = job.process.poll()
if return_code is None:
return {"RUNNING_MODAL"}
@@ -389,7 +418,8 @@ class RunJob(bpy.types.Operator):
def _finish(self, context, job, return_code):
settings = context.scene.snapform_settings
self._remove_timer(context)
job.log_file.close()
if job.log_file is not None:
job.log_file.close()
if job.cancelled:
settings.status = "Task cancelled"
@@ -403,9 +433,14 @@ class RunJob(bpy.types.Operator):
if return_code != 0 or error:
settings.status = "Task failed"
settings.progress = 0.0
fallback_error = (
"AI Host task failed"
if job.uses_host
else "Task subprocess failed"
)
self.report(
{"ERROR"},
f"{error or 'Task subprocess failed'}; log: {job.log_path}",
f"{error or fallback_error}; log: {job.log_path}",
)
close_active_job()
self._redraw_ui(context, force=True)
@@ -559,3 +594,46 @@ class CancelJob(bpy.types.Operator):
def execute(self, _context):
cancel_active_job()
return {"FINISHED"}
class ControlHost(bpy.types.Operator):
bl_idname = "snapform.control_host"
bl_label = "Control SnapForm AI Host"
action: bpy.props.EnumProperty(
items=(
("START", "Start", "Start or connect to the shared AI Host"),
("CLEAR", "Clear Models", "Release cached AI models"),
("RESTART", "Restart", "Restart the shared AI Host"),
("STOP", "Stop", "Stop the shared AI Host"),
("LOG", "Open Log", "Open the AI Host log"),
),
options={"HIDDEN", "SKIP_SAVE"},
)
def execute(self, _context):
try:
if self.action == "START":
runtime.enable_host()
message = "AI Host is ready"
elif self.action == "CLEAR":
runtime.clear_host_models()
message = "AI model cache is cleared"
elif self.action == "RESTART":
runtime.restart_host()
message = "AI Host is restarting"
elif self.action == "STOP":
runtime.stop_host()
message = "AI Host is stopping"
else:
log_path = runtime.host_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 AI Host log")
return {"FINISHED"}
except (OSError, RuntimeError) as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
self.report({"INFO"}, message)
return {"FINISHED"}