Skip to content
This repository was archived by the owner on Feb 4, 2026. It is now read-only.

Commit 01460ac

Browse files
committed
Final testing
1 parent 38ec8ea commit 01460ac

2 files changed

Lines changed: 137 additions & 29 deletions

File tree

AWS-RoseTTAFold.ipynb

Lines changed: 27 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -197,7 +197,7 @@
197197
"cell_type": "markdown",
198198
"metadata": {},
199199
"source": [
200-
"Get the names of the AWS Batch resources deployed in your account."
200+
"Select the job definitions and Batch queues for your job."
201201
]
202202
},
203203
{
@@ -207,9 +207,23 @@
207207
"outputs": [],
208208
"source": [
209209
"batch_resources = rfutils.get_rosettafold_batch_resources(region=region)\n",
210+
"\n",
211+
"cpu_queue = batch_resources[\"CPUJobQueue\"][0]\n",
212+
"gpu_queue = batch_resources[\"GPUJobQueue\"][0]\n",
213+
"cpu_data_prep_job_def = batch_resources[\"CPUDataPrepJobDefinition\"][0]\n",
214+
"cpu_predict_job_def = batch_resources[\"CPUPredictJobDefinition\"][0]\n",
215+
"gpu_predict_job_def = batch_resources[\"GPUPredictJobDefinition\"][0]\n",
216+
"\n",
210217
"batch_resources"
211218
]
212219
},
220+
{
221+
"cell_type": "markdown",
222+
"metadata": {},
223+
"source": [
224+
"Because our test sequence is small (less than 400 residues) we will run the prediction step on a GPU to decrease the job duration from hours to minutes."
225+
]
226+
},
213227
{
214228
"cell_type": "code",
215229
"execution_count": null,
@@ -219,16 +233,15 @@
219233
"two_step_response = rfutils.submit_2_step_job(\n",
220234
" bucket=bucket,\n",
221235
" job_name=job_name,\n",
222-
" data_prep_input_file=\"input.fa\",\n",
223-
" data_prep_job_definition=batch_resources[\"dataPrepJobDefinition\"][0],\n",
224-
" data_prep_queue=batch_resources[\"dataPrepJobQueue\"][0],\n",
225-
" data_prep_cpu=16,\n",
226-
" data_prep_mem=60,\n",
227-
" predict_job_definition=batch_resources[\"predictJobDefinition\"][0],\n",
228-
" predict_queue=batch_resources[\"predictJobQueue\"][0],\n",
229-
" predict_cpu=24,\n",
230-
" predict_mem=90,\n",
231-
" predict_gpu=1,\n",
236+
" data_prep_job_definition=cpu_data_prep_job_def,\n",
237+
" data_prep_queue=cpu_queue,\n",
238+
" data_prep_cpu=8,\n",
239+
" data_prep_mem=32,\n",
240+
" predict_job_definition=gpu_predict_job_def, # Change this to the cpu_predict_job_def for large proteins\n",
241+
" predict_queue=gpu_queue, # Change this to the cpu_queue for large proteins\n",
242+
" predict_cpu=4,\n",
243+
" predict_mem=16,\n",
244+
" predict_gpu=True, # Change this to False for large proteins\n",
232245
")\n",
233246
"data_prep_jobId = two_step_response[0][\"jobId\"]\n",
234247
"predict_jobId = two_step_response[1][\"jobId\"]"
@@ -248,9 +261,9 @@
248261
"outputs": [],
249262
"source": [
250263
"rfutils.get_rf_job_info(\n",
251-
" cpu_queue=batch_resources[\"dataPrepJobQueue\"][0],\n",
252-
" gpu_queue=batch_resources[\"predictJobQueue\"][0],\n",
253-
" hrs_in_past=3,\n",
264+
" cpu_queue,\n",
265+
" gpu_queue,\n",
266+
" hrs_in_past=1,\n",
254267
")"
255268
]
256269
},

CASP14-Analysis.ipynb

Lines changed: 110 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -35,8 +35,15 @@
3535
"outputs": [],
3636
"source": [
3737
"## Install dependencies\n",
38-
"!pip install -q -q -r requirements.txt\n",
39-
"\n",
38+
"%pip install -q -q -r requirements.txt"
39+
]
40+
},
41+
{
42+
"cell_type": "code",
43+
"execution_count": null,
44+
"metadata": {},
45+
"outputs": [],
46+
"source": [
4047
"## Import helper functions at rfutils/rfutils.py\n",
4148
"from rfutils import rfutils\n",
4249
"\n",
@@ -151,7 +158,14 @@
151158
"metadata": {},
152159
"outputs": [],
153160
"source": [
154-
"batch_resources = rfutils.get_rosettafold_batch_resources()\n",
161+
"batch_resources = rfutils.get_rosettafold_batch_resources(region=region)\n",
162+
"\n",
163+
"cpu_queue = batch_resources[\"CPUJobQueue\"][0]\n",
164+
"gpu_queue = batch_resources[\"GPUJobQueue\"][0]\n",
165+
"cpu_data_prep_job_def = batch_resources[\"CPUDataPrepJobDefinition\"][0]\n",
166+
"cpu_predict_job_def = batch_resources[\"CPUPredictJobDefinition\"][0]\n",
167+
"gpu_predict_job_def = batch_resources[\"GPUPredictJobDefinition\"][0]\n",
168+
"\n",
155169
"batch_resources"
156170
]
157171
},
@@ -168,27 +182,49 @@
168182
"metadata": {},
169183
"outputs": [],
170184
"source": [
171-
"protein_count = 50\n",
185+
"protein_count = 84 # Change this to analyze a smaller number of CASP14 targets\n",
186+
"job_name_list = []\n",
172187
"\n",
173188
"for row in casp14_df[:protein_count].itertuples(index=False):\n",
174189
" record = SeqRecord(row.seq, id=row.id, description=row.description)\n",
175190
" print(f\"Protein sequence for analysis is \\n{record.description}\")\n",
191+
" sequence_length = len(record.seq)\n",
192+
" print(f\"Sequence length is {sequence_length}\")\n",
193+
"\n",
194+
" if sequence_length < 400:\n",
195+
" prep_cpu = 8\n",
196+
" prep_mem = 32\n",
197+
" predict_cpu = 4\n",
198+
" predict_mem = 16\n",
199+
" predict_gpu = True\n",
200+
" predict_job_definition = gpu_predict_job_def\n",
201+
" predict_queue = gpu_queue\n",
202+
" else:\n",
203+
" prep_cpu = 8\n",
204+
" prep_mem = 64\n",
205+
" predict_cpu = 4\n",
206+
" predict_mem = 32\n",
207+
" predict_gpu = False\n",
208+
" predict_job_definition = cpu_predict_job_def\n",
209+
" predict_queue = cpu_queue\n",
210+
"\n",
176211
" job_name = rfutils.create_job_name(record.id)\n",
177212
" print(f\"Automatically-generated job name is: {job_name}\")\n",
213+
" job_name_list.append(job_name)\n",
178214
" input_uri = rfutils.upload_fasta_to_s3(record, bucket, job_name)\n",
179215
" two_step_response = rfutils.submit_2_step_job(\n",
180216
" bucket=bucket,\n",
181217
" job_name=job_name,\n",
182218
" data_prep_input_file=\"input.fa\",\n",
183-
" data_prep_job_definition=batch_resources[\"dataPrepJobDefinition\"][0],\n",
184-
" data_prep_queue=batch_resources[\"dataPrepJobQueue\"][0],\n",
185-
" data_prep_cpu=16,\n",
186-
" data_prep_mem=60,\n",
187-
" predict_job_definition=batch_resources[\"predictJobDefinition\"][0],\n",
188-
" predict_queue=batch_resources[\"predictJobQueue\"][0],\n",
189-
" predict_cpu=24,\n",
190-
" predict_mem=90,\n",
191-
" predict_gpu=1,\n",
219+
" data_prep_job_definition=cpu_data_prep_job_def,\n",
220+
" data_prep_queue=cpu_queue,\n",
221+
" data_prep_cpu=prep_cpu,\n",
222+
" data_prep_mem=prep_mem,\n",
223+
" predict_job_definition=predict_job_definition,\n",
224+
" predict_queue=predict_queue,\n",
225+
" predict_cpu=predict_cpu,\n",
226+
" predict_mem=predict_mem,\n",
227+
" predict_gpu=predict_gpu,\n",
192228
" )"
193229
]
194230
},
@@ -206,11 +242,70 @@
206242
"outputs": [],
207243
"source": [
208244
"rfutils.get_rf_job_info(\n",
209-
" cpu_queue=batch_resources[\"dataPrepJobQueue\"][0],\n",
210-
" gpu_queue=batch_resources[\"predictJobQueue\"][0],\n",
245+
" cpu_queue,\n",
246+
" gpu_queue,\n",
211247
" hrs_in_past=1,\n",
212248
")"
213249
]
250+
},
251+
{
252+
"cell_type": "code",
253+
"execution_count": null,
254+
"metadata": {},
255+
"outputs": [],
256+
"source": [
257+
"jobs = []\n",
258+
"for job_name in job_name_list:\n",
259+
" metrics = rfutils.get_rf_job_metrics(job_name, bucket, region)\n",
260+
" row = [\n",
261+
" job_name,\n",
262+
" metrics[\"DATA_PREP\"][\"JOB_ID\"],\n",
263+
" metrics[\"DATA_PREP\"][\"CPU\"],\n",
264+
" metrics[\"DATA_PREP\"][\"MEM\"],\n",
265+
" metrics[\"DATA_PREP\"][\"LENGTH\"],\n",
266+
" metrics[\"DATA_PREP\"][\"MSA_COUNT\"],\n",
267+
" metrics[\"DATA_PREP\"][\"TEMPLATE_COUNT\"],\n",
268+
" metrics[\"DATA_PREP\"][\"MSA_DURATION\"],\n",
269+
" metrics[\"DATA_PREP\"][\"SS_DURATION\"],\n",
270+
" metrics[\"DATA_PREP\"][\"TEMPLATE_DURATION\"],\n",
271+
" metrics[\"DATA_PREP\"][\"TOTAL_DATA_PREP_DURATION\"],\n",
272+
" metrics[\"PREDICT\"][\"JOB_ID\"],\n",
273+
" metrics[\"PREDICT\"][\"CPU\"],\n",
274+
" metrics[\"PREDICT\"][\"MEM\"],\n",
275+
" metrics[\"PREDICT\"][\"TOTAL_PREDICT_DURATION\"],\n",
276+
" ]\n",
277+
" jobs.append(row)\n",
278+
"metrics_df = pd.DataFrame(\n",
279+
" jobs,\n",
280+
" columns=[\n",
281+
" \"jobName\",\n",
282+
" \"dataPrepJobID\",\n",
283+
" \"dataPrepCPU\",\n",
284+
" \"dataPrepMEM\",\n",
285+
" \"sequenceLength\",\n",
286+
" \"MSACount\",\n",
287+
" \"templateCount\",\n",
288+
" \"MSADuration\",\n",
289+
" \"SSDuration\",\n",
290+
" \"templateDuration\",\n",
291+
" \"dataPrepDuration\",\n",
292+
" \"predictJobId\",\n",
293+
" \"predictCPU\",\n",
294+
" \"predictMEM\",\n",
295+
" \"predictDuration\",\n",
296+
" ],\n",
297+
")\n",
298+
"metrics_df.sort_values(by=[\"dataPrepCPU\", \"dataPrepMEM\", \"predictCPU\", \"predictMEM\"])"
299+
]
300+
},
301+
{
302+
"cell_type": "code",
303+
"execution_count": null,
304+
"metadata": {},
305+
"outputs": [],
306+
"source": [
307+
"metrics_df.to_csv(\"results.csv\")"
308+
]
214309
}
215310
],
216311
"metadata": {

0 commit comments

Comments
 (0)