Simplify store (#18234)

Co-authored-by: Jirka Borovec <[email protected]>
Co-authored-by: Justus Schock <[email protected]>

Author: 3 people

.azure/app-cloud-store.yml

@@ -302,18 +302,18 @@ subprojects:
 
   - id: "lightning: Store"
     paths:
-      - ".azure/app-cloud-store.yml"
+      - ".github/workflows/ci-tests-store.yml"
       - "src/lightning/__about__.py"
       - "src/lightning/__init__.py"
       - "src/lightning/__main__.py"
       - "src/lightning/__setup__.py"
       - "src/lightning/__version__.py"
       - "src/lightning/store/**"
-      - "tests/tests_cloud/**"
+      - "tests/tests_store/**"
     checks:
-      - "Store.cloud (test_store Linux)"
-      - "Store.cloud (test_store Mac)"
-      - "Store.cloud (test_store Windows)"
+      - "store-cpu (macOS-11, lightning, 3.10, 2.0)"
+      - "store-cpu (ubuntu-20.04, lightning, 3.10, 2.0)"
+      - "store-cpu (windows-2022, lightning, 3.10, 2.0)"
 
   # FixMe: re-enable when BE stabilize
   # - id: "lightning_app: Azure"

.github/checkgroup.yml

@@ -0,0 +1,92 @@
+name: Test Store
+
+# see: https://help.github.com/en/actions/reference/events-that-trigger-workflows
+on:
+  push:
+    branches: [master, "release/*"]
+  pull_request:
+    branches: [master, "release/*"]
+    types: [opened, reopened, ready_for_review, synchronize]  # added `ready_for_review` since draft is skipped
+    paths:
+      - ".actions/*"
+      - "requirements/store/**"
+      - "src/lightning/store/**"
+      - "tests/tests_store/**"
+      - "pyproject.toml"  # includes pytest config
+      - ".github/workflows/ci-tests-store.yml"
+      - "!requirements/*/docs.txt"
+      - "!*.md"
+      - "!**/*.md"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}-${{ github.head_ref }}
+  cancel-in-progress: ${{ ! (github.ref == 'refs/heads/master' || startsWith(github.ref, 'refs/heads/release/')) }}
+
+defaults:
+  run:
+    shell: bash
+
+jobs:
+  store-cpu:
+    runs-on: ${{ matrix.os }}
+    if: github.event.pull_request.draft == false
+    strategy:
+      fail-fast: false
+      matrix:
+        os: ["macOS-11", "ubuntu-20.04", "windows-2022"]
+        pkg-name: ["lightning"]
+        python-version: ["3.10"]
+        pytorch-version: ["2.0"]
+    timeout-minutes: 25  # because of building grpcio on Mac
+    env:
+      PACKAGE_NAME: ${{ matrix.pkg-name }}
+      FREEZE_REQUIREMENTS: ${{ ! (github.ref == 'refs/heads/master' || startsWith(github.ref, 'refs/heads/release/')) }}
+      # PYPI_CACHE_DIR: "_pip-wheels"
+      TORCH_URL_STABLE: "https://download.pytorch.org/whl/cpu/torch_stable.html"
+    steps:
+    - uses: actions/checkout@v3
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Adjust PyTorch versions in requirements files
+      if: ${{ matrix.requires != 'oldest' && matrix.release != 'pre' }}
+      run: |
+        pip install -q wget packaging
+        python -m wget https://raw.githubusercontent.com/Lightning-AI/utilities/main/scripts/adjust-torch-versions.py
+        for fpath in `ls requirements/store/*.txt`; do \
+          python ./adjust-torch-versions.py $fpath ${{ matrix.pytorch-version }}; \
+        done
+
+    - name: Install package & dependencies
+      run: |
+        python -m pip install -q pip -U
+        pip install -e ".[store-test]" "pytest-timeout" -U -f ${TORCH_URL} --prefer-binary
+        pip list
+
+    - name: Testing Store
+      working-directory: tests/tests_store
+      # NOTE: do not include coverage report here, see: https://github.com/nedbat/coveragepy/issues/1003
+      run: |
+        python -m coverage run --source lightning \
+          -m pytest -v --timeout=60 --durations=60
+
+    - name: Statistics
+      if: success()
+      working-directory: tests/tests_store
+      run: |
+        coverage report
+        coverage xml
+
+    - name: Upload coverage to Codecov
+      uses: codecov/codecov-action@v3
+      # see: https://github.com/actions/toolkit/issues/399
+      continue-on-error: true
+      with:
+        token: ${{ secrets.CODECOV_TOKEN }}
+        file: tests/tests_store/coverage.xml
+        flags: lightning,cpu,pytest,python${{ matrix.python-version }}
+        name: CPU-coverage
+        fail_ci_if_error: false

.github/workflows/ci-tests-store.yml

@@ -35,7 +35,7 @@ known_first_party = [
     "parity_pytorch",
     "integrations_app",
     "tests_app",
-    "tests_cloud",
+    "tests_store",
     "tests_fabric",
     "tests_pytorch",
 ]

pyproject.toml

@@ -0,0 +1,5 @@
+coverage ==7.2.7
+pytest ==7.4.0
+pytest-cov ==4.1.0
+pytest-rerunfailures ==10.3
+pytest-random-order ==1.1.0

requirements/store/test.txt

@@ -29,6 +29,7 @@
 from lightning.pytorch.trainer import Trainer  # noqa: E402
 
 import lightning.app  # isort: skip # noqa: E402
+import lightning.store  # isort: skip # noqa: E402
 
 lightning.app._PROJECT_ROOT = os.path.dirname(lightning.app._PROJECT_ROOT)
 

src/lightning/__init__.py

@@ -4,7 +4,6 @@
 - Store your models on the cloud \<-- simple call: `upload_model(...)`. 🗳️
 - Share it with your friends \<-- just share the "username/model_name" (and version if required) format. :handshake:
 - They download using a simple call: `download_model("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
@@ -13,87 +12,30 @@
 
 ```python
 import lightning as L
-from sample.model import LitAutoEncoder, Encoder, Decoder
 
-# Initialize your model here
-autoencoder = LitAutoEncoder(Encoder(), Decoder())
+# Upload a checkpoint:
+L.store.upload_model("mnist_model", "mnist_model.ckpt")
 
-# 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_model("unique_model_mnist", model=autoencoder, source_code_path="sample")
-
-# version:
-L.store.upload_model(
-    "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_model("krshrimali/unique_model_mnist", model=autoencoder)
-```
-
-You can also pass the checkpoint path: `upload_model("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_model(
-    "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_model("krshrimali/unique_model_mnist")
+# Optionally provide a version:
+L.store.upload_model("mnist_model", "mnist_model.ckpt", version="1.0.0")
 ```
 
-**Loading model**
-
-Then you can load the model to your program.
+**List your models**
 
 ```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)
+models = L.store.list_models()
 
-# 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)
+print([model.name for model in models])
+# ['username/mnist_model']
 ```
 
-**Loading model weights**
+**Downloading from the cloud**
 
 ```python
 import lightning as L
-from sample.model import LitAutoEncoder, Encoder, Decoder
-
-# If you had passed an `output_dir=...` to download_model(...), 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())
+# Download a checkpoint
+L.store.download_model("username/mnist_model", "any_path.ckpt")
 ```
-
-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/README.md

@@ -1,3 +1,3 @@
-from lightning.store.cloud_api import download_model, load_model, upload_model
+from lightning.store.store import download_model, list_models, upload_model
 
-__all__ = ["download_model", "load_model", "upload_model"]
+__all__ = ["download_model", "upload_model", "list_models"]