Debugging Axolotl Part 2: Attaching to Docker on a Remote Host

Hamel Husain shows how to debug Axolotl code in VS Code when running inside a Docker container on a remote GPU machine. He sets up SSH access, mounts your project and Hugging Face cache into the container, and uses VS Code's Dev Containers extension to attach the debugger. Breakpoints and local variables work the same as local debugging.
Chapters
Hi, I want to follow up to my last video about debugging in Axolotl with VS Code to show you how to debug with VS Code if you're using Docker and if you're using Docker on a remote host, because it's really common to use Docker on a remote host, and that's the way that lots of people like to develop. So there's a bit more setup required to make it happen, but you can do the same thing with VS Code that I showed before. So just to recap, in the Axolotl repository, we have a docs— under the docs folder, we have a debugging.md, and that goes through how you can set up VS Code to debug Axolotl along with a lot of best practices and tips and tricks on how to get started with debugging. Without stepping on footguns that a lot of newcomers step on. But, um, and there's already a video that shows you the basic setup steps of doing that, um, that doesn't involve Docker or remote host. What I want to take you through today is how to do the Docker and remote host part. So Just to give you some background, I personally develop Axolotl on a remote host. And I recommend that if you are using a remote host, you want to set up your SSH config to have that remote host in there in like some kind of convenient name. So for example, I have a my own GPU rig under my desk. And so I named it rig. And this is what my SSH config, that part of my SSH config, it looks like this. Basically, I call it rig. And then this is the IP address. This happens to be my local network IP address. But for you, if you're using cloud like AWS, GCP, this will actually be the IP address and then your username. And then you might also want to specify your keys, your SSH keys. And you can definitely do that. I do that here with the identity file. And just go into your SSH config and do it. And the reason I mentioned this, because when I'm demonstrating things, I'm going to refer to this rig thing. You might wonder where that's coming from. And that's where it's coming from. So I'm gonna go ahead and close my terminal. So I don't think I'll need that. And the first thing you want to do is start with VS Code. And what you want to do is you're going to want to remote into the host with the SSH remote feature. So if you do, if you use the command palette, so I'm going to do Command+Shift+P because I'm on a Mac. And it might be the same for Windows, I'm not really sure. And then you'll see you can— it happens to be at the top because I always use it, but you can type in remote SSH connect to host, and I'm just going to go and type in rig. That's already in my SSH config, so it's going to show up there. It's going to open a new window and it's going to go into this computer. Now What I'm going to do is, actually, let me, let me bring this back up. Just to show you, I'm going to go back into this computer. And I'm going to change into the axolotl directory. And I'm going to make sure I'm in the right branch here. So I'm going to change to the main branch. Let's say I'm trying to debug something in production. And I'm going to pull from the main branch. Okay, it's already up to date. And then what I'm going to do is I'm actually going to look at the instructions. So these are the instructions on debugging with Axolotl. There is a new section debugging with Docker that I just added. And so you can just copy and paste this specific Docker command is quite long. And all these flags are actually, most of these flags are flags that are recommended by NVIDIA if you're using NVIDIA container runtime. And there's some specific things here, like I'm mounting my current directory to the target directory in the, in the container. And I'm mounting my current Hugging Face cache to the container's Hugging Face cache. Cache, which is sensible in this case, but you can change that if you want to. Okay. And now I'm going to go ahead and run this container. So I'm just going to copy and paste this. And you'll see that now I'm in the container, like this container is running. The next thing I'm going to do per these instructions is I'm going to take— I'm going to make sure that there's an editable install. So I'm gonna come back here. I'm gonna install packaging in this second line, this pip install -e, that's doing an editable install. An editable install just means when you change your code, it'll be reflected right away in the Python package that's installed, the Axolotl Python package that's installed. That way you can debug it and develop on it. Now you are ready for the next step. And this is actually the last step. And the last step is So this is the, this is the VS Code window that is, is remoted into the host. Now I'm going to do the command palette again and do Dev Containers, Dev Containers, Attach to a Running Container. And if you want to, you can just type in Attach to Running Container. And this one says there's no running containers to attach to. So let me, let me see what's happened. Let me see if I can find out. Okay, looks like there is a running container. Maybe I need to connect Maybe I'm not connected properly to the host. Maybe this is the wrong one. So I'll do this again, and you might have to do the same thing. I may have gotten confused between my windows because different windows get popped up. Yeah, so it was probably this one, and I just got confused. But this is actually real. I'm not going to redo this video so you can see how I deal with it. Um, so this is— and then I'm going to do Attach to running container. And then there you go, you'll see this container. I've only had one running. And I'm going to select that container, it's going to open yet another window. It's a little bit of annoying thing that keeps opening new windows. It's not my favorite thing, I get confused sometimes. And it's going to start the container. Okay. And it's gonna— it might give you some warnings. You might have to, um, you might have to reinstall plugins if you like certain plugins, but kind of doesn't matter for this demonstration. And in this case, I want to set a breakpoint, um, right here, and I'm going to go ahead and run the debugger. Now the debugger is going to work because we have this launch.json, and I'm just going to use the one that's here. So I'm gonna go over to debug, and I'm going to start debugging. And you'll see that we are in the container. You can just check by looking at this lower left-hand corner. And that's a way to check to make sure you're in the right window. So it's gonna go ahead and start. And by the way, you can see the debugging command which is right here. And that's the debugging command. That's the command that debugger uses, um, from the launch.json file. And now you see, we are in the debugger. This is the debugger inside the Docker container on the remote host. And we can do debugging per usual. So I can go to the debug console, for example, and I can look at the variables. Like there's prompt. Notice how the local variables are here as well, which is like Um, one of my favorite features. Anyways, this is how you debug with VS Code in Docker on a remote host. Hope that helps.










