4
4
import subprocess
5
5
import tempfile
6
6
import unittest
7
+ import uuid
7
8
import yaml
8
9
9
10
from data_juicer .utils .unittest_utils import DataJuicerTestCaseBase
@@ -45,8 +46,7 @@ def setUp(self):
45
46
super ().setUp ()
46
47
47
48
self .tmp_dir = tempfile .TemporaryDirectory ().name
48
- if not osp .exists (self .tmp_dir ):
49
- os .makedirs (self .tmp_dir )
49
+ os .makedirs (self .tmp_dir , exist_ok = True )
50
50
51
51
def tearDown (self ):
52
52
super ().tearDown ()
@@ -101,36 +101,9 @@ class ProcessDataRayTest(DataJuicerTestCaseBase):
101
101
def setUp (self ):
102
102
super ().setUp ()
103
103
104
- # self._auto_create_ray_cluster()
105
- self .tmp_dir = f'/workspace/tmp/{ self .__class__ .__name__ } '
106
- if not osp .exists (self .tmp_dir ):
107
- os .makedirs (self .tmp_dir )
108
-
109
- def _auto_create_ray_cluster (self ):
110
- try :
111
- # ray cluster already exists, return
112
- run_in_subprocess ('ray status' )
113
- self .tmp_ray_cluster = False
114
- return
115
- except :
116
- pass
117
-
118
- self .tmp_ray_cluster = True
119
- head_port = '6379'
120
- head_addr = '127.0.0.1'
121
- rank = int (os .environ .get ('RANK' , 0 ))
122
-
123
- if rank == 0 :
124
- cmd = f"ray start --head --port={ head_port } --node-ip-address={ head_addr } "
125
- else :
126
- cmd = f"ray start --address={ head_addr } :{ head_port } "
127
-
128
- print (f"current rank: { rank } ; execute cmd: { cmd } " )
129
-
130
- run_in_subprocess (cmd )
131
-
132
- def _close_ray_cluster (self ):
133
- run_in_subprocess ('ray stop' )
104
+ cur_dir = osp .dirname (osp .abspath (__file__ ))
105
+ self .tmp_dir = osp .join (cur_dir , f'tmp_{ uuid .uuid4 ().hex } ' )
106
+ os .makedirs (self .tmp_dir , exist_ok = True )
134
107
135
108
def tearDown (self ):
136
109
super ().tearDown ()
@@ -141,9 +114,6 @@ def tearDown(self):
141
114
import ray
142
115
ray .shutdown ()
143
116
144
- # if self.tmp_ray_cluster:
145
- # self._close_ray_cluster()
146
-
147
117
def test_ray_image (self ):
148
118
tmp_yaml_file = osp .join (self .tmp_dir , 'config_0.yaml' )
149
119
tmp_out_path = osp .join (self .tmp_dir , 'output_0.json' )
0 commit comments