From 018e0ece75c64793ea774dd83ed36a8d3ab92f2e Mon Sep 17 00:00:00 2001 From: Saumya Pailwan Date: Fri, 12 Jun 2026 11:15:32 -0500 Subject: [PATCH 1/5] Make Cancel actually terminate process_fn via SIGTERM --- examples/midi_pitch_shifter/app.py | 63 +++++++++--------- examples/midi_synthesizer/app.py | 51 +++++++-------- examples/pitch_shifter/app.py | 59 ++++++++--------- pyharp/core.py | 100 ++++++++++++++++++++++++++++- 4 files changed, 185 insertions(+), 88 deletions(-) diff --git a/examples/midi_pitch_shifter/app.py b/examples/midi_pitch_shifter/app.py index 44d9087..ece30ee 100644 --- a/examples/midi_pitch_shifter/app.py +++ b/examples/midi_pitch_shifter/app.py @@ -28,37 +28,38 @@ def process_fn( return output_midi_path # Build Gradio endpoint -with gr.Blocks() as demo: - # Define input Gradio Components - input_components = [ - gr.File(type="filepath", - label="Input Midi", - file_types=[".mid", ".midi"]) - .harp_required(True), - gr.Slider( - minimum=-24, - maximum=24, - step=1, - value=7, - label="Pitch Shift (semitones)", - info="Controls the amount of pitch shift in semitones" - ), - ] +if __name__ == "__main__": + with gr.Blocks() as demo: + # Define input Gradio Components + input_components = [ + gr.File(type="filepath", + label="Input Midi", + file_types=[".mid", ".midi"]) + .harp_required(True), + gr.Slider( + minimum=-24, + maximum=24, + step=1, + value=7, + label="Pitch Shift (semitones)", + info="Controls the amount of pitch shift in semitones" + ), + ] - # Define output Gradio Components - output_components = [ - gr.File(type="filepath", - label="Output Midi", - file_types=[".mid", ".midi"]) - .set_info("The pitch-shifted MIDI."), - ] + # Define output Gradio Components + output_components = [ + gr.File(type="filepath", + label="Output Midi", + file_types=[".mid", ".midi"]) + .set_info("The pitch-shifted MIDI."), + ] - # Build a HARP-compatible endpoint - app = build_endpoint( - model_card=model_card, - input_components=input_components, - output_components=output_components, - process_fn=process_fn, - ) + # Build a HARP-compatible endpoint + app = build_endpoint( + model_card=model_card, + input_components=input_components, + output_components=output_components, + process_fn=process_fn, + ) -demo.queue().launch(share=True, show_error=False, pwa=True) + demo.queue().launch(share=True, show_error=False, pwa=True) diff --git a/examples/midi_synthesizer/app.py b/examples/midi_synthesizer/app.py index b998827..1ca992a 100644 --- a/examples/midi_synthesizer/app.py +++ b/examples/midi_synthesizer/app.py @@ -33,28 +33,29 @@ def process_fn(input_midi_path: str) -> str: return output_audio_path # Build Gradio endpoint -with gr.Blocks() as demo: - # Define input Gradio Components - input_components = [ - gr.File(type="filepath", - label="Input Midi", - file_types=[".mid", ".midi"]) - .harp_required(True), - ] - - # Define output Gradio Components - output_components = [ - gr.Audio(type="filepath", - label="Output Audio") - .set_info("The synthesized audio."), - ] - - # Build a HARP-compatible endpoint - app = build_endpoint( - model_card=model_card, - input_components=input_components, - output_components=output_components, - process_fn=process_fn, - ) - -demo.queue().launch(share=True, show_error=False, pwa=True) +if __name__ == "__main__": + with gr.Blocks() as demo: + # Define input Gradio Components + input_components = [ + gr.File(type="filepath", + label="Input Midi", + file_types=[".mid", ".midi"]) + .harp_required(True), + ] + + # Define output Gradio Components + output_components = [ + gr.Audio(type="filepath", + label="Output Audio") + .set_info("The synthesized audio."), + ] + + # Build a HARP-compatible endpoint + app = build_endpoint( + model_card=model_card, + input_components=input_components, + output_components=output_components, + process_fn=process_fn, + ) + + demo.queue().launch(share=True, show_error=False, pwa=True) diff --git a/examples/pitch_shifter/app.py b/examples/pitch_shifter/app.py index 2a61421..16815e5 100644 --- a/examples/pitch_shifter/app.py +++ b/examples/pitch_shifter/app.py @@ -38,35 +38,36 @@ def process_fn( # Build Gradio endpoint -with gr.Blocks() as demo: - # Define input Gradio Components - input_components = [ - gr.Audio(type="filepath", - label="Input Audio A") - .harp_required(True), - gr.Slider( - minimum=-24, - maximum=24, - step=1, - value=7, - label="Pitch Shift (semitones)", - info="Controls the amount of pitch shift in semitones" - ), - ] +if __name__ == "__main__": + with gr.Blocks() as demo: + # Define input Gradio Components + input_components = [ + gr.Audio(type="filepath", + label="Input Audio A") + .harp_required(True), + gr.Slider( + minimum=-24, + maximum=24, + step=1, + value=7, + label="Pitch Shift (semitones)", + info="Controls the amount of pitch shift in semitones" + ), + ] - # Define output Gradio Components - output_components = [ - gr.Audio(type="filepath", - label="Output Audio") - .set_info("The pitch-shifted audio."), - ] + # Define output Gradio Components + output_components = [ + gr.Audio(type="filepath", + label="Output Audio") + .set_info("The pitch-shifted audio."), + ] - # Build a HARP-compatible endpoint - app = build_endpoint( - model_card=model_card, - input_components=input_components, - output_components=output_components, - process_fn=process_fn, - ) + # Build a HARP-compatible endpoint + app = build_endpoint( + model_card=model_card, + input_components=input_components, + output_components=output_components, + process_fn=process_fn, + ) -demo.queue().launch(share=True, show_error=False, pwa=True) + demo.queue().launch(share=True, show_error=False, pwa=True) diff --git a/pyharp/core.py b/pyharp/core.py index e5950c9..88240d5 100644 --- a/pyharp/core.py +++ b/pyharp/core.py @@ -2,9 +2,15 @@ from dataclasses import dataclass, asdict from typing import List +import multiprocessing as mp +import threading import gradio as gr +# "spawn" avoids deadlocks with CUDA/PyTorch libraries that "fork" can cause on Linux (HuggingFace Spaces) +mp.set_start_method("spawn", force=True) + + __all__ = [ 'ModelCard', 'build_endpoint' @@ -174,8 +180,74 @@ def get_harp_component(gr_cmp: Component) -> HarpComponent: return harp_cmp +def _worker_entry(fn, args, result_q): + try: + result = fn(*args) + result_q.put(("ok", result)) + except Exception: + import traceback + result_q.put(("err", traceback.format_exc())) + +class JobSupervisor: + """ + Runs a callable in a subprocess and allows it to be cancelled or + timed out via SIGTERM, rather than running to completion server-side. + """ + + def __init__(self, timeout_s=300): + self.timeout_s = timeout_s + self._process = None + self._result_q = None + self._lock = threading.Lock() + + def run(self, fn, *args): + self.cancel() # single-flight: kill any previous job first + + with self._lock: + self._result_q = mp.Queue() + self._process = mp.Process( + target=_worker_entry, + args=(fn, args, self._result_q), + daemon=True, + ) + self._process.start() + process, result_q = self._process, self._result_q + + process.join(self.timeout_s) + + with self._lock: + if self._process is process and process.is_alive(): + self._terminate("timeout") + + status, payload = result_q.get() + + with self._lock: + if self._process is process: + self._cleanup() + + if status == "err": + raise RuntimeError(payload) + return payload + + def cancel(self): + with self._lock: + if self._process and self._process.is_alive(): + self._terminate("cancelled") + + def _terminate(self, reason): + self._process.terminate() + self._process.join() + if self._result_q is not None: + self._result_q.put(("err", f"Job {reason}")) + self._cleanup() + raise RuntimeError(f"Job {reason}") + + def _cleanup(self): + self._process = None + self._result_q = None + def build_endpoint(model_card: ModelCard, input_components: list, output_components: list, - process_fn: callable) -> tuple: + process_fn: callable, timeout_s: int = 300) -> tuple: """ Builds a Gradio endpoint compatible with HARP. @@ -197,6 +269,16 @@ def build_endpoint(model_card: ModelCard, input_components: list, output_compone - The function must accept the inputs in the same order as the inputs list. - The function must return the outputs in the same order as the outputs list, with a filepath string pointing to each output file. + - process_fn runs in a separate subprocess, so its arguments and + return values must be picklable (e.g. filepath strings, numbers, + booleans, JSON-serializable data). It cannot rely on gr.Progress + or other objects tied to the Gradio request context. + - If the Cancel button is pressed, or if process_fn runs longer + than timeout_s, the subprocess is killed (SIGTERM) and the job + is aborted. + timeout_s (int): Maximum time in seconds to let process_fn run before + it is forcibly killed. Defaults to 300 (5 minutes). Increase this + for models that need more time to process their inputs. Returns: app (dict): A dictionary containing: @@ -231,10 +313,22 @@ def fetch_model_info(): api_name="controls" ) + # Supervise process_fn in a subprocess so it can be killed on cancel/timeout + supervisor = JobSupervisor(timeout_s=timeout_s) + + def supervised_process(*args): + return supervisor.run(process_fn, *args) + + def cancel_handler(): + try: + supervisor.cancel() + except RuntimeError: + pass # expected -- job was cancelled successfully + # Create a button to begin processing process_button = gr.Button("Process") process_event = process_button.click( - fn=process_fn, + fn=supervised_process, inputs=input_components, outputs=output_components, api_name="process" @@ -243,7 +337,7 @@ def fetch_model_info(): # Create a button to cancel processing cancel_button = gr.Button("Cancel") cancel_button.click( - fn=lambda: None, + fn=cancel_handler, inputs=[], outputs=[], api_name="cancel", From 7e2fcbf4df83f6b991768857460ea7ce3594f7dd Mon Sep 17 00:00:00 2001 From: Saumya Pailwan Date: Fri, 12 Jun 2026 13:39:39 -0500 Subject: [PATCH 2/5] increased default timeout time --- pyharp/core.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyharp/core.py b/pyharp/core.py index 88240d5..779b828 100644 --- a/pyharp/core.py +++ b/pyharp/core.py @@ -194,7 +194,7 @@ class JobSupervisor: timed out via SIGTERM, rather than running to completion server-side. """ - def __init__(self, timeout_s=300): + def __init__(self, timeout_s=900): self.timeout_s = timeout_s self._process = None self._result_q = None @@ -247,7 +247,7 @@ def _cleanup(self): self._result_q = None def build_endpoint(model_card: ModelCard, input_components: list, output_components: list, - process_fn: callable, timeout_s: int = 300) -> tuple: + process_fn: callable, timeout_s: int = 900) -> tuple: """ Builds a Gradio endpoint compatible with HARP. @@ -277,7 +277,7 @@ def build_endpoint(model_card: ModelCard, input_components: list, output_compone than timeout_s, the subprocess is killed (SIGTERM) and the job is aborted. timeout_s (int): Maximum time in seconds to let process_fn run before - it is forcibly killed. Defaults to 300 (5 minutes). Increase this + it is forcibly killed. Defaults to 900 (15 minutes). Increase this for models that need more time to process their inputs. Returns: From 3165dd96f9d642a2ac33280df9033fabbbfd53a6 Mon Sep 17 00:00:00 2001 From: Saumya Pailwan Date: Wed, 17 Jun 2026 10:05:21 -0500 Subject: [PATCH 3/5] better error handling --- pyharp/core.py | 26 ++++++++++++++++++++------ setup.py | 4 ++-- 2 files changed, 22 insertions(+), 8 deletions(-) diff --git a/pyharp/core.py b/pyharp/core.py index 779b828..58eebb7 100644 --- a/pyharp/core.py +++ b/pyharp/core.py @@ -181,12 +181,19 @@ def get_harp_component(gr_cmp: Component) -> HarpComponent: return harp_cmp def _worker_entry(fn, args, result_q): + import traceback try: result = fn(*args) result_q.put(("ok", result)) - except Exception: - import traceback - result_q.put(("err", traceback.format_exc())) + except gr.Error as e: + traceback.print_exc() + result_q.put(( + "gr_error", + (e.message, e.duration, e.visible, e.title), + )) + except Exception as e: + tb = traceback.format_exc() + result_q.put(("err", (str(e), tb))) class JobSupervisor: """ @@ -217,7 +224,10 @@ def run(self, fn, *args): with self._lock: if self._process is process and process.is_alive(): - self._terminate("timeout") + try: + self._terminate("timeout") + except RuntimeError: + pass # sentinel pushed to result_q, handled below status, payload = result_q.get() @@ -225,8 +235,12 @@ def run(self, fn, *args): if self._process is process: self._cleanup() + if status == "gr_error": + message, duration, visible, title = payload + raise gr.Error(message, duration=duration, visible=visible, title=title) if status == "err": - raise RuntimeError(payload) + short_msg, tb = payload + raise RuntimeError(f"{short_msg}\n\n{tb}") return payload def cancel(self): @@ -238,7 +252,7 @@ def _terminate(self, reason): self._process.terminate() self._process.join() if self._result_q is not None: - self._result_q.put(("err", f"Job {reason}")) + self._result_q.put(("gr_error", (f"Job {reason}.", 10, True, reason.capitalize()))) self._cleanup() raise RuntimeError(f"Job {reason}") diff --git a/setup.py b/setup.py index 0970f92..4eb70d4 100644 --- a/setup.py +++ b/setup.py @@ -4,12 +4,12 @@ name='pyharp', version='0.3.0', url='https://github.com/TEAMuP-dev/pyharp', - author='Frank Cwitkowitz, Christodoulos Benetatos, Hugo Flores García, Patrick O\'Reilly, Nathan Pruyne, and Aldo Aguilar', + author='Frank Cwitkowitz, Christodoulos Benetatos, Hugo Flores García, Patrick O\'Reilly, Nathan Pruyne, Aldo Aguilar and Saumya Pailwan', author_email='fcwitkow@ur.rochester.edu', description='', packages=find_packages(), install_requires=[ - 'gradio==5.28.0', + 'gradio>=6.17.3,<7', 'descript-audiotools', 'symusic' ] From 3f0d3267de83307ef04507bc727427cf22d77566 Mon Sep 17 00:00:00 2001 From: Frank Cwitkowitz Date: Sun, 30 Aug 2026 08:29:52 -0400 Subject: [PATCH 4/5] Terminating processing on cancel by running process_fn in a reusable worker process, forwarding Gradio progress and messages back, and restructured examples around the __main__ guard this requires. --- README.md | 230 ++++++++++------ examples/midi_pitch_shifter/app.py | 1 + examples/midi_synthesizer/app.py | 1 + examples/pitch_shifter/app.py | 1 + examples/ui_tester/app.py | 218 +++++++-------- pyharp/core.py | 139 +++------- pyharp/worker.py | 412 +++++++++++++++++++++++++++++ setup.py | 2 +- 8 files changed, 707 insertions(+), 297 deletions(-) create mode 100644 pyharp/worker.py diff --git a/README.md b/README.md index 5b6e1bb..e601bdc 100644 --- a/README.md +++ b/README.md @@ -8,6 +8,7 @@ PyHARP is a **companion package** for [HARP](https://github.com/TEAMuP-dev/HARP) * **[PyHARP Apps](#pyharp-apps)** * **[Model Card](#model-card)** * **[Processing Code](#processing-code)** + * **[Worker Processes](#worker-processes)** * **[Pre-Trained Models](#pre-trained-models)** * **[Gradio Endpoint](#gradio-endpoint)** * **[Error Reporting](#error-reporting)** @@ -109,9 +110,50 @@ and returns: Note that by default PyHARP uses the [audiotools](https://github.com/descriptinc/audiotools) library from Descript (installation instructions can be found [here](https://github.com/descriptinc/audiotools#installation)) to load and save audio, but any standard method will work. +### Worker Processes + +`process_fn` runs in a separate worker process, so that pressing Cancel in HARP stops the work rather +than leaving it to run to completion server-side. `build_endpoint` also takes `timeout_s` (default +`900`, _i.e._ 15 minutes), after which a job is stopped the same way. Increase it for models that +legitimately run longer. One job runs at a time and starting a new one stops whatever was running. + +The worker is **reused between requests**, so whatever your `app.py` loads on the way to `process_fn` +is loaded once rather than once per job. Cancelling interrupts the job where it stands and keeps the +worker, models included. Only a job stuck inside a library call that refuses to be interrupted costs a +restart, and a replacement worker starts loading immediately, so it is usually ready again before the +next request. + +Being a separate process, the worker has to import your `app.py` to reach `process_fn`, and running +that file executes everything outside of a `__main__` guard, `launch()` included. The Gradio code +therefore belongs behind one, with `process_fn` defined above it, as every [example](#examples) does: + +```python +def process_fn(input_audio_path: str, pitch_shift_amount: int) -> str: + ... + +if __name__ == "__main__": + with gr.Blocks() as demo: + ... + demo.queue().launch() +``` + +An app without one still works, as PyHARP will suppress the second `launch()` with a warning. +However, in this case the interface is rebuilt in each worker, so the guard is worth adding. + +A few smaller notes: +- Arguments and return values are sent between processes, so they must be picklable. Filepath strings, + numbers, booleans and other plain data are fine. An open file handle or a live model object is not. +- `process_fn` must be reachable by import: defined at the top level of `app.py`, not as a `lambda`, + a closure, or inside the guard. +- `gr.Progress`, `gr.Info` and `gr.Warning` are carried back out of the worker, so they display just + as they would otherwise. The one exception is `gr.Progress().tqdm(...)`, which is not forwarded; + call `progress(...)` directly instead. +- Anything read from the request itself, such as a `gr.Request` parameter, is not available inside + `process_fn`. + ## Pre-Trained Models If you want to build an endpoint that utilizes a pre-trained model, we recommend the following: -- Load the model outside of `process_fn` so that it is only initialized once. Doing it inside would repeat the cost on every request, which usually dominates the runtime. Our [MIDI synthesizer](examples/midi_synthesizer/app.py) example demonstrates this with its soundfont, and the same applies to moving weights onto a GPU ([see below](#self-hosted-endpoints)). +- Load the model outside of `process_fn`, at the top level of `app.py`, so that it is only initialized once. Doing it inside would repeat the cost on every request, which usually dominates the runtime. Our [MIDI synthesizer](examples/midi_synthesizer/app.py) example demonstrates this with its soundfont, and the same applies to moving weights onto a GPU ([see below](#self-hosted-endpoints)). Note that this happens once per worker process rather than once per app, and that a worker is replaced if a job has to be killed to cancel it ([see above](#worker-processes)). - Store model weights within your app repository. Note that these cannot be committed to Git directly (see [Binary Files](#binary-files)). ## Gradio Endpoint @@ -126,41 +168,43 @@ from pyharp import build_endpoint import gradio as gr -# Build the Gradio endpoint -with gr.Blocks() as demo: - # Audio and MIDI components become tracks in HARP; everything else - # becomes a GUI control. Order must match the process_fn signature. - input_components = [ - gr.Audio( - type="filepath", - label="Input Audio" - ).harp_required(True), - gr.Slider( - minimum=-24, - maximum=24, - step=1, - value=7, - label="Pitch Shift (semitones)", - info="Amount to shift the pitch by." - ), - ] - - # Order must match the values returned by process_fn - output_components = [ - gr.Audio( - type="filepath", - label="Output Audio" - ).set_info("The pitch-shifted audio."), - ] - - app = build_endpoint( - model_card=model_card, - input_components=input_components, - output_components=output_components, - process_fn=process_fn, - ) +# The processing worker imports this file, so the app must not be built there +if __name__ == "__main__": + # Build the Gradio endpoint + with gr.Blocks() as demo: + # Audio and MIDI components become tracks in HARP; everything else + # becomes a GUI control. Order must match the process_fn signature. + input_components = [ + gr.Audio( + type="filepath", + label="Input Audio" + ).harp_required(True), + gr.Slider( + minimum=-24, + maximum=24, + step=1, + value=7, + label="Pitch Shift (semitones)", + info="Amount to shift the pitch by." + ), + ] + + # Order must match the values returned by process_fn + output_components = [ + gr.Audio( + type="filepath", + label="Output Audio" + ).set_info("The pitch-shifted audio."), + ] -demo.queue().launch(share=True, show_error=True, pwa=True) + app = build_endpoint( + model_card=model_card, + input_components=input_components, + output_components=output_components, + process_fn=process_fn, + ) + + demo.queue().launch(share=True, show_error=True, pwa=True) ``` A few requirements are easy to miss: @@ -203,30 +247,32 @@ def process_fn(input_midi_path, ...): return output_midi_path -# Build the Gradio endpoint -with gr.Blocks() as demo: - # A gr.File restricted to MIDI extensions becomes a MIDI track in HARP. - # Order must match the process_fn signature. - input_components = [ - gr.File( - type="filepath", - label="Input MIDI", - file_types=[".mid", ".midi"] - ).harp_required(True), - ... - ] - - # Order must match the values returned by process_fn - output_components = [ - gr.File( - type="filepath", - label="Output MIDI", - file_types=[".mid", ".midi"] - ).set_info("The transposed MIDI."), - ... - ] +# The processing worker imports this file, so the app must not be built there +if __name__ == "__main__": + # Build the Gradio endpoint + with gr.Blocks() as demo: + # A gr.File restricted to MIDI extensions becomes a MIDI track in HARP. + # Order must match the process_fn signature. + input_components = [ + gr.File( + type="filepath", + label="Input MIDI", + file_types=[".mid", ".midi"] + ).harp_required(True), + ... + ] - ... + # Order must match the values returned by process_fn + output_components = [ + gr.File( + type="filepath", + label="Output MIDI", + file_types=[".mid", ".midi"] + ).set_info("The transposed MIDI."), + ... + ] + + ... ``` Note that by default PyHARP uses the [symusic](https://github.com/Yikai-Liao/symusic) package to load and save MIDI, but any standard method will work. @@ -273,16 +319,19 @@ def process_fn(...): return ..., output_labels -with gr.Blocks() as demo: +# The processing worker imports this file, so the app must not be built there +if __name__ == "__main__": + # Build the Gradio endpoint + with gr.Blocks() as demo: - ... + ... - output_components = [ - ..., - gr.JSON(label="Output Labels") - ] + output_components = [ + ..., + gr.JSON(label="Output Labels") + ] - ... + ... ``` GUI elements corresponding to these labels will appear on the respective output tracks after processing in HARP. @@ -367,6 +416,10 @@ Rather than patching the model's source, keep the two apart: a **frontend** envi Our [BeatNet Space](https://huggingface.co/spaces/teamup-tech/BeatNet-dual) is a working example of this layout. +Cancellation reaches the backend as well. `process_fn` runs in a [worker process](#worker-processes), +which leads its own process group, so stopping a job stops whatever it started. Invoke the backend +with `subprocess.run` as below and it is torn down with the job rather than left running. + 1. Create a new [Hugging Face Space](https://huggingface.co/new-space). 2. Choose Docker as the SDK along with the blank template. 3. Select the desired hardware option. @@ -471,29 +524,32 @@ git push -u origin main return output_audio_path - with gr.Blocks() as demo: - input_components = [ - gr.Audio(type="filepath", label="Input Audio").harp_required(True), - ] - - output_components = [ - gr.Audio(type="filepath", label="Output Audio"), - ] - - app = build_endpoint( - model_card=model_card, - input_components=input_components, - output_components=output_components, - process_fn=process_fn, + # The processing worker imports this file, so the app must not be built there + if __name__ == "__main__": + # Build the Gradio endpoint + with gr.Blocks() as demo: + input_components = [ + gr.Audio(type="filepath", label="Input Audio").harp_required(True), + ] + + output_components = [ + gr.Audio(type="filepath", label="Output Audio"), + ] + + app = build_endpoint( + model_card=model_card, + input_components=input_components, + output_components=output_components, + process_fn=process_fn, + ) + + # The Space routes traffic to $PORT, and the app must bind to all interfaces + # so that requests can reach it from outside the container + demo.queue().launch( + server_name="0.0.0.0", + server_port=int(os.environ["PORT"]), + show_error=True ) - - # The Space routes traffic to $PORT, and the app must bind to all interfaces - # so that requests can reach it from outside the container - demo.queue().launch( - server_name="0.0.0.0", - server_port=int(os.environ["PORT"]), - show_error=True - ) ``` - `Dockerfile` diff --git a/examples/midi_pitch_shifter/app.py b/examples/midi_pitch_shifter/app.py index cedbdbb..d2712ab 100644 --- a/examples/midi_pitch_shifter/app.py +++ b/examples/midi_pitch_shifter/app.py @@ -42,6 +42,7 @@ def process_fn(input_midi_path: str, pitch_shift_amount: int) -> str: return output_midi_path +# The processing worker imports this file, so the app must not be built there if __name__ == "__main__": # Build the Gradio endpoint with gr.Blocks() as demo: diff --git a/examples/midi_synthesizer/app.py b/examples/midi_synthesizer/app.py index 359023d..5688ed5 100644 --- a/examples/midi_synthesizer/app.py +++ b/examples/midi_synthesizer/app.py @@ -57,6 +57,7 @@ def process_fn(input_midi_path: str) -> str: return output_audio_path +# The processing worker imports this file, so the app must not be built there if __name__ == "__main__": # Build the Gradio endpoint with gr.Blocks() as demo: diff --git a/examples/pitch_shifter/app.py b/examples/pitch_shifter/app.py index 6582d20..b489363 100644 --- a/examples/pitch_shifter/app.py +++ b/examples/pitch_shifter/app.py @@ -49,6 +49,7 @@ def process_fn(input_audio_path: str, pitch_shift_amount: int) -> str: return output_audio_path +# The processing worker imports this file, so the app must not be built there if __name__ == "__main__": # Build the Gradio endpoint with gr.Blocks() as demo: diff --git a/examples/ui_tester/app.py b/examples/ui_tester/app.py index 1bf6e9a..eacafe1 100644 --- a/examples/ui_tester/app.py +++ b/examples/ui_tester/app.py @@ -230,113 +230,115 @@ def process_fn( return output_audio_path, output_midi_path, output_labels, output_file_path -# Build the Gradio endpoint -with gr.Blocks() as demo: - # Audio and MIDI components become tracks in HARP. A gr.File with any - # other extension becomes a GUI file picker instead. Every input here is - # optional, so the app can be run without loading anything. - # Order must match the process_fn signature. - input_components = [ - gr.Audio( - type="filepath", - label="Input Audio" +# The processing worker imports this file, so the app must not be built there +if __name__ == "__main__": + # Build the Gradio endpoint + with gr.Blocks() as demo: + # Audio and MIDI components become tracks in HARP. A gr.File with any + # other extension becomes a GUI file picker instead. Every input here is + # optional, so the app can be run without loading anything. + # Order must match the process_fn signature. + input_components = [ + gr.Audio( + type="filepath", + label="Input Audio" + ) + .harp_required(False) + .set_info("Passed through unchanged. Bundled reference audio is used if empty."), + gr.File( + type="filepath", + label="Input MIDI", + file_types=[".mid", ".midi"] + ) + .harp_required(False) + .set_info("Passed through unchanged. Bundled reference MIDI is used if empty."), + gr.File( + type="filepath", + label="Input File", + file_types=[".txt", ".csv", ".json", ".nam"] + ) + .harp_required(False) + .set_info("A generic file. HARP shows this as a file picker, not a track."), + gr.Slider( + minimum=0, + maximum=60, + step=1, + value=0, + label="Processing Delay (s)", + info="Stalls processing, so the cancel button can be tested." + ), + gr.Slider( + minimum=0.0, + maximum=1.0, + step=0.01, + value=0.5, + label="Gain", + info="A fractional slider (unused)." + ), + gr.Number( + minimum=1, + maximum=16, + step=1, + value=4, + label="Repetitions", + info="A number box (unused)." + ), + gr.Dropdown( + choices=["first", "second", "third"], + value="second", + label="Mode", + info="A dropdown (unused)." + ), + gr.Dropdown( + choices=["reverb", "delay", "chorus"], + value=["reverb", "chorus"], + multiselect=True, + label="Effects", + info="A dropdown allowing more than one selection (unused)." + ), + gr.Checkbox( + value=True, + label="Audio Labels", + info="Emit output labels over the audio track." + ), + gr.Checkbox( + value=True, + label="MIDI Labels", + info="Emit output labels over the MIDI track." + ), + gr.Textbox( + value="Hello World", + label="Text Prompt", + info="A text box. Written to the output file when no input file is given." + ), + ] + + # A gr.JSON output receives the LabelList and is drawn over the tracks. + # Order must match the values returned by process_fn. + output_components = [ + gr.Audio( + type="filepath", + label="Output Audio" + ).set_info("The input audio, unchanged."), + gr.File( + type="filepath", + label="Output MIDI", + file_types=[".mid", ".midi"] + ).set_info("The input MIDI, unchanged."), + gr.JSON( + label="Output Labels" + ).set_info("Labels drawn over the output tracks."), + gr.File( + type="filepath", + label="Output File" + ).set_info("A generic file output."), + ] + + app = build_endpoint( + model_card=model_card, + input_components=input_components, + output_components=output_components, + process_fn=process_fn, ) - .harp_required(False) - .set_info("Passed through unchanged. Bundled reference audio is used if empty."), - gr.File( - type="filepath", - label="Input MIDI", - file_types=[".mid", ".midi"] - ) - .harp_required(False) - .set_info("Passed through unchanged. Bundled reference MIDI is used if empty."), - gr.File( - type="filepath", - label="Input File", - file_types=[".txt", ".csv", ".json", ".nam"] - ) - .harp_required(False) - .set_info("A generic file. HARP shows this as a file picker, not a track."), - gr.Slider( - minimum=0, - maximum=60, - step=1, - value=0, - label="Processing Delay (s)", - info="Stalls processing, so the cancel button can be tested." - ), - gr.Slider( - minimum=0.0, - maximum=1.0, - step=0.01, - value=0.5, - label="Gain", - info="A fractional slider (unused)." - ), - gr.Number( - minimum=1, - maximum=16, - step=1, - value=4, - label="Repetitions", - info="A number box (unused)." - ), - gr.Dropdown( - choices=["first", "second", "third"], - value="second", - label="Mode", - info="A dropdown (unused)." - ), - gr.Dropdown( - choices=["reverb", "delay", "chorus"], - value=["reverb", "chorus"], - multiselect=True, - label="Effects", - info="A dropdown allowing more than one selection (unused)." - ), - gr.Checkbox( - value=True, - label="Audio Labels", - info="Emit output labels over the audio track." - ), - gr.Checkbox( - value=True, - label="MIDI Labels", - info="Emit output labels over the MIDI track." - ), - gr.Textbox( - value="Hello World", - label="Text Prompt", - info="A text box. Written to the output file when no input file is given." - ), - ] - - # A gr.JSON output receives the LabelList and is drawn over the tracks. - # Order must match the values returned by process_fn. - output_components = [ - gr.Audio( - type="filepath", - label="Output Audio" - ).set_info("The input audio, unchanged."), - gr.File( - type="filepath", - label="Output MIDI", - file_types=[".mid", ".midi"] - ).set_info("The input MIDI, unchanged."), - gr.JSON( - label="Output Labels" - ).set_info("Labels drawn over the output tracks."), - gr.File( - type="filepath", - label="Output File" - ).set_info("A generic file output."), - ] - - app = build_endpoint( - model_card=model_card, - input_components=input_components, - output_components=output_components, - process_fn=process_fn, - ) -demo.queue().launch(share=True, show_error=True, pwa=True) + demo.queue().launch(share=True, show_error=True, pwa=True) diff --git a/pyharp/core.py b/pyharp/core.py index 9978964..23326b4 100644 --- a/pyharp/core.py +++ b/pyharp/core.py @@ -2,13 +2,10 @@ from dataclasses import dataclass, asdict from typing import List, Union -import multiprocessing as mp -import threading +import inspect import gradio as gr - -# "spawn" avoids deadlocks with CUDA/PyTorch libraries that "fork" can cause on Linux (HuggingFace Spaces) -mp.set_start_method("spawn", force=True) +from .worker import JobSupervisor __all__ = [ @@ -197,86 +194,6 @@ def get_harp_component(gr_cmp: Component) -> HarpComponent: return harp_cmp -def _worker_entry(fn, args, result_q): - import traceback - try: - result = fn(*args) - result_q.put(("ok", result)) - except gr.Error as e: - traceback.print_exc() - result_q.put(( - "gr_error", - (e.message, e.duration, e.visible, e.title), - )) - except Exception as e: - tb = traceback.format_exc() - result_q.put(("err", (str(e), tb))) - -class JobSupervisor: - """ - Runs a callable in a subprocess and allows it to be cancelled or - timed out via SIGTERM, rather than running to completion server-side. - """ - - def __init__(self, timeout_s=900): - self.timeout_s = timeout_s - self._process = None - self._result_q = None - self._lock = threading.Lock() - - def run(self, fn, *args): - self.cancel() # single-flight: kill any previous job first - - with self._lock: - self._result_q = mp.Queue() - self._process = mp.Process( - target=_worker_entry, - args=(fn, args, self._result_q), - daemon=True, - ) - self._process.start() - process, result_q = self._process, self._result_q - - process.join(self.timeout_s) - - with self._lock: - if self._process is process and process.is_alive(): - try: - self._terminate("timeout") - except RuntimeError: - pass # sentinel pushed to result_q, handled below - - status, payload = result_q.get() - - with self._lock: - if self._process is process: - self._cleanup() - - if status == "gr_error": - message, duration, visible, title = payload - raise gr.Error(message, duration=duration, visible=visible, title=title) - if status == "err": - short_msg, tb = payload - raise RuntimeError(f"{short_msg}\n\n{tb}") - return payload - - def cancel(self): - with self._lock: - if self._process and self._process.is_alive(): - self._terminate("cancelled") - - def _terminate(self, reason): - self._process.terminate() - self._process.join() - if self._result_q is not None: - self._result_q.put(("gr_error", (f"Job {reason}.", 10, True, reason.capitalize()))) - self._cleanup() - raise RuntimeError(f"Job {reason}") - - def _cleanup(self): - self._process = None - self._result_q = None - def build_endpoint(model_card: ModelCard, input_components: list, output_components: list, process_fn: callable, show_controls: bool = False, timeout_s: int = 900) -> tuple: """ @@ -300,13 +217,18 @@ def build_endpoint(model_card: ModelCard, input_components: list, output_compone - The function must accept the inputs in the same order as the inputs list. - The function must return the outputs in the same order as the outputs list, with a filepath string pointing to each output file. - - process_fn runs in a separate subprocess, so its arguments and - return values must be picklable (e.g. filepath strings, numbers, - booleans, JSON-serializable data). It cannot rely on gr.Progress - or other objects tied to the Gradio request context. - - If the Cancel button is pressed, or if process_fn runs longer - than timeout_s, the subprocess is killed (SIGTERM) and the job - is aborted. + - process_fn runs in a worker process, so its arguments and return + values must be picklable (e.g. filepath strings, numbers, booleans, + JSON-serializable data), and it must be reachable by import: defined + at the top level of the app file, not as a lambda, closure, or inside + a __main__ guard. + - gr.Progress, gr.Info, gr.Warning and gr.Error are forwarded out of + the worker and replayed here, so they behave as usual. + - The worker is reused between requests, so anything loaded when the + module is imported is loaded once rather than per job. + - If the Cancel button is pressed, or if process_fn runs longer than + timeout_s, the job is interrupted. A job that will not yield to an + interrupt has its worker replaced instead. show_controls (bool): Whether to show the "View Controls" button and the JSON box holding the control data. - These exist only so that HARP can read the model's interface, and mean @@ -315,9 +237,9 @@ def build_endpoint(model_card: ModelCard, input_components: list, output_compone to someone running the model from the Gradio page directly. - HARP is unaffected either way, since it calls the endpoints rather than clicking the buttons. - timeout_s (int): Maximum time in seconds to let process_fn run before - it is forcibly killed. Defaults to 900 (15 minutes). Increase this - for models that need more time to process their inputs. + timeout_s (int): Maximum time in seconds to let process_fn run before the + job is stopped. Defaults to 900 (15 minutes). Increase this for models + that need more time to process their inputs. Returns: app (dict): A dictionary containing: @@ -352,17 +274,32 @@ def fetch_model_info(): api_name="controls" ) - # Supervise process_fn in a subprocess so it can be killed on cancel/timeout + # Runs process_fn somewhere it can be stopped once it has started supervisor = JobSupervisor(timeout_s=timeout_s) def supervised_process(*args): - return supervisor.run(process_fn, *args) + *inputs, progress = args + + return supervisor.run(process_fn, *inputs, progress=progress) + + # Gradio injects a progress tracker bound to the current request into any handler + # that declares one, and finds it by scanning leading positional parameters. It + # stops at the first *args, so the parameters have to be advertised explicitly. + # The tracker arrives last, after one value per input component. + supervised_process.__signature__ = inspect.Signature( + [ + inspect.Parameter(f"input_{i}", inspect.Parameter.POSITIONAL_OR_KEYWORD) + for i in range(len(input_components)) + ] + + [ + inspect.Parameter( + "progress", inspect.Parameter.POSITIONAL_OR_KEYWORD, default=gr.Progress() + ) + ] + ) def cancel_handler(): - try: - supervisor.cancel() - except RuntimeError: - pass # expected -- job was cancelled successfully + supervisor.cancel() # Create a button to begin processing process_button = gr.Button("Process") diff --git a/pyharp/worker.py b/pyharp/worker.py new file mode 100644 index 0000000..301423d --- /dev/null +++ b/pyharp/worker.py @@ -0,0 +1,412 @@ +""" +Runs process_fn in a worker process, so that a job can be stopped once it has started. + +Gradio has no way to interrupt a running event handler, and a Python thread cannot be +killed, so cancelling one in place is not possible: the work would carry on holding the +GPU whatever the caller did. Running it in a separate process makes it stoppable, and +reusing that process across requests keeps whatever the app loaded on the way to +process_fn from being paid for again on every job. + +Only JobSupervisor is used elsewhere; everything else here supports it. +""" + +import multiprocessing as mp +import os +import queue +import signal +import threading +import time + +import gradio as gr + + +# "spawn" avoids deadlocks that "fork" can cause with CUDA/PyTorch on Linux (Hugging Face +# Spaces). A private context is used rather than mp.set_start_method(force=True), which +# would change the start method for the whole process and override whatever the app or +# any other library had chosen. +_MP = mp.get_context("spawn") + +# Grace period for a result already in flight when the worker exits. The queue is fed by +# a background thread in the worker, so a result put just before exit can arrive slightly +# after the process is gone. +_RESULT_GRACE_S = 10 + +# How often the supervisor wakes to re-check the worker while waiting for messages +_POLL_S = 0.1 + +# How long a cancelled job is given to unwind before the worker is killed outright +_INTERRUPT_GRACE_S = 2 + +# "spawn" re-imports the app module in the worker to reach process_fn, which for an app +# whose Gradio code is not behind an "if __name__" guard means launch() runs there too. +# Suppressing it keeps such an app working, at the cost of building its interface in the +# worker as well. +# +# multiprocessing sets _inheriting for exactly the span of that re-import - it is the flag +# behind its own "if __name__ == '__main__'" guidance - so it identifies a worker without +# any global state having to be set. Read defensively: if it ever goes away, nothing is +# suppressed and an unguarded app is back to needing the guard. +if getattr(mp.current_process(), "_inheriting", False): + + def _suppress_launch(self, *args, **kwargs): + import warnings + + warnings.warn( + "PyHARP suppressed a launch() call in a processing worker. Put the Gradio " + "code behind 'if __name__ == \"__main__\":', with process_fn defined above " + "it, to avoid rebuilding the interface in every worker.", + stacklevel=2, + ) + + return None, None, None + + gr.Blocks.launch = _suppress_launch + +def _redirect_context_calls(result_q, job_id): + """ + Points Gradio's progress and message helpers at the result queue. + + These normally reach the browser through the Gradio request context, which only + exists in the server process. Called from the worker they would find no context + and quietly print to the server log instead, so they are redirected here and + replayed by the supervisor, which does have that context. + + Both are patched at their single definition in gradio.helpers rather than at + gr.Info / gr.Progress, so that an app which imported the names directly is + redirected too. + """ + import gradio.helpers as helpers + + def forward_log(message, title, level="info", duration=10, visible=True): + result_q.put((job_id[0], "log", (str(message), str(title), level, duration, visible))) + + def forward_progress(self, progress, desc=None, total=None, unit="steps", _tqdm=None): + result_q.put((job_id[0], "progress", (progress, desc, total, unit))) + + # gr.Info, gr.Warning and gr.Success all funnel through log_message + helpers.log_message = forward_log + helpers.Progress.__call__ = forward_progress + + +def _worker_loop(jobs_q, results_q, job_done): + """ + Runs jobs one after another until the supervisor stops sending them. + + The worker outlives individual jobs, so a model loaded on the way to process_fn + is loaded once rather than once per request. A cancelled job arrives as + SIGINT, which unwinds Python-level work and leaves the worker - and everything it + has loaded - intact for the next job. + """ + import pickle + import traceback + + try: + # Leads its own process group, so that killing an uninterruptible job takes + # whatever it started with it. A model invoked as a subprocess - the layout the + # dual-environment Docker Spaces use - would otherwise be reparented and keep + # running after its worker is gone. + os.setsid() + except (AttributeError, OSError): + # Not available on this platform; the fallback in _end_process still applies + pass + + current_id = [None] + + _redirect_context_calls(results_q, current_id) + + running = threading.Event() + + def on_interrupt(signum, frame): + # Between jobs there is nothing to unwind, and raising would end the worker + if running.is_set(): + raise KeyboardInterrupt + + signal.signal(signal.SIGINT, on_interrupt) + + while True: + try: + job = jobs_q.get() + except KeyboardInterrupt: + # Arrived with no job to stop, so there is nothing to do but keep waiting + continue + except (EOFError, OSError): + # The queue is gone, so no further job can arrive + return + + job_id, fn, args = job + current_id[0] = job_id + + try: + # Set inside the try, so an interrupt landing here is caught below rather + # than unwinding out of the loop and ending the worker + running.set() + + result = fn(*args) + + # Pickling happens on a feeder thread once queued, where a failure would be + # invisible and the job would look like it never finished. Failing here + # instead reports it as the error it is. + pickle.dumps(result) + + results_q.put((job_id, "ok", result)) + except KeyboardInterrupt: + results_q.put((job_id, "gr_error", ("Job cancelled.", 10, True, "Cancelled"))) + except gr.Error as e: + traceback.print_exc() + results_q.put((job_id, "gr_error", (e.message, e.duration, e.visible, e.title))) + except Exception as e: + results_q.put((job_id, "err", (str(e), traceback.format_exc()))) + finally: + running.clear() + job_done.set() + + +class JobSupervisor: + """ + Runs process_fn in a worker process that can be cancelled or timed out, rather + than running to completion server-side. + + The worker is reused across requests so that whatever the app loads at import + time - model weights above all - is paid for once rather than per job. Cancelling + interrupts the job in place and keeps the worker; only a job stuck in a native + call that will not yield costs a restart, and the replacement is started + immediately so it is usually warm again before the next request. + + One supervisor is shared by every caller of an endpoint, so a Process or Cancel + from one visitor stops whatever job is running. Gradio serialises queued events + by default, which keeps that to a single job at a time; raising an event's + concurrency_limit above 1 would let visitors cancel each other. + """ + + def __init__(self, timeout_s=900): + self.timeout_s = timeout_s + self._worker = None + self._jobs_q = None + self._results_q = None + self._job_done = None + self._job_id = 0 + self._busy = False + self._lock = threading.Lock() + + def run(self, fn, *args, progress=None): + # Single-flight: a new request stops whatever was running, rather than queueing + # behind it, so that Process always starts the job the user just asked for + self.cancel() + + with self._lock: + worker, jobs_q, results_q, job_done = self._ensure_worker() + job_done.clear() + self._job_id += 1 + job_id = self._job_id + self._busy = True + + jobs_q.put((job_id, fn, args)) + + try: + status, payload = self._collect(worker, results_q, job_done, job_id, progress) + finally: + with self._lock: + self._busy = False + + if status == "gr_error": + message, duration, visible, title = payload + raise gr.Error(message, duration=duration, visible=visible, title=title) + if status == "err": + short_msg, tb = payload + raise RuntimeError(f"{short_msg}\n\n{tb}") + if status == "died": + raise gr.Error( + f"Processing stopped unexpectedly (exit code {payload}). This usually " + f"means the model ran out of memory or crashed.", + title="Processing failed", + ) + return payload + + def cancel(self): + with self._lock: + if not self._busy or self._worker is None or not self._worker.is_alive(): + return + + worker, results_q, job_done = self._worker, self._results_q, self._job_done + job_id = self._job_id + + try: + os.kill(worker.pid, signal.SIGINT) + except (ProcessLookupError, PermissionError, OSError): + pass + + # Released before waiting, so the job's own thread can finish collecting + if job_done.wait(timeout=_INTERRUPT_GRACE_S): + # The job unwound and the worker kept everything it had loaded + return + + with self._lock: + if self._worker is not worker or not worker.is_alive(): + return + + # Stuck somewhere that will not accept an interrupt, so nothing short of + # ending the process will stop it + self._discard_worker("cancelled", results_q, job_id) + + # Reload while the user decides what to do next, rather than on their next request + threading.Thread(target=self._warm_up, daemon=True).start() + + def _ensure_worker(self): + """Returns a live worker, starting one if needed. Called with the lock held.""" + if self._worker is not None and self._worker.is_alive(): + return self._worker, self._jobs_q, self._results_q, self._job_done + + self._jobs_q = _MP.Queue() + self._results_q = _MP.Queue() + self._job_done = _MP.Event() + self._worker = _MP.Process( + target=_worker_loop, + args=(self._jobs_q, self._results_q, self._job_done), + daemon=True, + ) + + self._worker.start() + + return self._worker, self._jobs_q, self._results_q, self._job_done + + def _discard_worker(self, reason, results_q, job_id): + """Ends the worker and leaves an outcome behind. Called with the lock held.""" + if self._worker is not None: + self._end_process(self._worker) + + if results_q is not None: + # Frees whichever request is waiting on this worker's queue + results_q.put( + (job_id, "gr_error", (f"Job {reason}.", 10, True, reason.capitalize())) + ) + + self._worker = None + self._jobs_q = None + self._results_q = None + self._job_done = None + + @staticmethod + def _end_process(worker): + """Ends a worker and everything it started.""" + try: + group = os.getpgid(worker.pid) + except (AttributeError, ProcessLookupError, OSError): + group = None + + # Only when the worker leads its own group, or this would signal the server too + if group is not None and group == worker.pid: + try: + os.killpg(group, signal.SIGKILL) + except (AttributeError, ProcessLookupError, PermissionError, OSError): + worker.kill() + else: + worker.kill() + + worker.join(timeout=5) + + def _warm_up(self): + with self._lock: + if self._worker is None: + self._ensure_worker() + + def _collect(self, worker, results_q, job_done, job_id, progress): + """ + Waits for the job's outcome, replaying progress and messages as they arrive. + + This runs in the server's request context, so the calls the worker could not + make itself are made here on its behalf. + """ + import gradio.helpers as helpers + + deadline = time.monotonic() + self.timeout_s + exited_at = None + finished_at = None + timed_out = False + + while True: + try: + message_id, kind, payload = results_q.get(timeout=_POLL_S) + except queue.Empty: + if time.monotonic() >= deadline: + timed_out = True + self._time_out(worker, results_q, job_id) + deadline = time.monotonic() + _RESULT_GRACE_S + continue + + if worker.is_alive(): + exited_at = None + + # The worker has moved on but nothing arrived for this job, which + # is what an unsendable return value looks like + if job_done.is_set(): + if finished_at is None: + finished_at = time.monotonic() + elif time.monotonic() - finished_at >= _RESULT_GRACE_S: + return "err", ( + "The job finished but sent nothing back. Check that " + "everything process_fn returns can be pickled.", + "", + ) + + continue + + # The worker is gone. Its result may still be in flight, since the + # queue is fed by a background thread, so allow a grace period before + # concluding that it stopped without posting one - which is what a + # segfault or an out-of-memory kill looks like. Waiting forever here + # would block this thread, and the queue behind it, for good. + if exited_at is None: + exited_at = time.monotonic() + elif time.monotonic() - exited_at >= _RESULT_GRACE_S: + return "died", worker.exitcode + + continue + + exited_at = None + + # A message from an earlier job, left behind when it was stopped + if message_id != job_id: + continue + + if kind == "progress": + if progress is not None: + value, desc, total, unit = payload + progress(value, desc=desc, total=total, unit=unit) + continue + + if kind == "log": + message, title, level, duration, visible = payload + helpers.log_message( + message, title=title, level=level, duration=duration, visible=visible + ) + continue + + if timed_out and kind == "gr_error": + # The worker reports an interrupt as a cancellation, since it cannot + # tell who asked for one. Here it is known to have been the clock. + _, duration, visible, _ = payload + return kind, ("Job timed out.", duration, visible, "Timed out") + + return kind, payload + + def _time_out(self, worker, results_q, job_id): + """Stops an overrunning job, interrupting first and ending the worker if that fails.""" + if not worker.is_alive(): + return + + try: + os.kill(worker.pid, signal.SIGINT) + except (ProcessLookupError, PermissionError, OSError): + pass + + with self._lock: + job_done = self._job_done + + if job_done is not None and job_done.wait(timeout=_INTERRUPT_GRACE_S): + return + + with self._lock: + if self._worker is worker: + self._discard_worker("timed out", results_q, job_id) + + threading.Thread(target=self._warm_up, daemon=True).start() diff --git a/setup.py b/setup.py index 099512e..8fefa51 100644 --- a/setup.py +++ b/setup.py @@ -13,7 +13,7 @@ # versions discard the error payload on the /gradio_api/call endpoint and # send a bare "data: null", which HARP cannot tell apart from a GPU quota # rejection (see TEAMuP-dev/HARP#349). - 'gradio>=6.17.3,<7', + 'gradio>=6.13.0,<7', 'descript-audiotools', # symusic 0.6.0 broke Synthesizer.render(): it raises "Unable to convert # function return value to a Python type" for its Eigen array return, From dca6a6710e8487d69bcab595f6160dd8068833e1 Mon Sep 17 00:00:00 2001 From: Frank Cwitkowitz Date: Sun, 30 Aug 2026 08:31:09 -0400 Subject: [PATCH 5/5] Added unit testing for worker processes / cancellation. --- README.md | 8 ++ setup.py | 6 +- tests/conftest.py | 71 ++++++++++++ tests/jobs.py | 103 ++++++++++++++++++ tests/test_worker.py | 252 +++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 439 insertions(+), 1 deletion(-) create mode 100644 tests/conftest.py create mode 100644 tests/jobs.py create mode 100644 tests/test_worker.py diff --git a/README.md b/README.md index e601bdc..1ee6999 100644 --- a/README.md +++ b/README.md @@ -140,6 +140,14 @@ if __name__ == "__main__": An app without one still works, as PyHARP will suppress the second `launch()` with a warning. However, in this case the interface is rebuilt in each worker, so the guard is worth adding. +The behaviour above is covered by a test suite, which is worth running after any change +to `pyharp/worker.py`: + +```bash +pip install -e ".[test]" +pytest tests +``` + A few smaller notes: - Arguments and return values are sent between processes, so they must be picklable. Filepath strings, numbers, booleans and other plain data are fine. An open file handle or a live model object is not. diff --git a/setup.py b/setup.py index 8fefa51..0abc015 100644 --- a/setup.py +++ b/setup.py @@ -22,5 +22,9 @@ # synthesis. Verified working on 0.5.9 against both numpy 1.26 and 2.5; # unpin once it is fixed upstream. 'symusic>=0.5.7,<0.6' - ] + ], + extras_require={ + # Run with "pytest tests" from the repository root + 'test': ['pytest'] + } ) diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..8e47161 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,71 @@ +import os +import sys +import threading + +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from pyharp.worker import JobSupervisor # noqa: E402 + + +@pytest.fixture +def supervisor(): + """A supervisor whose worker is always torn down, so tests leak no processes.""" + made = [] + + def build(timeout_s=30): + instance = JobSupervisor(timeout_s=timeout_s) + made.append(instance) + + return instance + + yield build + + for instance in made: + worker = instance._worker + + if worker is not None and worker.is_alive(): + JobSupervisor._end_process(worker) + + +@pytest.fixture +def collected_messages(monkeypatch): + """ + Captures what the supervisor replays on the worker's behalf. + + gr.Info, gr.Warning and gr.Success all reach the browser through log_message, + which needs a live request context. Replacing it records the calls instead. + """ + messages = [] + + def record(message, title, level="info", duration=10, visible=True): + messages.append({"message": message, "title": title, "level": level}) + + monkeypatch.setattr("gradio.helpers.log_message", record) + + return messages + + +class RecordingProgress: + """Stands in for the progress tracker Gradio injects into a request.""" + + def __init__(self): + self.updates = [] + + def __call__(self, value, desc=None, total=None, unit="steps"): + self.updates.append((round(value, 3), desc)) + + +@pytest.fixture +def progress(): + return RecordingProgress() + + +def cancel_after(supervisor, seconds): + """Cancels from another thread, as the Cancel endpoint does mid-request.""" + timer = threading.Timer(seconds, supervisor.cancel) + timer.daemon = True + timer.start() + + return timer diff --git a/tests/jobs.py b/tests/jobs.py new file mode 100644 index 0000000..773d7d3 --- /dev/null +++ b/tests/jobs.py @@ -0,0 +1,103 @@ +""" +Job functions for the worker tests. + +These live in an importable file rather than in the test module because the worker +is started with "spawn": it reaches a job by importing the file that defines it, so +a function written inside a test could not be found. + +WORKER_ID is generated once when this file is imported. Two jobs reporting the same +value therefore ran in the same worker, which is how the tests tell a reused worker +from a replaced one. +""" + +import os +import signal +import subprocess +import sys +import time +import uuid + +import gradio as gr + + +WORKER_ID = uuid.uuid4().hex + + +class Unsendable: + """Cannot cross a process boundary, to stand in for an accidental return value.""" + + def __reduce__(self): + raise TypeError("Unsendable cannot be pickled") + + +def identify(_=None): + """Reports which process and which worker ran the job.""" + return {"pid": os.getpid(), "worker_id": WORKER_ID} + + +def sleep_interruptibly(seconds): + """Spends its time in Python, so an interrupt unwinds it immediately.""" + deadline = time.monotonic() + seconds + + while time.monotonic() < deadline: + time.sleep(0.05) + + return "finished" + + +def sleep_ignoring_interrupts(seconds): + """Refuses to be interrupted, so stopping it needs the worker to be killed.""" + signal.signal(signal.SIGINT, signal.SIG_IGN) + + return sleep_interruptibly(seconds) + + +def spawn_child_then_sleep(marker_path, seconds): + """ + Starts a grandchild process and waits on it, as a dual-environment app does. + + The grandchild writes its own pid to marker_path so the test can check whether + it outlived the worker. + """ + signal.signal(signal.SIGINT, signal.SIG_IGN) + + script = ( + "import os, sys, time\n" + "open(sys.argv[1], 'w').write(str(os.getpid()))\n" + f"time.sleep({seconds})\n" + ) + + subprocess.run([sys.executable, "-c", script, marker_path], timeout=seconds + 30) + + return "finished" + + +def report_progress(steps): + for step in range(steps): + gr.Progress()((step + 1) / steps, desc=f"step {step + 1}") + + return "finished" + + +def report_messages(): + gr.Info("an informational message", title="Notice") + gr.Warning("a warning message") + + return "finished" + + +def raise_gradio_error(): + raise gr.Error("something the user should see", title="Bad Input", duration=7) + + +def raise_plain_error(): + raise ValueError("an unexpected failure") + + +def return_unsendable(): + return Unsendable() + + +def exit_abruptly(): + """Stands in for an out-of-memory kill: the process dies without reporting.""" + os._exit(1) diff --git a/tests/test_worker.py b/tests/test_worker.py new file mode 100644 index 0000000..9fbdbc3 --- /dev/null +++ b/tests/test_worker.py @@ -0,0 +1,252 @@ +""" +Tests for running process_fn in a worker process. + +The behaviour under test is mostly about what happens when things go wrong - +cancellation, timeouts, crashes - so most of these drive a failure deliberately and +assert on how it is reported. Each one keeps its own timings short; nothing here +should take more than a few seconds. +""" + +import os +import time + +import gradio as gr +import pytest + +from conftest import cancel_after + +import jobs + + +def is_running(pid): + try: + os.kill(pid, 0) + except OSError: + return False + + return True + + +# -------------------------------------------------------------------------------- +# Where the job runs +# -------------------------------------------------------------------------------- + + +def test_job_runs_outside_the_server_process(supervisor, progress): + result = supervisor().run(jobs.identify, progress=progress) + + assert result["pid"] != os.getpid() + + +def test_worker_is_reused_between_jobs(supervisor, progress): + sup = supervisor() + + first = sup.run(jobs.identify, progress=progress) + second = sup.run(jobs.identify, progress=progress) + + assert first["pid"] == second["pid"] + + # The same import, so anything loaded on the way to process_fn was loaded once + assert first["worker_id"] == second["worker_id"] + + +# -------------------------------------------------------------------------------- +# Results and failures +# -------------------------------------------------------------------------------- + + +def test_gradio_error_keeps_its_fields(supervisor, progress): + with pytest.raises(gr.Error) as raised: + supervisor().run(jobs.raise_gradio_error, progress=progress) + + assert raised.value.message == "something the user should see" + assert raised.value.title == "Bad Input" + assert raised.value.duration == 7 + + +def test_plain_exception_arrives_with_its_traceback(supervisor, progress): + with pytest.raises(RuntimeError) as raised: + supervisor().run(jobs.raise_plain_error, progress=progress) + + assert "an unexpected failure" in str(raised.value) + assert "raise_plain_error" in str(raised.value) + + +def test_unsendable_result_is_reported_rather_than_hanging(supervisor, progress): + started = time.monotonic() + + with pytest.raises((RuntimeError, gr.Error)): + supervisor(timeout_s=30).run(jobs.return_unsendable, progress=progress) + + # The point is that it does not wait for the timeout to notice + assert time.monotonic() - started < 15 + + +def test_worker_dying_without_reporting_is_not_a_hang(supervisor, progress): + with pytest.raises(gr.Error) as raised: + supervisor(timeout_s=60).run(jobs.exit_abruptly, progress=progress) + + assert "stopped unexpectedly" in raised.value.message + + +def test_supervisor_recovers_after_a_crash(supervisor, progress): + sup = supervisor(timeout_s=60) + + with pytest.raises(gr.Error): + sup.run(jobs.exit_abruptly, progress=progress) + + assert sup.run(jobs.identify, progress=progress)["pid"] != os.getpid() + + +# -------------------------------------------------------------------------------- +# Calls that need the request context, made from a process that does not have one +# -------------------------------------------------------------------------------- + + +def test_progress_updates_reach_the_request(supervisor, progress): + supervisor().run(jobs.report_progress, 3, progress=progress) + + assert progress.updates == [(0.333, "step 1"), (0.667, "step 2"), (1.0, "step 3")] + + +def test_info_and_warning_reach_the_request(supervisor, progress, collected_messages): + supervisor().run(jobs.report_messages, progress=progress) + + assert collected_messages == [ + {"message": "an informational message", "title": "Notice", "level": "info"}, + {"message": "a warning message", "title": "Warning", "level": "warning"}, + ] + + +def test_progress_is_optional(supervisor): + """A handler Gradio gave no tracker to must not fail when the job reports.""" + assert supervisor().run(jobs.report_progress, 2) == "finished" + + +# -------------------------------------------------------------------------------- +# Cancellation +# -------------------------------------------------------------------------------- + + +def test_cancel_stops_the_job_and_keeps_the_worker(supervisor, progress): + sup = supervisor() + + warm = sup.run(jobs.identify, progress=progress) + + cancel_after(sup, 1) + started = time.monotonic() + + with pytest.raises(gr.Error) as raised: + sup.run(jobs.sleep_interruptibly, 60, progress=progress) + + assert raised.value.message == "Job cancelled." + assert time.monotonic() - started < 10 + + # Interrupted in place, so whatever the worker had loaded is still loaded + assert sup.run(jobs.identify, progress=progress)["worker_id"] == warm["worker_id"] + + +def test_cancel_replaces_a_worker_that_ignores_interrupts(supervisor, progress): + sup = supervisor() + + warm = sup.run(jobs.identify, progress=progress) + + cancel_after(sup, 1) + + with pytest.raises(gr.Error) as raised: + sup.run(jobs.sleep_ignoring_interrupts, 60, progress=progress) + + assert raised.value.message == "Job cancelled." + + # Killing it is the only way to stop it, so the next job gets a fresh worker + assert sup.run(jobs.identify, progress=progress)["worker_id"] != warm["worker_id"] + + +def test_cancel_while_idle_does_nothing(supervisor, progress): + sup = supervisor() + + warm = sup.run(jobs.identify, progress=progress) + sup.cancel() + + assert sup.run(jobs.identify, progress=progress)["worker_id"] == warm["worker_id"] + + +def test_starting_a_job_stops_the_previous_one(supervisor, progress): + """Single-flight: Process runs what was just asked for, not what is queued.""" + sup = supervisor() + + outcome = {} + + def run_slow(): + try: + outcome["result"] = sup.run(jobs.sleep_interruptibly, 60, progress=progress) + except gr.Error as error: + outcome["error"] = error.message + + import threading + + slow = threading.Thread(target=run_slow, daemon=True) + slow.start() + time.sleep(2) + + assert sup.run(jobs.identify, progress=progress)["pid"] != os.getpid() + + slow.join(timeout=15) + assert outcome.get("error") == "Job cancelled." + + +def test_a_stopped_job_does_not_report_into_the_next_one(supervisor, progress): + """The sentinel left by a cancelled job must not be read as the next result.""" + sup = supervisor() + + cancel_after(sup, 1) + + with pytest.raises(gr.Error): + sup.run(jobs.sleep_interruptibly, 60, progress=progress) + + assert sup.run(jobs.identify, progress=progress)["pid"] != os.getpid() + + +def test_processes_the_job_started_are_stopped_with_it(supervisor, progress, tmp_path): + """A model invoked as a subprocess must not outlive the worker that ran it.""" + marker = tmp_path / "grandchild.pid" + + sup = supervisor() + cancel_after(sup, 3) + + with pytest.raises(gr.Error): + sup.run(jobs.spawn_child_then_sleep, str(marker), 60, progress=progress) + + assert marker.exists(), "the job never started its subprocess" + + grandchild = int(marker.read_text()) + + for _ in range(50): + if not is_running(grandchild): + break + time.sleep(0.1) + + assert not is_running(grandchild), f"process {grandchild} outlived its worker" + + +# -------------------------------------------------------------------------------- +# Timeouts +# -------------------------------------------------------------------------------- + + +def test_overrunning_job_is_stopped_at_the_limit(supervisor, progress): + started = time.monotonic() + + with pytest.raises(gr.Error) as raised: + supervisor(timeout_s=2).run(jobs.sleep_interruptibly, 60, progress=progress) + + elapsed = time.monotonic() - started + + # Reported as a timeout, not as the cancellation the worker sees + assert raised.value.message == "Job timed out." + assert raised.value.title == "Timed out" + assert 2 <= elapsed < 15 + + +def test_a_job_within_its_limit_is_left_alone(supervisor, progress): + assert supervisor(timeout_s=30).run(jobs.sleep_interruptibly, 1, progress=progress)