As an Ansible Lightspeed engineer, my team works on the cloud service that interacts with Watson Code Assistant large language models for Ansible task generation. Curious to learn more about the mechanics of training such a model, I set out to create my own, very basic Ansible tasks model. I decided to do this using Red Hat OpenShift Data Science, which made it easy to configure and launch an environment pre-configured with everything I needed to train my model. I’ll walk through the steps I took here.
To complete the steps below, you will need to ensure you have sufficient, dedicated compute available on your cluster. I provisioned a new g5.8xlarge AWS EC2 instance to support this project.
Creating an OpenShift Data Science project
To get started, we’ll create a data science project in Red Hat OpenShift Data Science. From the OpenShift Data Science dashboard, click Data Science Projects and then Create data science project. Provide a name, such as Ansible Playbooks Model, and Create.
From the newly created project, click Create workbench. You will see the Create workbench page shown in Figure 1.
Create your workbench with the following parameters (order of parameters may vary with your OpenShift Data Science version):
- Image selection: CUDA
- Deployment size>Container size: Medium
- Deployment size>Number of GPUs: 1
- Persistent storage size: 200 GiB
- Environment variables>Secret>Key/value:
HUGGINGFACE_TOKEN environment variable value, retrieve your HuggingFace access token from
https://huggingface.co/settings/tokens. You will need to be authenticated with HuggingFace in order to use The Stack dataset.
Once created, launch the workbench to begin fine-tuning.
Note: As you are working through the steps below, you might find at some points that your notebook fails with Out of Memory errors. I worked around this by shutting down the kernel of any notebooks I was not currently using.
Fine-tuning on YAML
With your workbench configured and launched, it is time to fine-tune SantaCoder on YAML. For this, we will use the YAML subset of The Stack dataset from BigCode. You can find two great code samples for fine-tuning SantaCoder in the santacoder-finetuning repo and this Google Colab, which fine-tunes on shell/bash. The GitHub repository provided for use with this tutorial is based on these two samples.
To clone the sample repository, click the Git icon in the OpenShift Data Science workbench left sidebar and click Clone a Repository. Specify the following URL and click Clone:
With the repo cloned, double-click into the
santacoder-ansible folder, then double-click the
train-yaml.ipynb Jupyter notebook file. You should see the contents of the santacoder-ansible repo in your File Browser and the Fine-tune SantaCoder on YAML Jupyter notebook opened and ready to run, as shown in Figure 2.
Run the file cell-by-cell, taking your time to understand each step and explore the various parameters involved. The trainer.train() step should take approximately one hour to complete. Once it's done, you'll see a new folder named
santacoder-finetuned-the-stack-yaml has been created. This contains a local copy of your newly trained model, which now can generate YAML code.
To test your new model, open the
generate.ipynb Jupyter notebook. Notice the first two code cells are for testing YAML or testing Ansible. You'll want to run the Test YAML cell and skip over the Test Ansible cell for now. After running this notebook, you'll see that YAML is successfully generated based on the
test_input string specified in the Test YAML cell. You can switch up this variable to see different generated YAML results.
Fine-tuning on Ansible
Now that we have a model that generates YAML, our last step is to fine-tune it on Ansible content. This step is very similar to the previous one, except we will create our own dataset from scratch this time. While the models used to power Ansible Lightspeed are trained on many thousands of Ansible samples, here we are exercising a HuggingFace capability to generate a data set from the contents of a text file, in this case containing a very small number of example Ansible tasks.
Double-click and run the
train-ansible.ipynb Jupyter notebook. As with the previous step, it should take about an hour to train your new model. Once complete, you'll see another folder created called
generate.ipynb Jupyter notebook again, this time skipping the Test YAML cell and running the Test Ansible cell. After running this notebook, you'll see that an Ansible task is successfully generated based on the test_input string specified in the Test Ansible cell (see Figure 3).
You can see in the
ansible-examples-ec2.txt file that the sample tasks are all based off the amazon.aws.ec2_instance Ansible module and are nearly identical. You can try adding more interesting examples and additional modules to this file, re-training your model, and seeing what outputs you receive.
While creating a model from this extremely small set of Ansible examples is a bit contrived, I hope this article has given you an easy way to get some beginner experience with fine-tuning large language models.
OpenShift Data Science learning resources
Explore learning paths to get started with OpenShift Data Science, including: