Init: Models store API (#15811)

Co-authored-by: Adrian Wälchli <[email protected]>
Co-authored-by: Jirka <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Jirka Borovec <[email protected]>

Author: 4 people

.azure/app-cloud-e2e.yml

@@ -117,6 +117,7 @@ jobs:
     workspace:
       clean: all
     variables:
+      FREEZE_REQUIREMENTS: "1"
       HEADLESS: '1'
       PACKAGE_LIGHTNING: '1'
       CLOUD: '1'
@@ -146,8 +147,6 @@ jobs:
     - bash: |
         pip install -e .[test] \
           -f https://download.pytorch.org/whl/cpu/torch_stable.html
-      env:
-        FREEZE_REQUIREMENTS: "1"
       displayName: 'Install Lightning & dependencies'
 
     - bash: |

.azure/app-cloud-store.yml

@@ -0,0 +1,67 @@
+# Python package
+# Create and test a Python package on multiple Python versions.
+# Add steps that analyze code, save the dist with the build record, publish to a PyPI-compatible index, and more:
+# https://docs.microsoft.com/azure/devops/pipelines/languages/python
+
+trigger:
+  tags:
+    include:
+      - '*'
+  branches:
+    include:
+      - "master"
+      - "release/*"
+      - "refs/tags/*"
+
+pr:
+  branches:
+    include:
+      - "master"
+      - "release/*"
+  paths:
+    include:
+      - ".actions/**"
+      - ".azure/app-cloud-store.yml"
+      - "src/lightning/store/**"
+      - "tests/tests_cloud/**"
+      - "setup.py"
+    exclude:
+      - "*.md"
+      - "**/*.md"
+
+jobs:
+  - job: test_store
+    pool:
+      vmImage: $(imageName)
+    strategy:
+      matrix:
+        Linux:
+          imageName: 'ubuntu-latest'
+        Mac:
+          imageName: 'macOS-latest'
+        Windows:
+          imageName: 'windows-latest'
+    timeoutInMinutes: "20"
+    cancelTimeoutInMinutes: "1"
+    workspace:
+      clean: all
+    variables:
+      FREEZE_REQUIREMENTS: "1"
+      TORCH_URL: "https://download.pytorch.org/whl/cpu/torch_stable.html"
+    steps:
+    - task: UsePythonVersion@0
+      inputs:
+        versionSpec: '3.9'
+
+    - bash: pip install -e .[test] -f $(TORCH_URL)
+      displayName: 'Install Lightning & dependencies'
+
+    - bash: |
+        python -m pytest -m "not cloud" tests_cloud --timeout=300 -v
+      workingDirectory: tests/
+      env:
+        API_KEY: $(LIGHTNING_API_KEY_PROD)
+        API_USERNAME: $(LIGHTNING_USERNAME_PROD)
+        PROJECT_ID: $(LIGHTNING_PROJECT_ID_PROD)
+        LIGHTNING_CLOUD_URL: $(LIGHTNING_CLOUD_URL_PROD)
+      displayName: 'Run the tests'

.gitignore

@@ -48,8 +48,10 @@ wheels/
 *.egg-info/
 .installed.cfg
 *.egg
-src/lightning/*/
 src/*/version.info
+src/lightning/app/
+src/lightning/fabric/
+src/lightning/pytorch/
 
 # PyInstaller
 #  Usually these files are written by a python script from a template

src/lightning/store/README.md

@@ -0,0 +1,99 @@
+## Getting Started
+
+- Login to lightning.ai (_optional_) \<-- takes less than a minute.  ⏩
+- Store your models on the cloud \<-- simple call: `upload_to_cloud(...)`. 🗳️
+- Share it with your friends \<-- just share the "username/model_name" (and version if required) format. :handshake:
+- They download using a simple call: `download_from_cloud("username/model_name", version="your_version")`. :wink:
+- They load your cool model. `load_model("username/model_name", version="your_version")`. :tada:
+- Lightning :zap: fast, isn't it?. :heart:
+
+## Usage
+
+**Storing to the cloud**
+
+```python
+import lightning as L
+from sample.model import LitAutoEncoder, Encoder, Decoder
+
+# Initialize your model here
+autoencoder = LitAutoEncoder(Encoder(), Decoder())
+
+# Pass the model object:
+# No need to pass the username (we'll deduce ourselves), just pass the model name you want as the first argument (with an optional version):
+# format: `model_name:version` (version can either be latest or combination of digits and full-stops: 1.0.0 for example)
+L.store.upload_to_cloud("unique_model_mnist", model=autoencoder, source_code_path="sample")
+
+# version:
+L.store.upload_to_cloud(
+    "unique_model_mnist",
+    version="1.0.0",
+    model=autoencoder,
+    source_code_path="sample/model.py",
+)
+
+# OR: (this will save the file which has the model defined)
+L.store.upload_to_cloud("krshrimali/unique_model_mnist", model=autoencoder)
+```
+
+You can also pass the checkpoint path: `to_lightning_cloud("model_name", version="latest", checkpoint_path=...)`.
+
+**Downloading from the cloud**
+
+At first, you need to download the model to your local machine.
+
+```python
+import lightning as L
+
+L.store.download_from_cloud(
+    "krshrimali/unique_model_mnist",
+    output_dir="your_output_dir",
+)
+# OR: (default to model_storage
+#        $HOME
+#         |- .lightning
+#         |  |- model_store
+#         |  |  |- username
+#         |  |  |  |- <model_name>
+#         |  |  |  |  |- version_<version_with_dots_replaced_by_underscores>
+#      folder)
+L.store.download_from_cloud("krshrimali/unique_model_mnist")
+```
+
+**Loading model**
+
+Then you can load the model to your program.
+
+```python
+import lightning as L
+
+# from <username>.<model_name>.version_<version_with_dots_replaced_by_underscores>.<model_source_file> import LitAutoEncoder, Encoder, Decoder
+model = L.store.load_model("<username>/<model_name>>", version="version")  # version is optional (defaults to latest)
+
+# OR: load weights or checkpoint (if they were uploaded)
+L.store.load_model(
+    "<username>/<model_name>", version="version", load_weights=True | False, load_checkpoint=True | False
+)
+print(model)
+```
+
+**Loading model weights**
+
+```python
+import lightning as L
+from sample.model import LitAutoEncoder, Encoder, Decoder
+
+# If you had passed an `output_dir=...` to download_from_lightning_cloud(...), then you can just do:
+# from output_dir.<model_source_file> import LitAutoEncoder, Encoder, Decoder
+
+model = LitAutoEncoder(Encoder(), Decoder())
+
+model = L.store.load_model(load_weights=True, model=model)
+print("State dict: ", model.state_dict())
+```
+
+Loading checkpoint is similar, just do: `load_checkpoint=True`.
+
+## Known limitations
+
+- missing web UI for user to brows his uploads
+- missing CLI/API to list and delete uploaded models

src/lightning/store/__init__.py

@@ -0,0 +1,3 @@
+from lightning.store.cloud_api import download_from_cloud, load_model, upload_to_cloud
+
+__all__ = ["download_from_cloud", "load_model", "upload_to_cloud"]

src/lightning/store/authentication.py

@@ -0,0 +1,64 @@
+# Copyright The Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import webbrowser
+
+import requests
+from requests.models import HTTPBasicAuth
+
+from lightning.app.core.constants import get_lightning_cloud_url
+from lightning.app.utilities.network import LightningClient
+
+_LIGHTNING_CLOUD_URL = get_lightning_cloud_url()
+
+
+def _get_user_details():
+    client = LightningClient()
+    user_details = client.auth_service_get_user()
+    return user_details.username, user_details.api_key
+
+
+def _get_username_from_api_key(api_key: str):
+    response = requests.get(url=f"{_LIGHTNING_CLOUD_URL}/v1/auth/user", auth=HTTPBasicAuth("lightning", api_key))
+    if response.status_code != 200:
+        raise ConnectionRefusedError(
+            "API_KEY provided is either invalid or wasn't found in the database."
+            " Please ensure that you passed the correct API_KEY."
+        )
+    return json.loads(response.content)["username"]
+
+
+def _check_browser_runnable():
+    try:
+        webbrowser.get()
+    except webbrowser.Error:
+        return False
+    return True
+
+
+def _authenticate(inp_api_key: str = ""):
+    # TODO: we have headless login now,
+    #  so it could be reasonable to just point to that if browser can't be opened / user can't be authed
+    if not inp_api_key:
+        if not _check_browser_runnable():
+            raise ValueError(
+                "Couldn't find a runnable browser in the current system/server."
+                " In order to run the commands on this system, we suggest passing the `api_key`"
+                " after logging into https://lightning.ai."
+            )
+        username, inp_api_key = _get_user_details()
+    else:
+        username = _get_username_from_api_key(inp_api_key)
+    return username, inp_api_key