valohai.yaml

- step:
    name: load-data
    image: nvcr.io/nvidia/clara/bionemo-framework:1.3
    environment: pro-trial-prod-oci-vm-gpu-a10-1
    command:
      - pip install valohai-utils
      - python prepare-data.py
    parameters:
      - name: data_version
        type: string
        default: v1.0
        description: Version of the dataset to use when saving preprocessed data
- step:
    name: predict-properties
    image: nvcr.io/nvidia/clara/bionemo-framework:2.4
    environment: pro-trial-prod-oci-vm-gpu-a10-1
    command:
      - pip install -r requirements.txt
      - python predict-properties.py
    environment-variables:
      - name: WORK_DIR
        default: "/valohai/outputs/esm2_inference_tutorial"
        description: Working directory for the inference step
    inputs:
      - name: dataset
        default: "https://drive.google.com/uc?export=download&id=1lKd6CNtZEho_rIXqRiw7zMeDOK5_pjSS"
    parameters:
      - name: num_gpus
        type: integer
        default: 1
        description: Number of GPUs to use
      - name: precision
        type: string
        default: fp16
        description: Precision to use (fp16 or fp32)
      - name: micro-batch-size
        type: integer
        default: 8
        description: Micro batch size for inference
- step:
    name: generate-proteins
    image: nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04
    environment: pro-trial-prod-oci-vm-gpu-a10-1
    command:
      - apt-get update && apt-get install -y python3-pip
      - pip3 install torch==2.7.1+cu118 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
      - pip3 install -r requirements.txt
      - python3 generate-proteins.py
    environment-variables:
      - name: WORK_DIR
        default: "/valohai/outputs/protein_generation"
        description: Working directory for the generation step
    parameters:
      - name: max_length
        type: integer
        default: 100
        description: Maximum length of generated sequences
      - name: top_k
        type: integer
        default: 950
        description: Top-k sampling
      - name: repetition_penalty
        type: float
        default: 1.2
      - name: num_return_sequences
        type: integer
        default: 10
        description: Number of sequences to generate
      - name: eos_token_id
        type: integer
        default: 0
        description: End-of-sequence token ID
- step:
    name: evaluate-generated
    image: docker.io/python:3.10
    environment: pro-trial-prod-oci-vm-gpu-a10-1
    command:
      - pip install torch==2.7.1
      - pip install -r requirements.txt
      - python evaluate-generated.py
    environment-variables:
      - name: WORK_DIR
        default: "/valohai/outputs/evaluate_generated"
        description: Working directory for the generation step
    inputs:
      - name: sequences
        default: "datum://01993303-06cb-3ff9-30b6-a78467ddd67b"
- step:
    name: visualize-similarity
    image: docker.io/python:3.10
    environment: pro-trial-prod-oci-vm-gpu-a10-1
    command:
      - pip install torch pandas scikit-learn matplotlib valohai-utils
      - python protein-similarity-visualization.py
    inputs:
      - name: embeddings
        default: "datum://01993283-53bf-e659-b9df-d77062d76d84"
    parameters:
      - name: query_idx
        type: integer
        default: 0
        description: Query index for similarity search
      - name: topk
        type: integer
        default: 20
        description: Number of top similar proteins to retrieve
- step:
    name: deploy-NIM
    image: docker.io/python:3.10
    environment: pro-trial-prod-oci-vm-gpu-a10-1
    command:
      - pip install valohai-utils transformers>=4.41 torch>=2.6.0
      - python convert-model.py
      - apt-get update && apt-get install -y --no-install-recommends openssh-client ca-certificates && rm -rf /var/lib/apt/lists/*
      # prepare the ssh key (Add your private key as a Valohai input)
      - cp /valohai/inputs/key/tomi-gpu-test-2025-ec2-unsecure.pem /tmp/key.pem
      # set the right permissions
      - chmod 400 /tmp/key.pem
      - ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o IdentitiesOnly=yes -i /tmp/key.pem ec2-user@{parameter-value:host_name} "echo connected"
      # ensure remote target folder exists
      - ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o IdentitiesOnly=yes -i /tmp/key.pem ec2-user@{parameter-value:host_name} "mkdir -p /home/ec2-user/models/protgpt2"
      # copy the model files into that folder
      - scp -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o IdentitiesOnly=yes -i /tmp/key.pem -r /valohai/outputs/my-output/ProtGPT2-safetensors/* ec2-user@{parameter-value:host_name}:/home/ec2-user/models/protgpt2/
      # restart NIM to load the new model
      - ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o IdentitiesOnly=yes -i /tmp/key.pem ec2-user@{parameter-value:host_name} "~/restart_nim.sh"
    parameters:
      - name: host_name
        type: string
        default: 51.44.10.163
    inputs:
      - name: key
        default: "datum://0199760b-5401-37b2-e22f-287f424bb6f3"
- pipeline:
    name: bionemo_end_to_end
    parameters:
      - name: num_gpus
        targets:
          - predict-properties.parameters.num_gpus
        default: 1
      - name: precision
        targets:
          - predict-properties.parameters.precision
        default: fp16
      - name: micro-batch-size
        targets:
          - predict-properties.parameters.micro-batch-size
        default: 8
      - name: max_length
        targets:
          - generate-proteins.parameters.max_length
        default: 100
      - name: top_k
        targets:
          - generate-proteins.parameters.top_k
        default: 950
      - name: repetition_penalty
        targets:
          - generate-proteins.parameters.repetition_penalty
        default: 1.2
      - name: num_return_sequences
        targets:
          - generate-proteins.parameters.num_return_sequences
        default: 10
      - name: eos_token_id
        targets:
          - generate-proteins.parameters.eos_token_id
        default: 0

      - name: query_idx
        targets:
          - visualize-similarity.parameters.query_idx
        default: 0
      - name: topk
        targets:
          - visualize-similarity.parameters.topk
        default: 20

    nodes:
      - name: load-data
        type: execution
        step: load-data
      - name: predict-properties
        type: execution
        step: predict-properties
      - name: visualize-similarity
        type: execution
        step: visualize-similarity
      - name: generate-proteins
        type: execution
        step: generate-proteins
      - name: evaluate-generated
        type: execution
        step: evaluate-generated
      - name: deploy-NIM
        type: execution
        step: deploy-NIM
        actions:
          - when: node-starting
            then: require-approval

    edges:
      - [load-data.outputs.*.zip, predict-properties.input.dataset]
      - [
          predict-properties.outputs.*.zip,
          visualize-similarity.input.embeddings,
        ]
      - [
          generate-proteins.outputs.sequences*.*,
          evaluate-generated.input.sequences,
        ]