@@ -132,49 +132,83 @@ def upsert_predictions_and_send_to_project(
132132 project_id : str ,
133133 priority : Optional [int ] = 5 ,
134134 ) -> 'MEAPredictionImport' : # type: ignore
135- """ Upload predictions and create a batch import to project.
135+ """
136+ Provides a convenient way to execute the following steps in a single function call:
137+ 1. Upload predictions to a Model
138+ 2. Create a batch from data rows that had predictions assocated with them
139+ 3. Attach the batch to a project
140+ 4. Add those same predictions to the project as MAL annotations
141+
142+ Note that partial successes are possible.
143+ If it is important that all stages are successful then check the status of each individual task
144+ with task.errors. E.g.
145+
146+ >>> mea_import_job, batch, mal_import_job = upsert_predictions_and_send_to_project(name, predictions, project_id)
147+ >>> # handle mea import job successfully created (check for job failure or partial failures)
148+ >>> print(mea_import_job.status, mea_import_job.errors)
149+ >>> if batch is None:
150+ >>> # Handle batch creation failure
151+ >>> if mal_import_job is None:
152+ >>> # Handle mal_import_job creation failure
153+ >>> else:
154+ >>> # handle mal import job successfully created (check for job failure or partial failures)
155+ >>> print(mal_import_job.status, mal_import_job.errors)
156+
157+
136158 Args:
137159 name (str): name of the AnnotationImport job as well as the name of the batch import
138160 predictions (Iterable):
139161 iterable of annotation rows
140162 project_id (str): id of the project to import into
141163 priority (int): priority of the job
142164 Returns:
143- (MEAPredictionImport, Batch, MEAToMALPredictionImport)
165+ Tuple[MEAPredictionImport, Batch, MEAToMALPredictionImport]
166+ If any of these steps fail the return value will be None.
167+
144168 """
145169 kwargs = dict (client = self .client , model_run_id = self .uid , name = name )
146170 project = self .client .get_project (project_id )
147171 import_job = self .add_predictions (name , predictions )
148172 prediction_statuses = import_job .statuses
149- mea_to_mal_data_rows_set = set ([
150- row ['dataRow' ]['id' ]
151- for row in prediction_statuses
152- if row ['status' ] == 'SUCCESS'
153- ])
154173 mea_to_mal_data_rows = list (
155- mea_to_mal_data_rows_set )[:DATAROWS_IMPORT_LIMIT ]
156-
157- if len (mea_to_mal_data_rows ) >= DATAROWS_IMPORT_LIMIT :
174+ set ([
175+ row ['dataRow' ]['id' ]
176+ for row in prediction_statuses
177+ if row ['status' ] == 'SUCCESS'
178+ ]))
179+
180+ if not mea_to_mal_data_rows :
181+ # 0 successful model predictions imported
182+ return import_job , None , None
158183
184+ elif len (mea_to_mal_data_rows ) >= DATAROWS_IMPORT_LIMIT :
185+ mea_to_mal_data_rows = mea_to_mal_data_rows [:DATAROWS_IMPORT_LIMIT ]
159186 logger .warning (
160- f"Got { len ( mea_to_mal_data_rows_set ) } rows to import , trimmed down to { DATAROWS_IMPORT_LIMIT }
187+ f"Exeeded max data row limit { len ( mea_to_mal_data_rows ) } { DATAROWS_IMPORT_LIMIT } . "
161188 )
162- if len (mea_to_mal_data_rows ) == 0 :
163- return import_job , None , None
164189
165190 try :
166191 batch = project .create_batch (name , mea_to_mal_data_rows , priority )
167- try :
168- mal_prediction_import = Entity .MEAToMALPredictionImport .create_for_model_run_data_rows (
169- data_row_ids = mea_to_mal_data_rows ,
170- project_id = project_id ,
171- ** kwargs )
172- return import_job , batch , mal_prediction_import
173- except :
174- return import_job , batch , None
175- except :
192+ except Exception as e :
193+ logger .warning (f"Failed to create batch. Messsage : { e } )
194+ # Unable to create batch
176195 return import_job , None , None
177196
197+ try :
198+ mal_prediction_import = Entity .MEAToMALPredictionImport .create_for_model_run_data_rows (
199+ data_row_ids = mea_to_mal_data_rows ,
200+ project_id = project_id ,
201+ ** kwargs )
202+ mal_prediction_import .wait_until_done ()
203+ except Exception as e :
204+ logger .warning (
205+ f"Failed to create MEA to MAL prediction import. Message : { e }
206+ )
207+ # Unable to create mea to mal prediction import
208+ return import_job , batch , None
209+
210+ return import_job , batch , mal_prediction_import
211+
178212 def add_predictions (
179213 self ,
180214 name : str ,
0 commit comments