1
- from typing import TYPE_CHECKING , Dict , Iterable , Union
1
+ from typing import TYPE_CHECKING , Dict , Iterable , Union , List , Optional , Any
2
2
from pathlib import Path
3
3
import os
4
4
import time
@@ -182,7 +182,11 @@ def delete_model_run_data_rows(self, data_row_ids):
182
182
})
183
183
184
184
@experimental
185
- def export_annotations (self , download : bool = False ) -> LabelGenerator :
185
+ def export_labels (
186
+ self ,
187
+ download : bool = False ,
188
+ timeout_seconds : int = 600
189
+ ) -> Optional [Union [str , List [Dict [Any , Any ]]]]:
186
190
"""
187
191
Experimental. To use, make sure client has enable_experimental=True.
188
192
@@ -191,49 +195,40 @@ def export_annotations(self, download: bool = False) -> LabelGenerator:
191
195
Args:
192
196
download (bool): Returns the url if False
193
197
Returns:
194
- URL of the data file that is generated from this ModelRun.
198
+ URL of the data file with this ModelRun's labels.
195
199
If download=True, this instead returns the contents as NDJSON format.
200
+ If the server didn't generate during the `timeout_seconds` period,
201
+ None is returned.
196
202
"""
197
-
203
+ sleep_time = 2
198
204
query_str = """
199
205
mutation exportModelRunAnnotationsPyApi($modelRunId: ID!) {
200
206
exportModelRunAnnotations(data: {modelRunId: $modelRunId}) {
201
207
downloadUrl createdAt status
202
208
}
203
209
}
204
210
"""
205
- url = self .client .execute (
206
- query_str , {'modelRunId' : self .model_id },
207
- experimental = True )['exportModelRunAnnotations' ]['downloadUrl' ]
208
-
209
- counter = 1
210
- while url is None :
211
- logger .info (f"Fetching model run annotations...attempt: { counter } " )
212
- counter += 1
213
- if counter == 10 :
214
- raise Exception (
215
- f"Unsuccessfully got downloadUrl after { counter } attempts." )
216
- time .sleep (10 )
211
+
212
+ while True :
217
213
url = self .client .execute (
218
- query_str , {'modelRunId' : self .model_id },
214
+ query_str , {'modelRunId' : self .uid },
219
215
experimental = True )['exportModelRunAnnotations' ]['downloadUrl' ]
220
216
221
- if not download :
222
- return url
223
- else :
224
- response = requests .get (url )
225
- response .raise_for_status ()
226
- contents = ndjson .loads (response .content )
227
- return contents
228
- """
229
- # Need to discuss how to properly get the asset data type from export
230
- # Until then, we can't do much with contents as the LBV1Converter may
231
- # not infer correct media type and error out
232
- # for now, return the contents as either url or ndjson content
233
- # """
234
-
235
- # labels = LBV1Converter.deserialize(contents)
236
- # return labels
217
+ if url :
218
+ if not download :
219
+ return url
220
+ else :
221
+ response = requests .get (url )
222
+ response .raise_for_status ()
223
+ return ndjson .loads (response .content )
224
+
225
+ timeout_seconds -= sleep_time
226
+ if timeout_seconds <= 0 :
227
+ return None
228
+
229
+ logger .debug ("ModelRun '%s' label export, waiting for server..." ,
230
+ self .uid )
231
+ time .sleep (sleep_time )
237
232
238
233
239
234
class ModelRunDataRow (DbObject ):
0 commit comments