k8s.py

import MDAnalysis as mda
from dask.distributed import Client, get_client
from MDAnalysis.analysis.backends import BackendBase
from MDAnalysis.analysis.dssp import DSSP
from dask_kubernetes.operator import KubeCluster
from dataclasses import dataclass
from dask.delayed import delayed

from typing import NamedTuple, Literal


class Resources(NamedTuple):
    memory: str
    cpu: str

    def as_dict(self):
        return {'memory': self.memory, 'cpu': self.cpu}

@dataclass(frozen=True)
class Config:
    """
    Example:
    --------

    ```python
    config = {
        "name": "dask-cluster-2",
        "image": "<redacted>",
        "n_workers": 7,
        "resources":{"requests": {"memory": "2Gi", "cpu": "800m"},
                     "limits": {"memory": "4Gi", "cpu": "1"}}
        }
    ```
    """

    name: str
    image: str
    n_workers: int
    resources: dict[Literal['requests', 'limits'], Resources]

    def as_dict(self):
        resources = {k:v.as_dict() for k, v in self.resources.items()}
        return {'name': self.name, 'image': self.image, 'n_workers': self.n_workers, 'resources': resources}


def start_cluster(config: Config) -> KubeCluster:
    return KubeCluster(**config.as_dict())

def get_client_for(cluster: KubeCluster) -> Client:
    return Client(cluster)


class DistributedBackend(BackendBase):
    def __init__(self, n_workers):
        super().__init__(n_workers)

    def assign_client(self, client):
        self.client = client

    def apply(self, func, computations):
        return self.client.compute([delayed(func)(c) for c in computations], sync=True)


if __name__ == '__main__':
    u = mda.Universe("YiiP_lipids.gro.gz", "YiiP_lipids.xtc")
    backend = DistributedBackend(n_workers=5)
    config = ...
    cluster = start_cluster(config)
    client = get_client(cluster)
    backend.assign_client(client)

    # test computation
    tasks = [delayed(lambda x: x + 1)(i) for i in range(11)]
    print(f"{client.compute(tasks, sync=True)=}")

    # actual MDAnalysis computation
    print(DSSP(u).run(backend=backend, unsupported_backend=True).results)