Skip to content

Scale Compute for PLM Embedding Generation #6

Description

@edwardchalstrey1

Pre-requisites:

Objective

Scale embedding generation by moving from local in-memory sequential processing to a distributed setup capable of handling large-scale Protein Language Models (PLMs) and extensive protein datasets.

To handle larger-scale compute and data effectively within our Dagster pipelines, we need to transition our embedding generation to a distributed architecture.

🏗️ Architecture for Scaled Compute

We will integrate Ray or Apache Spark with Dagster to distribute the embedding generation across multiple nodes (or GPUs in our Kubernetes cluster), rather than processing sequentially on a single machine.

✅ Implementation Checklist

  • Refactor the generate_embeddings function to use a distributed computing framework (e.g., Ray).
  • Configure Dagster resources to spin up a Ray cluster dynamically for the inference step.

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions