Hello everyone, We are running metaflow on AWS and have been doing good. Recently have been experimenting with PyG (Graph Based model) and when deploying on AWS it give me this warning.

NOTE: The SHMEM allocation limit is set to the default of 64MB. This may be

insufficient for PyG. NVIDIA recommends the use of the following flags:

docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 ...

and without this flag it is throwing bus error basically unable to allocate a large graph memory. How do i set

--ipc=host --ulimit memlock=-1 --ulimit stack=67108864

in our batch jobs?

What are you using for compute? AWS Batch or Kubernetes?

AWS Batch.

those are set with the batch job definition , under container properties . IIRC metaflow exposes this

@flaky-plumber-70709 I checked the parameters taken by batch decorator there is any keyword in named container properties or any kwargs, how can i pass the value any example would be appreciated.

so this isn't exposed unfortunately, you can see in the batch client where max swap / swappiness are configured: https://github.com/Netflix/metaflow/blob/master/metaflow/plugins/aws/batch/batch_client.py#L282 this is on the linuxParameters settings: https://docs.aws.amazon.com/batch/latest/APIReference/API_LinuxParameters.html in your case you'd prob want to be able to set

sharedMemorySize

For ulimit that is on the parent containerProperties, in either case this would need to be a PR I think to expose. Definitely would be nice to have typed data models for this kind of thing and expose the whole thing, or at least an escape hatch to get at containerProperties

from metaflow import FlowSpec, step, batch

class LightGCNFlow(FlowSpec):

    @batch(
        gpu=1,
        image="docker-image:latest",
        container_properties={
            "linuxParameters": {
                "sharedMemorySize": 2147483648  # 2GB shm
            },
            "ulimits": [
                {"name": "memlock", "softLimit": -1, "hardLimit": -1},
                {"name": "stack", "softLimit": 67108864, "hardLimit": 67108864}
            ]
        }
    )
    @step
    def train(self):
        import torch
        import torch_geometric as pyg
        print("Training LightGCN...")
        #  training code here

if __name__ == "__main__":
    LightGCNFlow()

is this correct ?

no - unfortunately it's not exposed so that won't work, I think it would require some discussion and a PR

@flaky-plumber-70709 sorry, couldn't respond back got a bit occupied. Since we both agree this as a useful feature, I write a github issue and start the discussion once people agree on it, I will write the PR for this