diff --git a/client/ayon_core/hosts/zbrush/api/lib.py b/client/ayon_core/hosts/zbrush/api/lib.py index 1020575795..95c7ee4376 100644 --- a/client/ayon_core/hosts/zbrush/api/lib.py +++ b/client/ayon_core/hosts/zbrush/api/lib.py @@ -2,6 +2,7 @@ import os import uuid import time import tempfile +import functools import logging from . import CommunicationWrapper @@ -38,10 +39,89 @@ def execute_zscript(zscript, communicator=None): return communicator.execute_zscript(zscript) +def wait_zscript(until=None, + wait: float = 0.1, + ping_wait: float = 2.0, + timeout: float = 15.0) -> int: + """Wait until the condition is met or until zbrush responds again. + + This periodically 'pings' zbrush by submitting a zscript for execution + that will write a temporary ping file. As soon as that file exists it is + assumed that Zbrush has responded. + + If the `until` callable is passed, then during the wait this function will + periodically be called, and when True it's will assume success and stop + waiting. + + Args: + until (callable): If a callable is provided, whenever it returns + True the wait is cancelled and assumed to have finished. + wait (float): The amount of seconds to wait in-between each file + existence check. + ping_wait (float): The amount of seconds between sending a new 'ping' + whether Zbrush is responding already - usually to detect whether + a zscript had finished processing. + timeout (float): The amount of seconds after which the script will be + assumed to have failed and raise an error. + + Returns: + int: -1 if callable `until` returned True. Otherwise returns the amount + of pings that were sent before Zbrush responded. + + """ + # It may occur that a zscript execution gets interrupted and thus a 'ping' + # gets lost. To avoid just long waits until the timeout in case previous + # pings got lost we periodically execute a new check ping zscript to see + # if that finishes rapidly + + ping_filepath = get_tempfile_path().replace("\\", "/") + var_name = str(uuid.uuid4()).replace("-", "_") + create_ping_file_zscript = f""" +[MemCreate, "AYON_{var_name}", 1, 0] +[MemWriteString, "AYON_{var_name}", "1", 0] +[MemSaveToFile, "AYON_{var_name}", "{ping_filepath}", 1] +[MemDelete, "AYON_{var_name}"] + """ + start_time = time.time() + timeout_time = start_time + timeout + last_ping_time = None + num_pings_sent = 0 + while True: + if until is not None and until(): + # We have reached the `until` condition + print("Condition met..") + return -1 + + t = time.time() + if last_ping_time is None or t - last_ping_time > ping_wait: + last_ping_time = t + num_pings_sent += 1 + execute_zscript(create_ping_file_zscript) + + # Check the periodic pings we have sent - check only the last pings + # up to the max amount. + if os.path.exists(ping_filepath): + print(f"Sent {num_pings_sent} pings. " + f"Received answer after {t-start_time} seconds.") + if os.path.isfile(ping_filepath): + os.remove(ping_filepath) + + return num_pings_sent + + if t > timeout_time: + raise RuntimeError( + "Timeout. Zscript took longer than " + f"{timeout}s to run." + ) + + time.sleep(wait) + + def execute_zscript_and_wait(zscript, check_filepath=None, - wait=0.1, - timeout=20): + wait: float = 0.1, + ping_wait: float = 2.0, + timeout: float = 10.0): """Execute ZScript and wait until ZScript finished processing. This actually waits until a particular file exists on disk. If your ZScript @@ -51,9 +131,6 @@ def execute_zscript_and_wait(zscript, finished. If no `check_filepath` is provided a few extra lines of ZScript will be appended to your - Warning: If your script errors in Zbrush and thus does not continue to - write the file then this function will wait around until the timeout. - Raises: RuntimeError: When timeout is reached. @@ -63,22 +140,41 @@ def execute_zscript_and_wait(zscript, wait until the timeout is reached if never found. wait (float): The amount of seconds to wait in-between each file existence check. + ping_wait (float): The amount of seconds between sending a new 'ping' + whether Zbrush is responding already - usually to detect whether + a zscript had finished processing. timeout (float): The amount of seconds after which the script will be assumed to have failed and raise an error. """ - execute_zscript(zscript) + if check_filepath is None: + var_name = str(uuid.uuid4()) + success_check_file = get_tempfile_path().replace("\\", "/") + zscript += f""" +[MemCreate, "AYON_{var_name}", 1, 0] +[MemWriteString, "AYON_{var_name}", "1", 0] +[MemSaveToFile, "AYON_{var_name}", "{success_check_file}", 1] +[MemDelete, "AYON_{var_name}"] + """ + else: + success_check_file = check_filepath - # Wait around until the zscript finished - time_taken = 0 - while not os.path.exists(check_filepath): - time.sleep(wait) - time_taken += wait - if time_taken > timeout: - raise RuntimeError( - "Timeout. Zscript took longer than " - f"{timeout}s to run." - ) + def wait_until(filepath): + if filepath and os.path.exists(filepath): + return True + + fn = functools.partial(wait_until, check_filepath) + + execute_zscript(zscript) + wait_zscript(until=fn, + wait=wait, + ping_wait=ping_wait, + timeout=timeout) + + if not os.path.exists(success_check_file): + raise RuntimeError( + f"Success file does not exist: {success_check_file}" + ) def get_workdir() -> str: @@ -97,16 +193,18 @@ def export_tool(filepath: str, subdivision_level: int = 0): subdivs - e.g. -1 is the highest available subdiv. """ + # TODO: If this overrides a tool's subdiv level it should actually revert + # it to the original level so that subsequent publishes behave the same filepath = filepath.replace("\\", "/") # Only set any subdiv level if subdiv level != 0 set_subdivs_script = "" if subdivision_level != 0: set_subdivs_script = f""" -[VarSet, max_subd, [IGetMax, "Tool:Geometry:SDiv"]] -[If, max_subd > 0, +[VarSet, maxsubd, [IGetMax, "Tool:Geometry:SDiv"]] +[If, #maxsubd > 0, [ISet, "Tool:Geometry:SDiv", {subdivision_level}, 0], - [ISet, "Tool:Geometry:SDiv", [VarGet, max_subd] - {subdivision_level}, 0] + [ISet, "Tool:Geometry:SDiv", #maxsubd - {subdivision_level}, 0] ]""" # Export tool @@ -114,13 +212,14 @@ def export_tool(filepath: str, subdivision_level: int = 0): [IFreeze, {set_subdivs_script} [FileNameSetNext, "{filepath}"] [IKeyPress, 13, [IPress, Tool:Export]] - ]""" # We do not check for the export file's existence because Zbrush might # write the file in chunks, as such the file might exist before the writing # to it has finished - execute_zscript_and_wait(export_tool_zscript, check_filepath=filepath) + execute_zscript_and_wait(export_tool_zscript) + if not os.path.exists(filepath): + raise RuntimeError(f"Export failed. File does not exist: {filepath}") def is_in_edit_mode() -> bool: