huangyanpei/VLA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[lint] Update lint (#16)

Browse Source 
* update lint

* update readme

* update ruff lint
This commit is contained in:
Lufang Chen
2025-09-11 13:18:33 +08:00
committed by GitHub
parent a89dce95aa
commit e9332a283d
28 changed files with 2406 additions and 1074 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wall_x/utils/constant.py
+228 -57
View File
@@ -9,133 +9,292 @@ action_statistic_dof = {
"min": [-3.6176],
"delta": [8.5015],
},
"follow_left_ee_cartesian_pos": {"min": [-0.036, -0.3241, -0.1245], "delta": [0.4389, 0.557, 0.479]},
"follow_left_ee_rotation": {"min": [-1.2373, -0.1929, -1.5182], "delta": [2.2009, 1.5669, 2.0936]},
"follow_left_ee_cartesian_pos": {
"min": [-0.036, -0.3241, -0.1245],
"delta": [0.4389, 0.557, 0.479],
},
"follow_left_ee_rotation": {
"min": [-1.2373, -0.1929, -1.5182],
"delta": [2.2009, 1.5669, 2.0936],
},
"follow_left_gripper": {"min": [-0.1196], "delta": [4.5226]},
"follow_right_ee_cartesian_pos": {"min": [-0.0326, -0.2273, -0.1377], "delta": [0.4574, 0.5704, 0.4743]},
"follow_right_ee_rotation": {"min": [-1.2201, -0.2611, -0.7427], "delta": [2.6623, 1.6622, 2.4186]},
"follow_right_ee_cartesian_pos": {
"min": [-0.0326, -0.2273, -0.1377],
"delta": [0.4574, 0.5704, 0.4743],
},
"follow_right_ee_rotation": {
"min": [-1.2201, -0.2611, -0.7427],
"delta": [2.6623, 1.6622, 2.4186],
},
"follow_right_gripper": {"min": [-0.1208], "delta": [4.5261]},
"height": {"min": [-0.0001], "delta": [0.5051]},
"head_actions": {"min": [-1.5000, -1.4167], "delta": [2.5000, 1.8879]},
"base_velocity": {"min": [-0.0359, -0.084, -0.0162], "delta": [0.1539, 0.1848, 0.0322]},
"base_velocity": {
"min": [-0.0359, -0.084, -0.0162],
"delta": [0.1539, 0.1848, 0.0322],
},
},
"DobbE": {
"follow_right_ee_cartesian_pos": {"min": [-0.6107, -0.3272, -0.4282], "delta": [1.2629, 1.5297, 0.8349]},
"follow_right_ee_rotation": {"min": [-1.7378, -1.4597, -1.8712], "delta": [2.7031, 2.8182, 3.5921]},
"follow_right_ee_cartesian_pos": {
"min": [-0.6107, -0.3272, -0.4282],
"delta": [1.2629, 1.5297, 0.8349],
},
"follow_right_ee_rotation": {
"min": [-1.7378, -1.4597, -1.8712],
"delta": [2.7031, 2.8182, 3.5921],
},
"follow_right_gripper": {"min": [0.0], "delta": [0.9983]},
},
"RH20T": {
"follow_right_ee_cartesian_pos": {"min": [0.3646, -0.2722, 0.0066], "delta": [0.3813, 0.5973, 0.3277]},
"follow_right_ee_rotation": {"min": [-1.8716, -0.4398, -3.1414], "delta": [3.4145, 1.0225, 6.2828]},
"follow_right_ee_cartesian_pos": {
"min": [0.3646, -0.2722, 0.0066],
"delta": [0.3813, 0.5973, 0.3277],
},
"follow_right_ee_rotation": {
"min": [-1.8716, -0.4398, -3.1414],
"delta": [3.4145, 1.0225, 6.2828],
},
"follow_right_gripper": {"min": [0.0], "delta": [95.0]},
},
"agibotworld_alpha": {
"follow_left_ee_cartesian_pos": {"min": [0.4954, 0.0166, 0.1729], "delta": [0.3336, 0.5123, 0.9189]},
"follow_left_ee_rotation": {"min": [-3.1064, -1.2629, -3.1238], "delta": [6.2127, 2.5923, 6.2496]},
"follow_left_ee_cartesian_pos": {
"min": [0.4954, 0.0166, 0.1729],
"delta": [0.3336, 0.5123, 0.9189],
},
"follow_left_ee_rotation": {
"min": [-3.1064, -1.2629, -3.1238],
"delta": [6.2127, 2.5923, 6.2496],
},
"follow_left_gripper": {"min": [34.6222], "delta": [86.1921]},
"follow_right_ee_cartesian_pos": {"min": [0.4615, -0.5975, 0.1638], "delta": [0.3823, 0.5577, 0.8873]},
"follow_right_ee_rotation": {"min": [-3.0891, -1.0739, -2.5091], "delta": [6.1707, 2.3074, 3.8533]},
"follow_right_ee_cartesian_pos": {
"min": [0.4615, -0.5975, 0.1638],
"delta": [0.3823, 0.5577, 0.8873],
},
"follow_right_ee_rotation": {
"min": [-3.0891, -1.0739, -2.5091],
"delta": [6.1707, 2.3074, 3.8533],
},
"follow_right_gripper": {"min": [34.6222], "delta": [85.7635]},
"height": {"min": [0.0], "delta": [0.4535]},
"head_actions": {"min": [-0.1746, 0.0523], "delta": [0.2444, 0.4713]},
},
"austin_buds": {
"follow_right_ee_cartesian_pos": {"min": [0.3496, -0.2855, 0.0105], "delta": [0.3748, 0.492, 0.3116]},
"follow_right_ee_rotation": {"min": [-3.1405, -0.151, -0.0737], "delta": [6.2813, 0.3218, 0.1536]},
"follow_right_ee_cartesian_pos": {
"min": [0.3496, -0.2855, 0.0105],
"delta": [0.3748, 0.492, 0.3116],
},
"follow_right_ee_rotation": {
"min": [-3.1405, -0.151, -0.0737],
"delta": [6.2813, 0.3218, 0.1536],
},
"follow_right_gripper": {"min": [0.0076], "delta": [0.0724]},
},
"austin_sailor": {
"follow_right_ee_cartesian_pos": {"min": [0.387, -0.3165, 0.0244], "delta": [0.2999, 0.5252, 0.2308]},
"follow_right_ee_rotation": {"min": [-3.1402, -0.1618, -1.5918], "delta": [6.2804, 0.337, 2.9478]},
"follow_right_ee_cartesian_pos": {
"min": [0.387, -0.3165, 0.0244],
"delta": [0.2999, 0.5252, 0.2308],
},
"follow_right_ee_rotation": {
"min": [-3.1402, -0.1618, -1.5918],
"delta": [6.2804, 0.337, 2.9478],
},
"follow_right_gripper": {"min": [0.0005], "delta": [0.0773]},
},
"austin_sirius": {
"follow_right_ee_cartesian_pos": {"min": [0.0, -0.1182, 0.0], "delta": [0.5329, 0.3812, 0.2723]},
"follow_right_ee_rotation": {"min": [-3.1407, -0.1243, -1.7434], "delta": [6.2823, 0.1975, 1.8073]},
"follow_right_ee_cartesian_pos": {
"min": [0.0, -0.1182, 0.0],
"delta": [0.5329, 0.3812, 0.2723],
},
"follow_right_ee_rotation": {
"min": [-3.1407, -0.1243, -1.7434],
"delta": [6.2823, 0.1975, 1.8073],
},
"follow_right_gripper": {"min": [0.0334], "delta": [0.046]},
},
"bc_z": {
"follow_right_ee_cartesian_pos": {"min": [-0.3883, -0.1116, 0.6113], "delta": [0.7199, 0.4288, 0.3709]},
"follow_right_ee_rotation": {"min": [-1.056, -1.0587, -2.6295], "delta": [1.9142, 1.9455, 4.8064]},
"follow_right_ee_cartesian_pos": {
"min": [-0.3883, -0.1116, 0.6113],
"delta": [0.7199, 0.4288, 0.3709],
},
"follow_right_ee_rotation": {
"min": [-1.056, -1.0587, -2.6295],
"delta": [1.9142, 1.9455, 4.8064],
},
"follow_right_gripper": {"min": [0.2], "delta": [0.8]},
},
"berkeley_autolab_ur5": {
"follow_right_ee_cartesian_pos": {"min": [0.3018, -0.2129, -0.1888], "delta": [0.3121, 0.52, 0.3107]},
"follow_right_ee_rotation": {"min": [-3.1396, -0.2278, 1.1413], "delta": [6.279, 0.454, 0.9841]},
"follow_right_ee_cartesian_pos": {
"min": [0.3018, -0.2129, -0.1888],
"delta": [0.3121, 0.52, 0.3107],
},
"follow_right_ee_rotation": {
"min": [-3.1396, -0.2278, 1.1413],
"delta": [6.279, 0.454, 0.9841],
},
"follow_right_gripper": {"min": [0.0], "delta": [1.0]},
},
"berkeley_cable_routing": {
"follow_right_ee_cartesian_pos": {"min": [0.4617, -0.28, 0.03], "delta": [0.1838, 0.5665, 0.1272]},
"follow_right_ee_rotation": {"min": [-3.1413, -0.0299, -0.7665], "delta": [6.2826, 0.0692, 3.322]},
"follow_right_ee_cartesian_pos": {
"min": [0.4617, -0.28, 0.03],
"delta": [0.1838, 0.5665, 0.1272],
},
"follow_right_ee_rotation": {
"min": [-3.1413, -0.0299, -0.7665],
"delta": [6.2826, 0.0692, 3.322],
},
},
"berkeley_fanuc_manipulation": {
"follow_right_ee_cartesian_pos": {"min": [0.3718, -0.4072, 0.0184], "delta": [0.3483, 0.7201, 0.5229]},
"follow_right_ee_rotation": {"min": [-3.1399, -1.0166, -1.6988], "delta": [6.2802, 1.4498, 3.2074]},
"follow_right_ee_cartesian_pos": {
"min": [0.3718, -0.4072, 0.0184],
"delta": [0.3483, 0.7201, 0.5229],
},
"follow_right_ee_rotation": {
"min": [-3.1399, -1.0166, -1.6988],
"delta": [6.2802, 1.4498, 3.2074],
},
"follow_right_gripper": {"min": [0.0], "delta": [1.0]},
},
"bridge_data_v2": {
"follow_right_ee_cartesian_pos": {"min": [0.1498, -0.2178, -0.0901], "delta": [0.3012, 0.469, 0.298]},
"follow_right_ee_rotation": {"min": [-0.3279, -0.6105, -1.0578], "delta": [0.7378, 1.0353, 2.2552]},
"follow_right_ee_cartesian_pos": {
"min": [0.1498, -0.2178, -0.0901],
"delta": [0.3012, 0.469, 0.298],
},
"follow_right_ee_rotation": {
"min": [-0.3279, -0.6105, -1.0578],
"delta": [0.7378, 1.0353, 2.2552],
},
"follow_right_gripper": {"min": [0.0692], "delta": [0.9426]},
},
"dlr_edan_shared_control": {
"follow_right_ee_cartesian_pos": {"min": [-0.8387, 0.1473, -0.3934], "delta": [0.6579, 0.6025, 1.1566]},
"follow_right_ee_rotation": {"min": [-3.1217, -1.5197, -2.2516], "delta": [6.2505, 1.5594, 4.2831]},
"follow_right_ee_cartesian_pos": {
"min": [-0.8387, 0.1473, -0.3934],
"delta": [0.6579, 0.6025, 1.1566],
},
"follow_right_ee_rotation": {
"min": [-3.1217, -1.5197, -2.2516],
"delta": [6.2505, 1.5594, 4.2831],
},
"follow_right_gripper": {"min": [0.0], "delta": [1.0]},
},
"droid": {
"follow_right_ee_cartesian_pos": {"min": [0.2667, -0.4396, -0.0472], "delta": [0.5159, 0.8806, 0.8331]},
"follow_right_ee_rotation": {"min": [-3.1374, -1.216, -2.1741], "delta": [6.2749, 2.1075, 4.2259]},
"follow_right_ee_cartesian_pos": {
"min": [0.2667, -0.4396, -0.0472],
"delta": [0.5159, 0.8806, 0.8331],
},
"follow_right_ee_rotation": {
"min": [-3.1374, -1.216, -2.1741],
"delta": [6.2749, 2.1075, 4.2259],
},
"follow_right_gripper": {"min": [0.0], "delta": [0.9912]},
},
"fmb": {
"follow_right_ee_cartesian_pos": {"min": [0.3554, -0.2844, 0.0354], "delta": [0.336, 0.4961, 0.2943]},
"follow_right_ee_rotation": {"min": [-3.1404, -0.9302, -0.0599], "delta": [6.2807, 1.724, 1.8284]},
"follow_right_ee_cartesian_pos": {
"min": [0.3554, -0.2844, 0.0354],
"delta": [0.336, 0.4961, 0.2943],
},
"follow_right_ee_rotation": {
"min": [-3.1404, -0.9302, -0.0599],
"delta": [6.2807, 1.724, 1.8284],
},
"follow_right_gripper": {"min": [0.0], "delta": [1.0]},
},
"fractal": {
"follow_right_ee_cartesian_pos": {"min": [0.3242, -0.2836, 0.1405], "delta": [0.5518, 0.4963, 0.9328]},
"follow_right_ee_rotation": {"min": [-3.1308, -0.2421, -2.9685], "delta": [6.2609, 1.7343, 5.819]},
"follow_right_ee_cartesian_pos": {
"min": [0.3242, -0.2836, 0.1405],
"delta": [0.5518, 0.4963, 0.9328],
},
"follow_right_ee_rotation": {
"min": [-3.1308, -0.2421, -2.9685],
"delta": [6.2609, 1.7343, 5.819],
},
"follow_right_gripper": {"min": [0.0], "delta": [1.0]},
},
"furniture_bench": {
"follow_right_ee_cartesian_pos": {"min": [0.3691, -0.181, 0.0058], "delta": [0.2962, 0.3582, 0.1775]},
"follow_right_ee_rotation": {"min": [-3.1394, -0.6121, -1.9958], "delta": [6.2786, 1.6114, 3.7748]},
"follow_right_ee_cartesian_pos": {
"min": [0.3691, -0.181, 0.0058],
"delta": [0.2962, 0.3582, 0.1775],
},
"follow_right_ee_rotation": {
"min": [-3.1394, -0.6121, -1.9958],
"delta": [6.2786, 1.6114, 3.7748],
},
"follow_right_gripper": {"min": [0.0035], "delta": [0.0762]},
},
"jaco_play": {
"follow_right_ee_cartesian_pos": {"min": [-0.3787, -0.6294, 0.1682], "delta": [0.5898, 0.3587, 0.2183]},
"follow_right_ee_rotation": {"min": [0.9792, -0.0668, -0.0498], "delta": [0.0175, 0.1277, 0.0686]},
"follow_right_ee_cartesian_pos": {
"min": [-0.3787, -0.6294, 0.1682],
"delta": [0.5898, 0.3587, 0.2183],
},
"follow_right_ee_rotation": {
"min": [0.9792, -0.0668, -0.0498],
"delta": [0.0175, 0.1277, 0.0686],
},
"follow_right_gripper": {"min": [0.0791], "delta": [0.1033]},
},
"nyu_rot": {
"follow_right_ee_cartesian_pos": {"min": [0.25, -1.0, -0.2], "delta": [0.75, 2.0, 1.2]},
"follow_right_ee_rotation": {"min": [-3.1416, -3.1416, 6.2831], "delta": [9.4248, 4.1416, 0.0]},
"follow_right_ee_cartesian_pos": {
"min": [0.25, -1.0, -0.2],
"delta": [0.75, 2.0, 1.2],
},
"follow_right_ee_rotation": {
"min": [-3.1416, -3.1416, 6.2831],
"delta": [9.4248, 4.1416, 0.0],
},
"follow_right_gripper": {"min": [0.0], "delta": [1.0]},
},
"stanford_hydra": {
"follow_right_ee_cartesian_pos": {"min": [0.2068, -0.274, 0.1317], "delta": [0.4929, 0.4981, 0.4588]},
"follow_right_ee_rotation": {"min": [-3.1321, -0.7496, -3.0269], "delta": [6.2658, 1.5176, 5.8261]},
"follow_right_ee_cartesian_pos": {
"min": [0.2068, -0.274, 0.1317],
"delta": [0.4929, 0.4981, 0.4588],
},
"follow_right_ee_rotation": {
"min": [-3.1321, -0.7496, -3.0269],
"delta": [6.2658, 1.5176, 5.8261],
},
"follow_right_gripper": {"min": [0.0], "delta": [0.0811]},
},
"stanford_kuka_multimodal": {
"follow_right_ee_cartesian_pos": {"min": [0.4781, -0.0659, 0.3424], "delta": [0.0868, 0.0864, 0.1863]},
"follow_right_ee_rotation": {"min": [-3.136, -0.0521, -3.1413], "delta": [6.2727, 0.1109, 6.2825]},
"follow_right_ee_cartesian_pos": {
"min": [0.4781, -0.0659, 0.3424],
"delta": [0.0868, 0.0864, 0.1863],
},
"follow_right_ee_rotation": {
"min": [-3.136, -0.0521, -3.1413],
"delta": [6.2727, 0.1109, 6.2825],
},
"follow_right_gripper": {"min": [-0.4713], "delta": [0.9485]},
},
"taco_play": {
"follow_right_ee_cartesian_pos": {"min": [0.1375, -0.4291, 0.2052], "delta": [0.5327, 1.0237, 0.3913]},
"follow_right_ee_rotation": {"min": [-3.1391, -0.6946, -1.2808], "delta": [6.2784, 0.8196, 3.0856]},
"follow_right_ee_cartesian_pos": {
"min": [0.1375, -0.4291, 0.2052],
"delta": [0.5327, 1.0237, 0.3913],
},
"follow_right_ee_rotation": {
"min": [-3.1391, -0.6946, -1.2808],
"delta": [6.2784, 0.8196, 3.0856],
},
"follow_right_gripper": {"min": [0.0001], "delta": [0.0806]},
},
"utaustin_mutex": {
"follow_right_ee_cartesian_pos": {"min": [0.3213, -0.4734, 0.0141], "delta": [0.2108, 0.8471, 0.5644]},
"follow_right_ee_rotation": {"min": [-3.1404, -0.2202, -1.5489], "delta": [6.2805, 0.582, 1.9282]},
"follow_right_ee_cartesian_pos": {
"min": [0.3213, -0.4734, 0.0141],
"delta": [0.2108, 0.8471, 0.5644],
},
"follow_right_ee_rotation": {
"min": [-3.1404, -0.2202, -1.5489],
"delta": [6.2805, 0.582, 1.9282],
},
"follow_right_gripper": {"min": [0.0019], "delta": [0.0738]},
},
"viola": {
"follow_right_ee_cartesian_pos": {"min": [0.4011, -0.2521, 0.0103], "delta": [0.2444, 0.4305, 0.4355]},
"follow_right_ee_rotation": {"min": [-3.1403, -0.2737, -1.8626], "delta": [6.2804, 0.4901, 2.0618]},
"follow_right_ee_cartesian_pos": {
"min": [0.4011, -0.2521, 0.0103],
"delta": [0.2444, 0.4305, 0.4355],
},
"follow_right_ee_rotation": {
"min": [-3.1403, -0.2737, -1.8626],
"delta": [6.2804, 0.4901, 2.0618],
},
"follow_right_gripper": {"min": [0.0002], "delta": [0.0773]},
},
"kuka": {
@@ -179,11 +338,23 @@ action_statistic_dof = {
},
},
"agibotworld_beta": {
"follow_left_ee_cartesian_pos": {"min": [0.4954, 0.0166, 0.1729], "delta": [0.3336, 0.5123, 0.9189]},
"follow_left_ee_rotation": {"min": [-3.1064, -1.2629, -3.1238], "delta": [6.2127, 2.5923, 6.2496]},
"follow_left_ee_cartesian_pos": {
"min": [0.4954, 0.0166, 0.1729],
"delta": [0.3336, 0.5123, 0.9189],
},
"follow_left_ee_rotation": {
"min": [-3.1064, -1.2629, -3.1238],
"delta": [6.2127, 2.5923, 6.2496],
},
"follow_left_gripper": {"min": [34.6222], "delta": [86.1921]},
"follow_right_ee_cartesian_pos": {"min": [0.4615, -0.5975, 0.1638], "delta": [0.3823, 0.5577, 0.8873]},
"follow_right_ee_rotation": {"min": [-3.0891, -1.0739, -2.5091], "delta": [6.1707, 2.3074, 3.8533]},
"follow_right_ee_cartesian_pos": {
"min": [0.4615, -0.5975, 0.1638],
"delta": [0.3823, 0.5577, 0.8873],
},
"follow_right_ee_rotation": {
"min": [-3.0891, -1.0739, -2.5091],
"delta": [6.1707, 2.3074, 3.8533],
},
"follow_right_gripper": {"min": [34.6222], "delta": [85.7635]},
"height": {"min": [0.0], "delta": [0.4535]},
"head_actions": {"min": [-0.1746, 0.0523], "delta": [0.2444, 0.4713]},
wall_x/utils/timers.py
+56 -37
View File
@@ -4,23 +4,28 @@ from torch.cuda import nvtx
from abc import ABC, abstractmethod
from typing import List
def _is_distributed():
return torch.distributed.is_available() and torch.distributed.is_initialized()
def _get_world_size():
if _is_distributed():
return torch.distributed.get_world_size()
return 1
def _get_rank():
if _is_distributed():
return torch.distributed.get_rank()
return 0
def _barrier(group=None):
if _is_distributed():
torch.distributed.barrier(group=group)
if torch.distributed.is_available():
try:
dist_all_gather_func = torch.distributed.all_gather_into_tensor
@@ -29,6 +34,7 @@ if torch.distributed.is_available():
else:
dist_all_gather_func = None
class TimerBase(ABC):
"""Timer base class."""
@@ -76,7 +82,7 @@ class DummyTimer(TimerBase):
"""Dummy Timer."""
def __init__(self):
super().__init__('dummy timer')
super().__init__("dummy timer")
def start(self, barrier=False, nvtx_push=False):
return
@@ -89,8 +95,8 @@ class DummyTimer(TimerBase):
def elapsed(self, reset=True, barrier=False):
raise Exception(
'dummy timer should not be used to calculate elapsed time, '
'check if timer\'s log_level <= self._log_level.'
"dummy timer should not be used to calculate elapsed time, "
"check if timer's log_level <= self._log_level."
)
def active_time(self):
@@ -98,8 +104,8 @@ class DummyTimer(TimerBase):
Note: Not supported for DummyTimer.
"""
raise Exception(
'active timer should not be used to calculate elapsed time, '
'check if timer\'s log_level <= self._log_level.'
"active timer should not be used to calculate elapsed time, "
"check if timer's log_level <= self._log_level."
)
@@ -144,7 +150,7 @@ class Timer(TimerBase):
Args:
barrier (bool, optional): Synchronizes ranks before starting. Defaults to False.
"""
assert not self._started, 'timer has already been started'
assert not self._started, "timer has already been started"
if barrier:
_barrier(group=self._barrier_group)
if torch.cuda.is_available():
@@ -154,7 +160,6 @@ class Timer(TimerBase):
if nvtx_push:
nvtx.range_push("{}".format(self.name))
self.nvtx = True
def stop(self, barrier=False, sync=False):
"""Stop the timer.
@@ -164,7 +169,7 @@ class Timer(TimerBase):
"""
if self.nvtx:
nvtx.range_pop()
assert self._started, 'timer is not started'
assert self._started, "timer is not started"
if barrier:
_barrier(group=self._barrier_group)
if torch.cuda.is_available() and sync:
@@ -221,10 +226,10 @@ class Timers:
Allowed: ['max', 'minmax', 'all'].
"""
self._log_level = log_level
allowed_log_options = set(['max', 'minmax', 'all'])
allowed_log_options = set(["max", "minmax", "all"])
assert (
log_option in allowed_log_options
), 'input log option {} is invalid. It must be one of {}'.format(
), "input log option {} is invalid. It must be one of {}".format(
log_option, allowed_log_options
)
self._log_option = log_option
@@ -240,8 +245,10 @@ class Timers:
if name in self._timers:
if log_level is not None:
assert log_level == self._log_levels[name], (
'input log level {} does not match already existing '
'log level {} for {} timer'.format(log_level, self._log_levels[name], name)
"input log level {} does not match already existing "
"log level {} for {} timer".format(
log_level, self._log_levels[name], name
)
)
return self._timers[name]
# If timer does not exist and no log level is provided,
@@ -250,7 +257,7 @@ class Timers:
log_level = self._max_log_level
assert (
log_level <= self._max_log_level
), 'log level {} is larger than max supported log level {}'.format(
), "log level {} is larger than max supported log level {}".format(
log_level, self._max_log_level
)
# Now if the input log level is larger than the one set for
@@ -284,7 +291,7 @@ class Timers:
if torch.cuda.is_available():
device = torch.cuda.current_device()
else:
device = torch.device('cpu')
device = torch.device("cpu")
rank_name_to_time = torch.zeros(
(world_size, len(names)), dtype=torch.float, device=device
@@ -296,7 +303,9 @@ class Timers:
if world_size > 1 and _is_distributed() and dist_all_gather_func is not None:
try:
dist_all_gather_func(rank_name_to_time.view(-1), rank_name_to_time[rank, :].view(-1))
dist_all_gather_func(
rank_name_to_time.view(-1), rank_name_to_time[rank, :].view(-1)
)
except Exception as e:
print(f"Warning: all_gather failed: {e}. Using single rank timing.")
@@ -319,30 +328,38 @@ class Timers:
)
return name_to_min_max_time
def _get_global_min_max_time_string(self, names, reset, barrier, normalizer, max_only):
def _get_global_min_max_time_string(
self, names, reset, barrier, normalizer, max_only
):
"""Report strings for max/minmax times across all ranks."""
name_to_min_max_time = self._get_global_min_max_time(names, reset, barrier, normalizer)
name_to_min_max_time = self._get_global_min_max_time(
names, reset, barrier, normalizer
)
if not name_to_min_max_time:
return None
world_size = _get_world_size()
if world_size == 1:
output_string = 'time (ms):'
output_string = "time (ms):"
for name in name_to_min_max_time:
_, max_time = name_to_min_max_time[name]
output_string += '\n {}: {:.2f}'.format((name + ' ').ljust(48, '.'), max_time)
output_string += "\n {}: {:.2f}".format(
(name + " ").ljust(48, "."), max_time
)
else:
if max_only:
output_string = 'max time across ranks (ms):'
output_string = "max time across ranks (ms):"
else:
output_string = '(min, max) time across ranks (ms):'
output_string = "(min, max) time across ranks (ms):"
for name in name_to_min_max_time:
min_time, max_time = name_to_min_max_time[name]
if max_only:
output_string += '\n {}: {:.2f}'.format((name + ' ').ljust(48, '.'), max_time)
output_string += "\n {}: {:.2f}".format(
(name + " ").ljust(48, "."), max_time
)
else:
output_string += '\n {}: ({:.2f}, {:.2f})'.format(
(name + ' ').ljust(48, '.'), min_time, max_time
output_string += "\n {}: ({:.2f}, {:.2f})".format(
(name + " ").ljust(48, "."), min_time, max_time
)
return output_string
@@ -351,7 +368,7 @@ class Timers:
rank_name_to_time = self._get_elapsed_time_all_ranks(names, reset, barrier)
world_size = _get_world_size()
output_string = 'times across ranks (ms):'
output_string = "times across ranks (ms):"
no_reported_timing = True
for i, name in enumerate(names):
not_yet_found = True
@@ -360,13 +377,13 @@ class Timers:
no_reported_timing = False
if not_yet_found:
not_yet_found = False
output_string += '\n {}:'.format(name)
output_string += "\n {}:".format(name)
if world_size == 1:
output_string += '\n {:.2f}'.format(
output_string += "\n {:.2f}".format(
rank_name_to_time[rank, i] / normalizer
)
else:
output_string += '\n rank {:2d}: {:.2f}'.format(
output_string += "\n rank {:2d}: {:.2f}".format(
rank, rank_name_to_time[rank, i] / normalizer
)
if no_reported_timing:
@@ -398,23 +415,23 @@ class Timers:
str: Formatted string with the timer values.
"""
if names == None: # get all registered timers
if names is None: # get all registered timers
names = list(self._timers.keys())
assert normalizer > 0.0
if self._log_option in ['max', 'minmax']:
if self._log_option in ["max", "minmax"]:
max_only = False
if self._log_option == 'max':
if self._log_option == "max":
max_only = True
output_string = self._get_global_min_max_time_string(
names, reset, barrier, normalizer / 1000.0, max_only
)
elif self._log_option == 'all':
elif self._log_option == "all":
output_string = self._get_all_ranks_time_string(
names, reset, barrier, normalizer / 1000.0
)
else:
raise Exception('unknown timing log option {}'.format(self._log_option))
raise Exception("unknown timing log option {}".format(self._log_option))
return output_string
def log(
@@ -444,7 +461,7 @@ class Timers:
# If no input rank is provided, log on last rank.
world_size = _get_world_size()
current_rank = _get_rank()
if rank is None:
rank = world_size - 1
if rank == current_rank and output_string is not None:
@@ -476,8 +493,10 @@ class Timers:
# torch.utils.add_scalars makes each timer its own run, which
# polutes the runs list, so we just add each as a scalar
assert normalizer > 0.0
name_to_min_max_time = self._get_global_min_max_time(names, reset, barrier, normalizer)
name_to_min_max_time = self._get_global_min_max_time(
names, reset, barrier, normalizer
)
if writer is not None:
for name in name_to_min_max_time:
_, max_time = name_to_min_max_time[name]
writer.add_scalar(name + '-time', max_time, iteration)
writer.add_scalar(name + "-time", max_time, iteration)