Merge branch 'fix-actions' of https://github.com/nathaniel-hudson/FLoX …

…into fix-actions

.github/workflows/ci.yaml

@@ -35,7 +35,7 @@ jobs:
       - name: install requirements
         run: |
           python -m pip install --upgrade pip setuptools
-          python -m pip install './flox'
+          python -m pip install '.[all]'
           python -m pip install safety
       - name: run safety check
         run: safety check

flox/backends/transfer/proxystore.py

@@ -1,3 +1,6 @@
+from typing import cast
+from uuid import UUID
+
 from proxystore.connectors.endpoint import EndpointConnector
 from proxystore.proxy import Proxy
 from proxystore.store import Store
@@ -17,10 +20,10 @@ def __init__(self, flock: Flock, store: str = "endpoint", name: str = "default")
             )
 
         self.connector = EndpointConnector(
-            endpoints=[node.proxystore_endpoint for node in flock.nodes()]
+            endpoints=[cast(UUID, node.proxystore_endpoint) for node in flock.nodes()]
         )
-        store = Store(name=name, connector=self.connector)
-        self.config = store.config()
+        store_instance = Store(name=name, connector=self.connector)
+        self.config = store_instance.config()
 
     def report(
         self, node_state, node_idx, node_kind, state_dict, history

flox/data/utils.py

@@ -4,7 +4,7 @@
 import matplotlib.pyplot as plt
 import numpy as np
 from matplotlib.axes import Axes
-from scipy import stats
+from scipy import stats  # type: ignore
 from torch.utils.data import DataLoader, Dataset, Subset
 
 from flox.data import FederatedSubsets

flox/run/fit_sync.py

@@ -5,6 +5,7 @@
 from typing import Literal, TypeAlias
 
 import pandas as pd
+from torch.utils.data import Dataset, Subset
 from tqdm import tqdm
 
 from flox.backends.launcher import GlobusComputeLauncher, LocalLauncher
@@ -139,6 +140,7 @@ def sync_flock_traverse(
 
     # If the current node is a worker node, then Launch the LOCAL FITTING job.
     if flock.get_kind(node) is FlockNodeKind.WORKER:
+        dataset: Dataset | Subset
         if isinstance(transfer, ProxyStoreTransfer):
             dataset = transfer.proxy(datasets[node.idx])
         else:

pyproject.toml

@@ -38,6 +38,8 @@ dependencies = [
 [project.optional-dependencies]
 dev = ["black", "coverage", "jupyterlab", "matplotlib", "numpy", "pytest", "seaborn", "tensorboard", "torchvision", "matplotlib-stubs", "pandas-stubs", "networkx-stubs"]
 monitoring = ["tensorboard"]
+proxystore = ["proxystore"]
+all = ["flox[dev,monitoring,proxystore]"]
 
 [tool.pytest.ini_options]
 addopts = [

tox.ini

@@ -16,7 +16,9 @@ commands =
 [testenv:mypy]
 deps =
     mypy>=1.6.1
-extras = dev
+extras =
+    dev
+    proxystore
 commands = mypy --install-types --non-interactive -p flox {posargs}