Spaces:
Runtime error
Runtime error
File size: 923 Bytes
5f5d58c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
# Test file to ensure that in general certain situational setups for notebooks work.
import os
from pytest import raises
from accelerate import PartialState, notebook_launcher
from accelerate.test_utils import require_bnb
from accelerate.utils import is_bnb_available
def basic_function():
# Just prints the PartialState
print(f"PartialState:\n{PartialState()}")
NUM_PROCESSES = os.environ.get("ACCELERATE_NUM_PROCESSES", 1)
def test_can_initialize():
notebook_launcher(basic_function, (), num_processes=NUM_PROCESSES)
@require_bnb
def test_problematic_imports():
with raises(AssertionError, match="Please keep these imports"):
notebook_launcher(basic_function, (), num_processes=NUM_PROCESSES)
def main():
print("Test basic notebook can be ran")
test_can_initialize()
if is_bnb_available():
print("Test problematic imports (bnb)")
test_problematic_imports()
|