LightningCLI: Support config files on object stores

[leezu]

🚀 Feature

In addition to python trainer.py --config config.yaml, LightningCLI should enable python trainer.py --config s3://bucket/config.yaml

Motivation

pytorch-lightning has good support for remote filesystems thanks to #2164 #3320 (and possibly others). Thanks to that, it's straightforward to deploy lightning-based scripts to compute cluster environments without having to write extra code preparing / post-processing the local file-systems https://pytorch-lightning.readthedocs.io/en/latest/clouds/cluster.html It would be nice to extend this to LightningCLI.

When using LightningCLI however, it's currently necessary to deploy the config.yaml file to each node, instead of referencing a copy in an object store such as S3.

cc @mauvilsa?

[mauvilsa]

@leezu I think this is a good idea, will look into it.

[edgarriba]

@leezu @mauvilsa thanks for bringing this up. Please, let us know if you need any support with the CLI integration /cc @carmocca

[mauvilsa]

Had a quick look at this and I have the following thoughts/doubts:

• The idea would be to use fsspec, which means this should be implemented in such a way that any protocol registered in fsspec in principle could work.
• The fsspec filesystems might require additional settings, e.g. key, secret, token, etc. for s3fs. Having just s3://bucket/config.yaml is not enough in common use cases. How would these additional settings be specified?
• I see that fsspec is a requirement of lightning. But in the tests I only see simple http loading, e.g. load_from_checkpoint, nothing more complex like s3 or non-url based additional options.

[carmocca]

The idea would be to use fsspec

Do you think this should be done here or in upstream jsonargparse?

How would these additional settings be specified?

I guess environment variables are the best option here

nothing more complex like s3 or non-url based additional options.

Not sure if it's worth for us to have those tests instead of just trusting upstream fsspec

[leezu]

The fsspec filesystems might require additional settings, e.g. key, secret, token, etc. for s3fs. Having just s3://bucket/config.yaml is not enough in common use cases. How would these additional settings be specified?

I guess environment variables are the best option here

Manually setting environment variables is possible, but may not be the so convenient :) If your code is running within an AWS network (eg. EC2), the credentials can be automatically retrieved (assuming your instance is configured to have access to the bucket in question) and fsspec will "just work". See https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/iam-roles-for-amazon-ec2.html#instance-metadata-security-credentials

[mauvilsa]

The idea would be to use fsspec

Do you think this should be done here or in upstream jsonargparse?

jsonargparse already supports http/https for the config, though requires optional packages and be explicitly enabled. The fsspec support should be added extending that. Here what we can do is enable fsspec by default, since it can be assumed that fsspec is installed.

nothing more complex like s3 or non-url based additional options.

Not sure if it's worth for us to have those tests instead of just trusting upstream fsspec

I am interested in looking at code and tests using fsspec. What I was able to find in lightning is not enough to know how to best do this. I have not used fsspec. It looks very powerful but the documentation seems to lack a bit.

The fsspec filesystems might require additional settings, e.g. key, secret, token, etc. for s3fs. Having just s3://bucket/config.yaml is not enough in common use cases. How would these additional settings be specified?

I guess environment variables are the best option here

Manually setting environment variables is possible, but may not be the so convenient :) If your code is running within an AWS network (eg. EC2), the credentials can be automatically retrieved (assuming your instance is configured to have access to the bucket in question) and fsspec will "just work". See https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/iam-roles-for-amazon-ec2.html#instance-metadata-security-credentials

That is nice. Though I think the scope of this shouldn't be to make it work only for s3 on aws. If someone wants to use this in some private cluster which uses a minio server as storage, would fsspec "just work"? Or what if someone wants to use some other file system supported by fsspec? How would users know how to do this or that this is possible? The documentation in lightning and in fsspec is not enough. Or maybe I missed where these things are explained better.

@leezu when you say "environment variables is possible" do you mean that we could implement it like this, or fsspec already supports environment variables? Even if we initially implement it so that it works only for s3 on aws it would be good to have an idea how it can be extended without needing completely change the original implementation.

[leezu]

when you say "environment variables is possible" do you mean that we could implement it like this, or fsspec already supports environment variables? Even if we initially implement it so that it works only for s3 on aws it would be good to have an idea how it can be extended without needing completely change the original implementation.

fsspec just calls the AWS SDK for Python which then attempts to find credentials. Above example just serves as documentation to understand how the SDK does that. Sorry if it was confusing. The takeaway is that you don't need to worry about authentication in jsonargparse/pytorch-lightning and that each of the fsspec backends should have their own ways (via environment variables or other features) to ensure authentication is taken care of.

[mauvilsa]

I have implemented the fsspec support in jsonargparse. @leezu would you mind trying it out with LightningCLI and an s3 bucket? To make it work would be installing jsonargparse from the branch fsspec-support and then in the cli tool add

from jsonargparse import set_config_read_mode
set_config_read_mode(fsspec_enabled=True)

[leezu]

Thank you @mauvilsa. It works great!

[edenlightning]

Closing this issue, feel free to reopen if there's anything to be done on the lightning front.

[leezu]

@edenlightning it'd be helpful for lightning to call

from jsonargparse import set_config_read_mode
set_config_read_mode(fsspec_enabled=True)

as lightning already depends on fsspec

[carmocca]

@leezu Want to open a PR with it?

[mauvilsa]

There isn't yet a jsonargparse release with the fsspec support. When this gets released I can create the PR.