Hi I am testing out my GPU Compute Environment on AWS Batch via this example. Mamba is not playing well and keeps throwing the following dependency error:

Bootstrapping conda environment...(this could take a few minutes)
    Conda ran into an error while setting up environment.:
    Step: start, Error: command '['/Users/yravindranath/mambaforge/condabin/mamba', 'create', '--yes', '--no-default-packages', '--name', 'metaflow_GPUFlow_linux-64_62c77af6e7b48ded7385bc9cc1c37d614649108d', '--quiet', b'python==3.8', b'requests==>=2.21.0', b'boto3==>=1.14.0', b'pytorch==1.11.0', b'torchvision==0.12.0']' returned error (1): b'Could not solve for environment specs\nThe following packages are incompatible\n\xe2\x94\x9c\xe2\x94\x80 pytorch 1.11.0  is uninstallable because there are no viable options\n\xe2\x94\x82  \xe2\x94\x9c\xe2\x94\x80 pytorch 1.11.0 would require\n\xe2\x94\x82  \xe2\x94\x82  \xe2\x94\x94\xe2\x94\x80 __glibc >=2.17,<3.0.a0 , which is missing on the system;\n\xe2\x94\x82  \xe2\x94\x94\xe2\x94\x80 pytorch 1.11.0 would require\n\xe2\x94\x82     \xe2\x94\x94\xe2\x94\x80 __cuda  , which is missing on the system;\n\xe2\x94\x94\xe2\x94\x80 torchvision 0.12.0  is uninstallable because it requires\n   \xe2\x94\x94\xe2\x94\x80 __glibc >=2.17,<3.0.a0 , which is missing on the system.\n{\n    "success": false\n}\n', stderr=b''

Any idea how to fix it? I tried: 1. Removing

.metaflow

2. Removing

metaflow_GPUFlow_linux-64_62c77af6e7b48ded7385bc9cc1c37d614649108d

3.

mamba clean -a

I am running this on a

g5.xlarge

instance.

Cuda is a PITA in general when it comes to library compatibility. Your best bet is indeed to use an image that comes with drivers (and maybe key libraries) pre installed

Thanks @straight-shampoo-11124 do you have some suggestions on images to use?

there are plenty of (semi)official pytorch images out there, e.g. https://hub.docker.com/r/huggingface/transformers-pytorch-deepspeed-latest-gpu/

Hi @curved-island-17262 I get a very similar error as yours when I try to use @conda_base. Can I know how you use the docker image to fix it?

Bootstrapping conda environment...(this could take a few minutes)
    Conda ran into an error while setting up environment.:
    Step: start, Error: command '['/opt/conda/condabin/mamba', 'create', '--yes', '--no-default-packages', '--name', 'metaflow_T5TensorFlowFlow_linux-64_cfc6e3fe2150d4dc900110bfefecc567883aac18', '--quiet', b'python==3.9.16', b'requests==>=2.21.0', b'boto3==>=1.14.0', b'datasets==2.13.0', b'transformers==4.29.2', b'tensorflow==2.11.1', b'cudatoolkit==11.8.0', b'sentencepiece==0.1.99']' returned error (-9): b'', stderr=b''

And @straight-shampoo-11124 do you have suggestions to use tensorflow and corresponding cudatoolkit with

@conda_base

?

@fast-dog-88545 just pass in the image to your compute backend, for example I am using AWS Batch so I pass in the image to use for the steps:

class GPUFlow(FlowSpec):
    learning_rates = Parameter(
        "learning-rates", default=json.dumps([0.01, 0.001]), type=JSONType
    )

    @batch(image="anibali/pytorch:2.0.0-cuda11.8", queue="staging-gpu-queue", gpu=1)
    @step
    def train(self):
        import torch

        import torch_steps

        gpu_available = torch.cuda.is_available()
        gpu_count = torch.cuda.device_count()
        gpu_name = torch.cuda.get_device_name(torch.cuda.current_device())
        print(f"{gpu_available=}")
        print(f"{gpu_count=}")
        print(f"{gpu_name=}")
        self.model = torch_steps.train_model(self.trainloader, lr=self.input)
        self.next(self.evaluate_model)