Drop python-3.8 unittest in favor of adding 3.12. Update to python-3.10 for other workflows (#1020)

Summary:

PyTorch dropped support for python-3.8 a while ago. This change makes TorchX consistent with PyTorch. This should also fix gh workflow breakages in https://hud.pytorch.org/hud/pytorch/torchx/main

Changes:

1. Drops python-3.8 in unittest matrix
2. Adds python-3.12 to unittest matrix
3. Upgrades a few libraries (torch, torchvision, etc) in dev-requirements.txt
4. Upgrades to python-3.10 in non-unittest workflows

Reviewed By: manav-a

Differential Revision: D71475986

Author: kiukchung

.github/workflows/aws-batch-integration-tests.yaml

@@ -16,7 +16,7 @@ jobs:
       - name: Setup Python
         uses: actions/setup-python@v2
         with:
-          python-version: 3.9
+          python-version: "3.10"
           architecture: x64
       - name: Checkout TorchX
         uses: actions/checkout@v2

.github/workflows/components-integration-tests.yaml

@@ -35,7 +35,7 @@ jobs:
       - name: Setup Python
         uses: actions/setup-python@v2
         with:
-          python-version: 3.8
+          python-version: "3.10"
           architecture: x64
       - name: Checkout TorchX
         uses: actions/checkout@v2

.github/workflows/container.yaml

@@ -15,7 +15,7 @@ jobs:
       - name: Setup Python
         uses: actions/setup-python@v2
         with:
-          python-version: 3.8
+          python-version: "3.10"
           architecture: x64
       - name: Checkout TorchX
         uses: actions/checkout@v2

.github/workflows/gcp-batch-integration-tests.yaml

@@ -17,7 +17,7 @@ jobs:
       - name: Setup Python
         uses: actions/setup-python@v2
         with:
-          python-version: 3.9
+          python-version: "3.10"
           architecture: x64
       - name: Checkout TorchX
         uses: actions/checkout@v2

.github/workflows/kfp-integration-tests.yaml

@@ -13,7 +13,7 @@ jobs:
       - name: Setup Python
         uses: actions/setup-python@v2
         with:
-          python-version: 3.8
+          python-version: "3.10"
           architecture: x64
       - name: Checkout TorchX
         uses: actions/checkout@v2

.github/workflows/nightly.yaml

@@ -12,7 +12,7 @@ jobs:
       - name: Setup Python 3.10
         uses: actions/setup-python@v2
         with:
-          python-version: '3.10'
+          python-version: "3.10"
           architecture: x64
       - name: Checkout TorchX
         uses: actions/checkout@v2
@@ -22,7 +22,7 @@ jobs:
           pip install -e .[dev]
           pip install twine
       - name: Run tests
-        run: pytest --cov=./ --cov-report=xml 
+        run: pytest --cov=./ --cov-report=xml
       - name: Push nightly
         env:
           PYPI_TOKEN: ${{ secrets.PYPI_TOKEN }}

.github/workflows/pyre.yaml

@@ -21,7 +21,7 @@ jobs:
         run: |
           set -eux
           pip install -e .[dev]
-
+      - name: Init Lint Runner
           lintrunner init
       - name: Run Pyre
         run: scripts/pyre.sh

.github/workflows/python-unittests.yaml

@@ -10,7 +10,7 @@ jobs:
   unittest:
     strategy:
       matrix:
-        python-version: [3.8, 3.9, "3.10", 3.11]
+        python-version: [3.9, "3.10", 3.11, 3.12]
         platform: ["linux.20_04.4x"]
         include:
           - python-version: 3.9

.github/workflows/slurm-local-integration-tests.yaml

@@ -19,7 +19,7 @@ jobs:
       - name: Setup Python
         uses: actions/setup-python@v2
         with:
-          python-version: 3.8
+          python-version: "3.10"
           architecture: x64
       - name: Checkout TorchX
         uses: actions/checkout@v2

dev-requirements.txt

@@ -24,11 +24,11 @@ pytorch-lightning==2.3.1
 tensorboard==2.14.0
 sagemaker==2.224.4
 torch-model-archiver>=0.4.2
-torch==2.2.1
-torchmetrics==0.10.3
+torch>=2.6.0
+torchmetrics==1.6.3
 torchserve>=0.10.0
-torchtext==0.17.1
-torchvision==0.17.1
+torchtext==0.18.0
+torchvision==0.21.0
 ts==0.5.1
 ray[default]
 wheel

torchx/distributed/__init__.py

@@ -83,9 +83,7 @@ def local_device() -> torch.device:
     if dist.is_initialized():
         default_pg = _get_default_group()
         return (
-            local_cuda_device()
-            if default_pg.options.backend == "nccl"
-            else torch.device("cpu")
+            local_cuda_device() if default_pg.name() == "nccl" else torch.device("cpu")
         )
     else:
         return torch.device("cuda") if has_cuda_devices() else torch.device("cpu")

torchx/examples/apps/lightning/model.py

@@ -23,7 +23,7 @@
 import torch
 import torch.jit
 from torch.nn import functional as F
-from torchmetrics import Accuracy
+from torchmetrics.classification import MulticlassAccuracy
 from torchvision.models.resnet import BasicBlock, ResNet
 
 
@@ -44,13 +44,12 @@ def __init__(
 
         # We use the torchvision resnet model with some small tweaks to match
         # TinyImageNet.
-        m = ResNet(BasicBlock, layer_sizes)
+        m = ResNet(BasicBlock, layer_sizes, num_classes=200)
         m.avgpool = torch.nn.AdaptiveAvgPool2d(1)
-        m.fc.out_features = 200
         self.model: ResNet = m
 
-        self.train_acc = Accuracy()
-        self.val_acc = Accuracy()
+        self.train_acc = MulticlassAccuracy(num_classes=m.fc.out_features)
+        self.val_acc = MulticlassAccuracy(num_classes=m.fc.out_features)
 
     # pyre-fixme[14]
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -71,7 +70,7 @@ def validation_step(
     def _step(
         self,
         step_name: str,
-        acc_metric: Accuracy,
+        acc_metric: MulticlassAccuracy,
         batch: Tuple[torch.Tensor, torch.Tensor],
         batch_idx: int,
     ) -> torch.Tensor:

torchx/examples/apps/lightning/train.py

@@ -60,6 +60,7 @@
 import torch
 from pytorch_lightning.callbacks import ModelCheckpoint
 from pytorch_lightning.loggers import TensorBoardLogger
+from torch.distributed.elastic.multiprocessing import errors
 from torchx.examples.apps.lightning.data import (
     create_random_data,
     download_data,
@@ -85,7 +86,12 @@ def parse_args(argv: List[str]) -> argparse.Namespace:
     parser.add_argument(
         "--batch_size", type=int, default=32, help="batch size to use for training"
     )
-    parser.add_argument("--num_samples", type=int, default=10, help="num_samples")
+    parser.add_argument(
+        "--num_samples",
+        type=int,
+        default=32,
+        help="number of samples in the dataset",
+    )
     parser.add_argument(
         "--data_path",
         type=str,
@@ -126,6 +132,7 @@ def get_model_checkpoint(args: argparse.Namespace) -> Optional[ModelCheckpoint]:
     )
 
 
+@errors.record
 def main(argv: List[str]) -> None:
     with tempfile.TemporaryDirectory() as tmpdir:
         args = parse_args(argv)
@@ -138,7 +145,7 @@ def main(argv: List[str]) -> None:
         if not args.data_path:
             data_path = os.path.join(tmpdir, "data")
             os.makedirs(data_path)
-            create_random_data(data_path)
+            create_random_data(data_path, args.num_samples)
         else:
             data_path = download_data(args.data_path, tmpdir)
 

torchx/runner/test/api_test.py

@@ -601,7 +601,7 @@ def test_get_schedulers(self, json_dumps_mock: MagicMock, _) -> None:
             )
             app = AppDef("sleeper", roles=[role])
             runner.run(app, scheduler="local")
-            local_sched_mock.submit.called_once_with(app, {})
+            local_sched_mock.schedule.assert_called_once()
 
     def test_run_from_module(self, _: str) -> None:
         runner = get_runner(name="test_session")