dynamic max job duration

app.py

@@ -63,6 +63,13 @@ with gr.Blocks(title="ESM2 Protein Embeddings") as demo:
             value=32,
             label="Batch Size"
         )
+        max_duration = gr.Number(
+            value=3600,
+            label="Max Duration (seconds)",
+            precision=0,
+            minimum=1,
+            maximum=7199
+        )
     
     with gr.Row():
         with gr.Column():
@@ -102,34 +109,34 @@ with gr.Blocks(title="ESM2 Protein Embeddings") as demo:
                     )
 
 
-    def run_pipeline_with_selected_model(fasta_files, model_key, batch_size_value, task="embedding"):
+    def run_pipeline_with_selected_model(fasta_files, model_key, batch_size_value, max_duration, task="embedding"):
         """Wrapper to run pipeline with selected model from dropdown."""
 
         if not fasta_files:
             return gr.update(), "No FASTA files uploaded. Please upload at least one FASTA file for inference."
         model, tokenizer = models_and_tokenizers[model_key]
         if task == "embedding":
-            return full_embedding_pipeline(fasta_files, model, tokenizer, batch_size_value)
+            return full_embedding_pipeline(fasta_files, model, tokenizer, batch_size_value, max_duration)
         elif task == "ppl":
-            return full_ppl_pipeline(fasta_files, model, tokenizer, batch_size_value, mask_percentage=None)
+            return full_ppl_pipeline(fasta_files, model, tokenizer, batch_size_value, None, max_duration)
         elif task == "ppl-approx":
-            return full_ppl_pipeline(fasta_files, model, tokenizer, batch_size_value, mask_percentage=0.1)
+            return full_ppl_pipeline(fasta_files, model, tokenizer, batch_size_value, 0.1, max_duration)
 
     submit_btn.click(
         fn=run_pipeline_with_selected_model,
-        inputs=[input_files, model_dropdown, batch_size, gr.State("embedding")],
+        inputs=[input_files, model_dropdown, batch_size, max_duration, gr.State("embedding")],
         outputs=[download_output, status_output]
     )
 
     ppl_button.click(
         fn=run_pipeline_with_selected_model,
-        inputs=[input_files, model_dropdown, batch_size, gr.State("ppl")],
+        inputs=[input_files, model_dropdown, batch_size, max_duration, gr.State("ppl")],
         outputs=[ppl_download, ppl_status]
     )
 
     ppl_approx_button.click(
         fn=run_pipeline_with_selected_model,
-        inputs=[input_files, model_dropdown, batch_size, gr.State("ppl-approx")],
+        inputs=[input_files, model_dropdown, batch_size, max_duration, gr.State("ppl-approx")],
         outputs=[ppl_download, ppl_status]
     )
 

utils/pipelines.py

@@ -8,8 +8,11 @@ import random
 import os
 import pandas as pd
 
-@spaces.GPU(duration=240)
-def generate_embeddings(sequences_batch, model, tokenizer):
+def get_duration_embeddings(sequences_batch, model, tokenizer, max_duration):
+    return max_duration
+
+@spaces.GPU(duration=get_duration_embeddings)
+def generate_embeddings(sequences_batch, model, tokenizer, max_duration):
     """Generate embeddings for ESM models using the transformers library. 
 
     Parameters:
@@ -55,8 +58,11 @@ def generate_embeddings(sequences_batch, model, tokenizer):
     
     return np.array(sequence_embeddings)
 
-@spaces.GPU(duration=240)
-def generate_ppl_scores(sequences_batch, model, tokenizer):
+def get_duration_ppl(sequences_batch, model, tokenizer, max_duration):
+    return max_duration
+
+@spaces.GPU(duration=get_duration_ppl)
+def generate_ppl_scores(sequences_batch, model, tokenizer, max_duration):
     """Generate pseudo-perplexity scores for ESM models using batched masking across all sequences.
 
     Parameters:
@@ -146,9 +152,11 @@ def generate_ppl_scores(sequences_batch, model, tokenizer):
     
     return ppl_scores
 
+def get_duration_ppl_approx(sequences_batch, model, tokenizer, mask_percentage, max_duration):
+    return max_duration
 
-@spaces.GPU(duration=240)
-def generate_ppl_scores_approx(sequences_batch, model, tokenizer, mask_percentage=0.15):
+@spaces.GPU(duration=get_duration_ppl_approx)
+def generate_ppl_scores_approx(sequences_batch, model, tokenizer, mask_percentage=0.15, max_duration=240):
     """Generate approximate pseudo-perplexity scores for ESM models using chunked masking.
 
     Parameters:
@@ -258,7 +266,7 @@ def generate_ppl_scores_approx(sequences_batch, model, tokenizer, mask_percentag
     
     return ppl_scores
 
-def full_embedding_pipeline(fasta_files, model, tokenizer, batch_size):
+def full_embedding_pipeline(fasta_files, model, tokenizer, batch_size, max_duration):
     """Full pipeline to process FASTA files and generate embeddings from desired model.
     
     Parameters:
@@ -291,7 +299,7 @@ def full_embedding_pipeline(fasta_files, model, tokenizer, batch_size):
         batch = sequences_info[i:i + batch_size]
         batch_sequences = [seq for _, seq, _ in batch]
         
-        embeddings = generate_embeddings(batch_sequences, model, tokenizer)
+        embeddings = generate_embeddings(batch_sequences, model, tokenizer, max_duration)
         status_string += f"Generated {len(embeddings)} embeddings for batch {i // batch_size + 1}/{n_batches}\n"
         all_embeddings.extend(embeddings)
 
@@ -314,7 +322,7 @@ def full_embedding_pipeline(fasta_files, model, tokenizer, batch_size):
     
     return all_file_paths, status_string
     
-def full_ppl_pipeline(fasta_files, model, tokenizer, batch_size, mask_percentage=None):
+def full_ppl_pipeline(fasta_files, model, tokenizer, batch_size, mask_percentage=None, max_duration=240):
     """Full pipeline to process FASTA files and generate embeddings from desired model.
     
     Parameters:
@@ -350,10 +358,10 @@ def full_ppl_pipeline(fasta_files, model, tokenizer, batch_size, mask_percentage
         batch = sequences_info[i:i + batch_size]
         batch_sequences = [seq for _, seq, _ in batch]
         if mask_percentage is None:
-            ppl_scores = generate_ppl_scores(batch_sequences, model, tokenizer)
+            ppl_scores = generate_ppl_scores(batch_sequences, model, tokenizer, max_duration)
             status_string += f"Generated {len(ppl_scores)} pseudo-perplexity scores for batch {i // batch_size + 1}/{n_batches}\n"
         else:
-            ppl_scores = generate_ppl_scores_approx(batch_sequences, model, tokenizer, mask_percentage=mask_percentage)
+            ppl_scores = generate_ppl_scores_approx(batch_sequences, model, tokenizer, mask_percentage=mask_percentage, max_duration=max_duration)
             status_string += f"Generated {len(ppl_scores)} approximate pseudo-perplexity scores for batch {i // batch_size + 1}/{n_batches} with mask percentage {mask_percentage*100:.1f}%\n"
         all_ppl.extend(ppl_scores)