Coverage for tests/simulator/test_simulator_actions.py: 100%

1import sys 

2import os 

3import pytest 

4 

5# Add current path 

6sys.path.append(os.getcwd()) 

7 

8from tests.test_utils import temp_sys_path 

9 

10with temp_sys_path("simulator", "streamwise"): 

11 from sim_types import Action 

12 from sim_types import ActionName 

13 from sim_types import GPUType 

14 from sim_types import Model 

15 from sim_types import Result 

16 

17 

18def test_action() -> None: 

19 action = Action( 

20 name=ActionName.ADD_DEVICE_REPLICA, 

21 model=Model.FLUX, 

22 gpu_type=GPUType.A100, 

23 models={GPUType.A100: {Model.FLUX: [2, 1]}}, # 2 devices, 1 replica 

24 action_result=Result( 

25 total_time_s=10.0, 

26 cost=5.0, 

27 total_energy=2.0, 

28 ), 

29 arrival_time_s=0.0, 

30 ) 

31 assert action is not None 

32 assert action.model == Model.FLUX 

33 assert action.name == ActionName.ADD_DEVICE_REPLICA 

34 assert action.gpu_type == GPUType.A100 

35 assert action.models == {GPUType.A100: {Model.FLUX: [2, 1]}} 

36 assert action.time == 10.0 

37 assert action.cost == 5.0 

38 

39 

40def test_action_errors() -> None: 

41 with pytest.raises(TypeError, match="missing 5 required positional arguments"): 

42 Action() # missing parameters 

43 

44 with pytest.raises(ValueError, match="Model flux .* not supported"): 

45 Action( 

46 name=ActionName.ADD_DEVICE_REPLICA, 

47 model="flux", 

48 gpu_type=GPUType.A100, 

49 models={GPUType.A100: {Model.FLUX: [2, 1]}}, # 2 devices, 1 replica 

50 action_result=Result( 

51 total_time_s=10.0, 

52 cost=5.0, 

53 ), 

54 arrival_time_s=0.0, 

55 ) 

56 

57 with pytest.raises(ValueError, match="Device type A100 .* not supported"): 

58 Action( 

59 name=ActionName.ADD_DEVICE_REPLICA, 

60 model=Model.FT, 

61 gpu_type="A100", 

62 models={GPUType.A100: {Model.FT: [1, 1]}}, # 1 device, 1 replica 

63 action_result=Result( 

64 total_time_s=10.0, 

65 cost=5.0, 

66 ), 

67 arrival_time_s=0.0, 

68 ) 

69 

70 with pytest.raises(ValueError, match="Action name .* not supported"): 

71 Action( 

72 name="add_device_replica", 

73 model=Model.FT, 

74 gpu_type=GPUType.A100, 

75 models={GPUType.A100: {Model.FT: [1, 1]}}, # 1 device, 1 replica 

76 action_result=Result( 

77 total_time_s=10.0, 

78 cost=5.0, 

79 ), 

80 arrival_time_s=0.0, 

81 )