Added support for AWS Trainium accelerator

[mmcclean-aws]

Added support for the AWS Trainium Accelerator

Tested (run the relevant ones):

• Code formatting: bash format.sh
• Any manual or new tests for this PR (please specify below)
• All smoke tests: pytest tests/test_smoke.py
• Relevant individual smoke tests: pytest tests/test_smoke.py::test_fill_in_the_name
• Backward compatibility tests: bash tests/backward_comaptibility_tests.sh

[cblmemo]

Thanks for contributing to SkyPilot! It is exciting ; ) Left some comments 🙂
I tried

$ sky show-gpus Trainium
Resources 'Trainium' not found. Try 'sky show-gpus --all' to show available accelerators.

and got the error above (Inferentia seems to have the same problem). Could you resolve this?
Also, seems like the workflow has failed. Do we need to update the skypilot-catalog repo too? cc @Michaelvll for a look

[cblmemo]

Shall we add a reference link for this instance type? Such as https://aws.amazon.com/ec2/instance-types/trn1/

[mmcclean-aws]

yes, good idea

[mmcclean-aws]

@cblmemo I see the catalog csv file here is not up to date with the Trainium accelerator type. How can this be regenerated ?

[Michaelvll]

@cblmemo I see the catalog csv file here is not up to date with the Trainium accelerator type. How can this be regenerated ?

Hi @mmcclean-aws, sorry for the delay! The catalog file is generated by the command: python -m sky.clouds.service_catalog.data_fetchers.fetch_aws

If you would like to update the information for the trainium, it would be nice to update the file https://github.com/skypilot-org/skypilot/blob/master/sky/clouds/service_catalog/data_fetchers/fetch_aws.py accordingly.

[Michaelvll]

I just merged the change of fetch_aws.py change into the master first so that our catalog can be automatically updated with trainium. Merging this PR now. Thanks for adding the support for Trainium @mmcclean-aws!

Just tried sky launch --gpus trainium and installed the neuron-ls by directly log into the VM. It seems working!

ubuntu@ip-172-31-35-125:~$ /opt/aws/neuron/bin/neuron-ls
+--------+--------+--------+---------+
| NEURON | NEURON | NEURON |   PCI   |
| DEVICE | CORES  | MEMORY |   BDF   |
+--------+--------+--------+---------+
| 0      | 2      | 32 GB  | 00:1d.0 |
+--------+--------+--------+---------+

Thanks for the contribution @mmcclean-aws!