change all 20.04

.github/workflows/code-quality.yaml

@@ -16,7 +16,7 @@ defaults:
     working-directory: .
 jobs:
   code-quality:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-22.04
     timeout-minutes: 15
     if: github.repository_owner == 'mosaicml'
     strategy:

.github/workflows/daily.yaml

@@ -21,52 +21,52 @@ jobs:
       matrix:
         include:
         - name: cpu-3.11-2.3
-          container: mosaicml/pytorch:2.3.1_cpu-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.3.1_cpu-python3.11-ubuntu22.04
           markers: not daily and (remote or not remote) and not gpu and not doctest
           pytest_command: coverage run -m pytest
           composer_package_name: mosaicml
         - name: cpu-3.11-2.4
-          container: mosaicml/pytorch:2.4.1_cpu-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.4.1_cpu-python3.11-ubuntu22.04
           markers: not daily and (remote or not remote) and not gpu and not doctest
           pytest_command: coverage run -m pytest
           composer_package_name: mosaicml
         - name: cpu-3.11-2.5
-          container: mosaicml/pytorch:2.5.1_cpu-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.5.1_cpu-python3.11-ubuntu22.04
           markers: not daily and (remote or not remote) and not gpu and not doctest
           pytest_command: coverage run -m pytest
           composer_package_name: mosaicml
         - name: cpu-3.11-2.5-composer
-          container: mosaicml/pytorch:2.5.1_cpu-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.5.1_cpu-python3.11-ubuntu22.04
           markers: not daily and (remote or not remote) and not gpu and not doctest
           pytest_command: coverage run -m pytest
           composer_package_name: composer
         - name: cpu-doctest
-          container: mosaicml/pytorch:2.4.1_cpu-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.4.1_cpu-python3.11-ubuntu22.04
           markers: not daily and (remote or not remote) and not gpu and doctest
           pytest_command: coverage run -m pytest tests/test_docs.py
           composer_package_name: mosaicml
         - name: daily-cpu-3.11-2.3
-          container: mosaicml/pytorch:2.3.1_cpu-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.3.1_cpu-python3.11-ubuntu22.04
           markers: daily and (remote or not remote) and not gpu and not doctest
           pytest_command: coverage run -m pytest
           composer_package_name: mosaicml
         - name: daily-cpu-3.11-2.4
-          container: mosaicml/pytorch:2.4.1_cpu-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.4.1_cpu-python3.11-ubuntu22.04
           markers: daily and (remote or not remote) and not gpu and not doctest
           pytest_command: coverage run -m pytest
           composer_package_name: mosaicml
         - name: daily-cpu-3.11-2.5
-          container: mosaicml/pytorch:2.5.1_cpu-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.5.1_cpu-python3.11-ubuntu22.04
           markers: daily and (remote or not remote) and not gpu and not doctest
           pytest_command: coverage run -m pytest
           composer_package_name: mosaicml
         - name: daily-cpu-3.11-2.5-composer
-          container: mosaicml/pytorch:2.5.1_cpu-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.5.1_cpu-python3.11-ubuntu22.04
           markers: daily and (remote or not remote) and not gpu and not doctest
           pytest_command: coverage run -m pytest
           composer_package_name: composer
         - name: daily-cpu-doctest
-          container: mosaicml/pytorch:2.5.1_cpu-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.5.1_cpu-python3.11-ubuntu22.04
           markers: daily and (remote or not remote) and not gpu and doctest
           pytest_command: coverage run -m pytest tests/test_docs.py
           composer_package_name: mosaicml
@@ -108,55 +108,55 @@ jobs:
         # Unlike CPU tests, we run daily tests together with GPU tests to minimize launch time
         # on MCLOUD and not eat up all GPUs at once
         - name: "gpu-3.11-2.3-1-gpu"
-          container: mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu22.04
           markers: "(daily or not daily) and (remote or not remote) and gpu and (doctest or not doctest)"
           pytest_command: "coverage run -m pytest"
           composer_package_name: "mosaicml"
           gpu_num: 1
         - name: "gpu-3.11-2.4-1-gpu"
-          container: mosaicml/pytorch:2.4.1_cu124-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.4.1_cu124-python3.11-ubuntu22.04
           markers: "(daily or not daily) and (remote or not remote) and gpu and (doctest or not doctest)"
           pytest_command: "coverage run -m pytest"
           composer_package_name: "mosaicml"
           gpu_num: 1
         - name: "gpu-3.11-2.5-1-gpu"
-          container: mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu22.04
           markers: "(daily or not daily) and (remote or not remote) and gpu and (doctest or not doctest)"
           pytest_command: "coverage run -m pytest"
           composer_package_name: "mosaicml"
           gpu_num: 1
         - name: "gpu-3.11-2.3-2-gpu"
-          container: mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu22.04
           markers: "(daily or not daily) and (remote or not remote) and gpu and (doctest or not doctest)"
           pytest_command: "coverage run -m pytest"
           composer_package_name: "mosaicml"
           gpu_num: 2
         - name: "gpu-3.11-2.4-2-gpu"
-          container: mosaicml/pytorch:2.4.1_cu124-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.4.1_cu124-python3.11-ubuntu22.04
           markers: "(daily or not daily) and (remote or not remote) and gpu and (doctest or not doctest)"
           pytest_command: "coverage run -m pytest"
           composer_package_name: "mosaicml"
           gpu_num: 2
         - name: "gpu-3.11-2.5-2-gpu"
-          container: mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu22.04
           markers: "(daily or not daily) and (remote or not remote) and gpu and (doctest or not doctest)"
           pytest_command: "coverage run -m pytest"
           composer_package_name: "mosaicml"
           gpu_num: 2
         - name: "gpu-3.11-2.3-4-gpu"
-          container: mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu22.04
           markers: "(daily or not daily) and (remote or not remote) and gpu and (doctest or not doctest)"
           pytest_command: "coverage run -m pytest"
           composer_package_name: "mosaicml"
           gpu_num: 4
         - name: "gpu-3.11-2.4-4-gpu"
-          container: mosaicml/pytorch:2.4.1_cu124-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.4.1_cu124-python3.11-ubuntu22.04
           markers: "(daily or not daily) and (remote or not remote) and gpu and (doctest or not doctest)"
           pytest_command: "coverage run -m pytest"
           composer_package_name: "mosaicml"
           gpu_num: 4
         - name: "gpu-3.11-2.5-4-gpu"
-          container: mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu22.04
           markers: "(daily or not daily) and (remote or not remote) and gpu and (doctest or not doctest)"
           pytest_command: "coverage run -m pytest"
           composer_package_name: "mosaicml"

.github/workflows/pr-cpu.yaml

@@ -17,19 +17,19 @@ jobs:
       matrix:
         include:
         - name: cpu-3.11-2.3
-          container: mosaicml/pytorch:2.3.1_cpu-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.3.1_cpu-python3.11-ubuntu22.04
           markers: not daily and not remote and not gpu and not doctest
           pytest_command: coverage run -m pytest
         - name: cpu-3.11-2.4
-          container: mosaicml/pytorch:2.4.1_cpu-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.4.1_cpu-python3.11-ubuntu22.04
           markers: not daily and not remote and not gpu and not doctest
           pytest_command: coverage run -m pytest
         - name: cpu-3.11-2.5
-          container: mosaicml/pytorch:2.5.1_cpu-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.5.1_cpu-python3.11-ubuntu22.04
           markers: not daily and not remote and not gpu and not doctest
           pytest_command: coverage run -m pytest
         - name: cpu-doctest
-          container: mosaicml/pytorch:2.5.1_cpu-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.5.1_cpu-python3.11-ubuntu22.04
           markers: not daily and not remote and not gpu and doctest
           pytest_command: coverage run -m pytest tests/test_docs.py
     steps:

.github/workflows/pr-gpu.yaml

@@ -16,7 +16,7 @@ jobs:
       matrix:
         include:
         - name: gpu-3.11-2.5-1
-          container: mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu22.04
           markers: not daily and not remote and gpu and (doctest or not doctest)
           pytest_command: coverage run -m pytest
           composer_package_name: mosaicml
@@ -45,7 +45,7 @@ jobs:
       matrix:
         include:
         - name: gpu-3.11-2.5-2
-          container: mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu22.04
           markers: not daily and not remote and gpu and (doctest or not doctest)
           pytest_command: coverage run -m pytest
           composer_package_name: mosaicml
@@ -75,7 +75,7 @@ jobs:
       matrix:
         include:
         - name: gpu-3.11-2.5-4
-          container: mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu20.04
+          container: mosaicml/pytorch:2.5.1_cu124-python3.11-ubuntu22.04
           markers: not daily and not remote and gpu and (doctest or not doctest)
           pytest_command: coverage run -m pytest
           composer_package_name: mosaicml

.github/workflows/release.yaml

@@ -8,7 +8,7 @@ on:
 
 jobs:
   code-quality:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-22.04
     strategy:
       matrix:
         python_version:

.github/workflows/smoketest.yaml

@@ -17,7 +17,7 @@ defaults:
     working-directory: .
 jobs:
   smoketest:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-22.04
     timeout-minutes: 10
     if: github.repository_owner == 'mosaicml'
     strategy:

.readthedocs.yaml

@@ -7,7 +7,7 @@ version: 2
 
 # Specify build system and tool dependencies
 build:
-  os: "ubuntu-20.04"
+  os: "ubuntu-22.04"
   tools:
     python: "3.10"
 

docs/source/tutorials/train_resnet50_on_aws.md

@@ -40,7 +40,7 @@ The starting point for this tutorial are the training recipes we present in our
         - PyTorch Version: 1.11.0
         - CUDA Version: 11.3
         - Python Version: 3.9
-        - Ubuntu Version: 20.04
+        - Ubuntu Version: 22.04
 
 ## Launching an AWS EC2 Instance
 

tests/conftest.py

@@ -152,6 +152,9 @@ def pytest_sessionfinish(session: pytest.Session, exitstatus: int):
 
 
 @pytest.fixture(autouse=True)
-def setup_mlflow_tracking(monkeypatch):
+def setup_mlflow_tracking(monkeypatch, tmp_path):
     mlflow = pytest.importorskip('mlflow')
-    monkeypatch.setenv(mlflow.environment_variables.MLFLOW_TRACKING_URI.name, 'databricks')
+    # Use a temporary directory instead of 'databricks'
+    tracking_uri = str(tmp_path / "mlruns")
+    monkeypatch.setenv(mlflow.environment_variables.MLFLOW_TRACKING_URI.name, tracking_uri)
+    os.makedirs(tracking_uri, exist_ok=True)