Hey all. Has anyone been able to get gpu support i...
# ask-metaflows
square-accountant-83677
04/07/2023, 4:56 PMHey all. Has anyone been able to get gpu support in aws batch working? I’m trying to use Jax and it’s not working (various errors or it just doesn’t find the gpu)
✅ 1
c
crooked-jordan-29960
04/07/2023, 5:39 PMFirst thing I'd make sure of is that there are GPUs available in your AWS batch compute environment. Can you see the instance types available to the compute env on AWS Batch dashboard?
s
square-accountant-83677
04/07/2023, 5:53 PMYeah it’s a gpu instance
square-accountant-83677
04/07/2023, 5:54 PMThat was my first thought too
c
crooked-jordan-29960
04/07/2023, 5:54 PMhmm does
jax.devices()crooked-jordan-29960
04/07/2023, 5:55 PMapologies if you have already gone through this basic routine. if you can share stack trace and/or the metaflow step code, I can try to run a test that reproduces what you are seeing..
👍 1
s
square-accountant-83677
04/07/2023, 5:56 PMThat’s a good point. I’ve been checking the platform but I should see if it’s even seeing the gpus at all
square-accountant-83677
04/07/2023, 5:56 PMI’ll see if i can create a minimal example to share
👍 1
square-accountant-83677
04/07/2023, 5:59 PMBefore I go down this route, do you know of any working examples that already exist that I can try to replicate?
c
crooked-jordan-29960
04/07/2023, 6:00 PMNot for jax off top of my head. I'll put together a hello world to verify. Message back here within the hour
s
square-accountant-83677
04/07/2023, 6:00 PMNo kidding? Man that would be awesome
square-accountant-83677
04/07/2023, 6:13 PMThe latest attempt only sees the cpu. I’m using the nvidia cuda12 runtime image, with jax[cuda12_local] installed
c
crooked-jordan-29960
04/07/2023, 6:17 PMvery helpful, will try with this image
s
square-accountant-83677
04/07/2023, 6:18 PMI’m not married to that image, it’s just the latest one I’ve been trying lol
square-accountant-83677
04/07/2023, 6:19 PMI’m just trying to get something that’ll work
c
crooked-jordan-29960
04/07/2023, 6:19 PMnice, even more helpful haha
out of the box it seems that conda + jax don't play too well together, might need to get creative here.
s
square-accountant-83677
04/07/2023, 6:19 PMYeah I noticed that
square-accountant-83677
04/07/2023, 6:19 PMThey have some workarounds, but they’re not too great so far
square-accountant-83677
04/07/2023, 6:20 PMLet me ask you… what is metaflow doing under the hood regarding environments?
c
crooked-jordan-29960
04/07/2023, 6:23 PMmetaflow is making a docker image to run on batch, we can also pass a custom image to batch decorator (described in this metaflow/docker 101 post) if we want more control. lets see if we can use an Nvidia image and then pip install jax on the fly, since using
@conda👍 1
s
square-accountant-83677
04/07/2023, 7:32 PMI found one bug. They changed their install instructions last week. Testing to see if that’s a fix
square-accountant-83677
04/07/2023, 7:32 PMIt worked!
woohoo 1
square-accountant-83677
04/07/2023, 7:33 PMIt was the new install instructions that did it
c
crooked-jordan-29960
04/07/2023, 7:34 PMhow are you installing it now?
s
square-accountant-83677
04/07/2023, 7:35 PMI can’t copy paste (my job blocks slack, so I’m using my phone), but I’m letting Jax install cuda and cudnn
square-accountant-83677
04/07/2023, 7:35 PMThe example in their GitHub readme
c
crooked-jordan-29960
04/07/2023, 7:35 PMnice! thanks. glad you got it sorted out 🙂
s
square-accountant-83677
04/07/2023, 7:36 PMJax[cuda11_pip]
square-accountant-83677
04/07/2023, 7:36 PMThanks for your help!
c
crooked-jordan-29960
04/07/2023, 7:36 PMAnytime