modal_bioemu.py

import modal

image = (
    modal.Image.micromamba(python_version="3.10")
    .apt_install('gcc')
    .apt_install('git', 'wget')
    .micromamba_install('openmm==8.0.0', 'pdbfixer','kalign2=2.04', 'hhsuite=3.3.0',
                        'mmseqs2', channels=['conda-forge', 'bioconda'])
    .run_commands('pip install "colabfold[alphafold-minus-jax] @ git+https://github.com/sokrypton/ColabFold"')
    .run_commands('pip install "colabfold[alphafold]"')
    .run_commands('pip install --upgrade "jax[cuda12]"==0.4.35')
    .run_commands('pip install --upgrade tensorflow')
    .run_commands('pip install silence_tensorflow')
    .run_commands('python -m colabfold.download')
    #btw: stop here and you would just have regular colabfold. below you could
    #then run colabfold_batch, e.g., with subprocess.run(['colabfold_batch', './', './out'])
    .pip_install('bioemu')
    .run_commands('mkdir /opt/conda/envs')
    #create a fake conda env to fool bioemu:
    .run_commands('ln -s /opt/conda/ /opt/conda/envs/colabfold-bioemu')
    #perform bioemu's patches
    .run_commands('patch /opt/conda/envs/colabfold-bioemu/lib/python3.10/site-packages/alphafold/model/modules.py /opt/conda/lib/python3.10/site-packages/bioemu/colabfold_setup/modules.patch')
    .run_commands('patch /opt/conda/lib/python3.10/site-packages/colabfold/batch.py /opt/conda/lib/python3.10/site-packages/bioemu/colabfold_setup/batch.patch')
    .run_commands('touch /opt/conda/.COLABFOLD_PATCHED')
)

app = modal.App(name="bioemu_run", image=image)

@app.function(
    timeout = 60*30, # you'll need longer times if you want more samples or larger proteins. 
    gpu="L40S" 
)
def bioemu_run(): #t4 lysozyme
    seq = 'MNIFEMLRIDEGLRLKIYKDTEGYYTIGIGHLLTKSPSLNAAKSELDKAIGRNCNGVITKDEAEKLFNQDVDAAVRGILRNAKLKPVYDSLDAVRRCALINMVFQMGETGVAGFTNSLRMLQQKRWDEAAVNLAKSRWYNQTPNRAKRVITTFRTGTWDAYKNL'
    import subprocess
    import os
    # Define the command
    os.mkdir('./out')
    command = [
        "python",
        "-m",
        "bioemu.sample",
        "--sequence", seq,
        "--num_samples", "50",
        "--output_dir", os.path.expanduser("./out")
    ]
    result = subprocess.run(command, capture_output=True, text=True)
    # check success
    if result.returncode == 0:
        print("Command successful")
        print(f"Output: {result.stdout}")
    else:
        print(f"Failed, with return code {result.returncode}")
        print(f"Error: {result.stderr}")

    #tar up the results
    import tarfile
    import os

    def make_targz(output_filename, source_dir):
        with tarfile.open(output_filename, "w:gz") as tar:
            tar.add(source_dir, arcname=os.path.basename(source_dir))

            print(f"Created {output_filename} from {source_dir}")

    make_targz("archive.tar.gz", "./out")

    #return the tar.gz as binary data.
    filedat = open('archive.tar.gz', 'rb').read()
    return filedat

@app.local_entrypoint()
def main():
    #computer go:
    outdat = bioemu_run.remote()
    #save output:
    with open('output.tar.gz', 'wb') as f:
        f.write(outdat)