Skip to content

Commit e2aa629

Browse files
SK-2777: fix sdk metrics
1 parent 75949f0 commit e2aa629

3 files changed

Lines changed: 33 additions & 30 deletions

File tree

‎skyflow/utils/_utils.py‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -177,8 +177,12 @@ def get_data_from_content_type(data, content_type):
177177
return converted_data, files
178178

179179

180+
_CACHED_METRICS: dict = {}
181+
180182
def get_metrics():
181-
sdk_name_version = "skyflow-python@" + SDK_VERSION
183+
global _CACHED_METRICS
184+
if _CACHED_METRICS:
185+
return _CACHED_METRICS
182186

183187
try:
184188
sdk_client_device_model = platform.node()
@@ -195,13 +199,13 @@ def get_metrics():
195199
except Exception:
196200
sdk_runtime_details = ""
197201

198-
details_dic = {
199-
'sdk_name_version': sdk_name_version,
202+
_CACHED_METRICS = {
203+
'sdk_name_version': "skyflow-python@" + SDK_VERSION,
200204
'sdk_client_device_model': sdk_client_device_model,
201205
'sdk_client_os_details': sdk_client_os_details,
202206
'sdk_runtime_details': "Python " + sdk_runtime_details,
203207
}
204-
return details_dic
208+
return _CACHED_METRICS
205209

206210
def parse_insert_response(api_response, continue_on_error):
207211
# Retrieve the headers and data from the API response

‎skyflow/vault/controller/_detect.py‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ def __poll_for_processed_file(self, run_id, max_wait_time=64):
6262
current_wait_time = 1 # Start with 1 second
6363
try:
6464
while True:
65-
response = files_api.get_run(run_id, vault_id=self.__vault_client.get_vault_id(), request_options=self.__get_headers()).data
65+
response = files_api.get_run(run_id, vault_id=self.__vault_client.get_vault_id(), request_options={'additional_headers': self.__get_headers()}).data
6666
status = response.status
6767
if status == 'IN_PROGRESS':
6868
if current_wait_time >= max_wait_time:
@@ -228,7 +228,7 @@ def deidentify_text(self, request: DeidentifyTextRequest) -> DeidentifyTextRespo
228228
restrict_regex=deidentify_text_body['restrict_regex'],
229229
token_type=deidentify_text_body['token_type'],
230230
transformations=deidentify_text_body['transformations'],
231-
request_options=self.__get_headers()
231+
request_options={'additional_headers': self.__get_headers()}
232232
)
233233
deidentify_text_response = parse_deidentify_text_response(api_response)
234234
log_info(SkyflowMessages.Info.DEIDENTIFY_TEXT_SUCCESS.value, self.__vault_client.get_logger())
@@ -252,7 +252,7 @@ def reidentify_text(self, request: ReidentifyTextRequest) -> ReidentifyTextRespo
252252
vault_id=self.__vault_client.get_vault_id(),
253253
text=reidentify_text_body['text'],
254254
format=reidentify_text_body['format'],
255-
request_options=self.__get_headers()
255+
request_options={'additional_headers': self.__get_headers()}
256256
)
257257
reidentify_text_response = parse_reidentify_text_response(api_response)
258258
log_info(SkyflowMessages.Info.REIDENTIFY_TEXT_SUCCESS.value, self.__vault_client.get_logger())
@@ -296,7 +296,7 @@ def deidentify_file(self, request: DeidentifyFileRequest):
296296
'allow_regex': request.allow_regex_list,
297297
'restrict_regex': request.restrict_regex_list,
298298
'transformations': self.__get_transformations(request),
299-
'request_options': self.__get_headers()
299+
'request_options': {'additional_headers': self.__get_headers()}
300300
}
301301

302302
elif file_extension in ['mp3', 'wav']:
@@ -316,7 +316,7 @@ def deidentify_file(self, request: DeidentifyFileRequest):
316316
'bleep_frequency': getattr(request, 'bleep', None).frequency if getattr(request, 'bleep', None) is not None else None,
317317
'bleep_start_padding': getattr(request, 'bleep', None).start_padding if getattr(request, 'bleep', None) is not None else None,
318318
'bleep_stop_padding': getattr(request, 'bleep', None).stop_padding if getattr(request, 'bleep', None) is not None else None,
319-
'request_options': self.__get_headers()
319+
'request_options': {'additional_headers': self.__get_headers()}
320320
}
321321

322322
elif file_extension == 'pdf':
@@ -331,7 +331,7 @@ def deidentify_file(self, request: DeidentifyFileRequest):
331331
'restrict_regex': request.restrict_regex_list,
332332
'max_resolution': getattr(request, 'max_resolution', None),
333333
'density': getattr(request, 'pixel_density', None),
334-
'request_options': self.__get_headers()
334+
'request_options': {'additional_headers': self.__get_headers()}
335335
}
336336

337337
elif file_extension in ['jpeg', 'jpg', 'png', 'bmp', 'tif', 'tiff']:
@@ -347,7 +347,7 @@ def deidentify_file(self, request: DeidentifyFileRequest):
347347
'masking_method': getattr(request, 'masking_method', None),
348348
'output_ocr_text': getattr(request, 'output_ocr_text', None),
349349
'output_processed_image': getattr(request, 'output_processed_image', None),
350-
'request_options': self.__get_headers()
350+
'request_options': {'additional_headers': self.__get_headers()}
351351
}
352352

353353
elif file_extension in ['ppt', 'pptx']:
@@ -360,7 +360,7 @@ def deidentify_file(self, request: DeidentifyFileRequest):
360360
'token_type': self.__get_token_format(request),
361361
'allow_regex': request.allow_regex_list,
362362
'restrict_regex': request.restrict_regex_list,
363-
'request_options': self.__get_headers()
363+
'request_options': {'additional_headers': self.__get_headers()}
364364
}
365365

366366
elif file_extension in ['csv', 'xls', 'xlsx']:
@@ -373,7 +373,7 @@ def deidentify_file(self, request: DeidentifyFileRequest):
373373
'token_type': self.__get_token_format(request),
374374
'allow_regex': request.allow_regex_list,
375375
'restrict_regex': request.restrict_regex_list,
376-
'request_options': self.__get_headers()
376+
'request_options': {'additional_headers': self.__get_headers()}
377377
}
378378

379379
elif file_extension in ['doc', 'docx']:
@@ -386,7 +386,7 @@ def deidentify_file(self, request: DeidentifyFileRequest):
386386
'token_type': self.__get_token_format(request),
387387
'allow_regex': request.allow_regex_list,
388388
'restrict_regex': request.restrict_regex_list,
389-
'request_options': self.__get_headers()
389+
'request_options': {'additional_headers': self.__get_headers()}
390390
}
391391

392392
elif file_extension in ['json', 'xml']:
@@ -400,7 +400,7 @@ def deidentify_file(self, request: DeidentifyFileRequest):
400400
'allow_regex': request.allow_regex_list,
401401
'restrict_regex': request.restrict_regex_list,
402402
'transformations': self.__get_transformations(request),
403-
'request_options': self.__get_headers()
403+
'request_options': {'additional_headers': self.__get_headers()}
404404
}
405405

406406
else:
@@ -414,7 +414,7 @@ def deidentify_file(self, request: DeidentifyFileRequest):
414414
'allow_regex': request.allow_regex_list,
415415
'restrict_regex': request.restrict_regex_list,
416416
'transformations': self.__get_transformations(request),
417-
'request_options': self.__get_headers()
417+
'request_options': {'additional_headers': self.__get_headers()}
418418
}
419419

420420
log_info(SkyflowMessages.Info.DETECT_FILE_REQUEST_RESOLVED.value, self.__vault_client.get_logger())
@@ -448,7 +448,7 @@ def get_detect_run(self, request: GetDetectRunRequest):
448448
response = files_api.get_run(
449449
run_id,
450450
vault_id=self.__vault_client.get_vault_id(),
451-
request_options=self.__get_headers()
451+
request_options={'additional_headers': self.__get_headers()}
452452
)
453453
if response.data.status == 'IN_PROGRESS':
454454
parsed_response = self.__parse_deidentify_file_response(DeidentifyFileResponse(run_id=run_id, status='IN_PROGRESS'))

‎skyflow/vault/controller/_vault.py‎

Lines changed: 12 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -89,10 +89,9 @@ def __get_file_for_file_upload(self, request: FileUploadRequest) -> Optional[Fil
8989
return None
9090

9191
def __get_headers(self):
92-
headers = {
93-
SKY_META_DATA_HEADER: json.dumps(get_metrics())
94-
}
95-
return headers
92+
if not hasattr(self, '_cached_headers'):
93+
self._cached_headers = {SKY_META_DATA_HEADER: json.dumps(get_metrics())}
94+
return self._cached_headers
9695

9796
def insert(self, request: InsertRequest):
9897
log_info(SkyflowMessages.Info.VALIDATE_INSERT_REQUEST.value, self.__vault_client.get_logger())
@@ -106,11 +105,11 @@ def insert(self, request: InsertRequest):
106105
log_info(SkyflowMessages.Info.INSERT_TRIGGERED.value, self.__vault_client.get_logger())
107106
if request.continue_on_error:
108107
api_response = records_api.record_service_batch_operation(self.__vault_client.get_vault_id(),
109-
records=insert_body, continue_on_error=request.continue_on_error, byot=request.token_mode.value, request_options=self.__get_headers())
108+
records=insert_body, continue_on_error=request.continue_on_error, byot=request.token_mode.value, request_options={'additional_headers': self.__get_headers()})
110109

111110
else:
112111
api_response = records_api.record_service_insert_record(self.__vault_client.get_vault_id(),
113-
request.table, records=insert_body,tokenization= request.return_tokens, upsert=request.upsert, homogeneous=request.homogeneous, byot=request.token_mode.value, request_options=self.__get_headers())
112+
request.table, records=insert_body,tokenization= request.return_tokens, upsert=request.upsert, homogeneous=request.homogeneous, byot=request.token_mode.value, request_options={'additional_headers': self.__get_headers()})
114113

115114
insert_response = parse_insert_response(api_response, request.continue_on_error)
116115
log_info(SkyflowMessages.Info.INSERT_SUCCESS.value, self.__vault_client.get_logger())
@@ -138,7 +137,7 @@ def update(self, request: UpdateRequest):
138137
record=record,
139138
tokenization=request.return_tokens,
140139
byot=request.token_mode.value,
141-
request_options = self.__get_headers()
140+
request_options={'additional_headers': self.__get_headers()}
142141
)
143142
log_info(SkyflowMessages.Info.UPDATE_SUCCESS.value, self.__vault_client.get_logger())
144143
update_response = parse_update_record_response(api_response)
@@ -159,7 +158,7 @@ def delete(self, request: DeleteRequest):
159158
self.__vault_client.get_vault_id(),
160159
request.table,
161160
skyflow_ids=request.ids,
162-
request_options=self.__get_headers()
161+
request_options={'additional_headers': self.__get_headers()}
163162
)
164163
log_info(SkyflowMessages.Info.DELETE_SUCCESS.value, self.__vault_client.get_logger())
165164
delete_response = parse_delete_response(api_response)
@@ -189,7 +188,7 @@ def get(self, request: GetRequest):
189188
download_url=request.download_url,
190189
column_name=request.column_name,
191190
column_values=request.column_values,
192-
request_options=self.__get_headers()
191+
request_options={'additional_headers': self.__get_headers()}
193192
)
194193
log_info(SkyflowMessages.Info.GET_SUCCESS.value, self.__vault_client.get_logger())
195194
get_response = parse_get_response(api_response)
@@ -209,7 +208,7 @@ def query(self, request: QueryRequest):
209208
api_response = query_api.query_service_execute_query(
210209
self.__vault_client.get_vault_id(),
211210
query=request.query,
212-
request_options=self.__get_headers()
211+
request_options={'additional_headers': self.__get_headers()}
213212
)
214213
log_info(SkyflowMessages.Info.QUERY_SUCCESS.value, self.__vault_client.get_logger())
215214
query_response = parse_query_response(api_response)
@@ -237,7 +236,7 @@ def detokenize(self, request: DetokenizeRequest):
237236
self.__vault_client.get_vault_id(),
238237
detokenization_parameters=tokens_list,
239238
continue_on_error = request.continue_on_error,
240-
request_options=self.__get_headers()
239+
request_options={'additional_headers': self.__get_headers()}
241240
)
242241
log_info(SkyflowMessages.Info.DETOKENIZE_SUCCESS.value, self.__vault_client.get_logger())
243242
detokenize_response = parse_detokenize_response(api_response)
@@ -262,7 +261,7 @@ def tokenize(self, request: TokenizeRequest):
262261
api_response = tokens_api.record_service_tokenize(
263262
self.__vault_client.get_vault_id(),
264263
tokenization_parameters=records_list,
265-
request_options=self.__get_headers()
264+
request_options={'additional_headers': self.__get_headers()}
266265
)
267266
tokenize_response = parse_tokenize_response(api_response)
268267
log_info(SkyflowMessages.Info.TOKENIZE_SUCCESS.value, self.__vault_client.get_logger())
@@ -285,7 +284,7 @@ def upload_file(self, request: FileUploadRequest):
285284
file=self.__get_file_for_file_upload(request),
286285
skyflow_id=request.skyflow_id,
287286
return_file_metadata= False,
288-
request_options=self.__get_headers()
287+
request_options={'additional_headers': self.__get_headers()}
289288
)
290289
log_info(SkyflowMessages.Info.FILE_UPLOAD_REQUEST_RESOLVED.value, self.__vault_client.get_logger())
291290
log_info(SkyflowMessages.Info.FILE_UPLOAD_SUCCESS.value, self.__vault_client.get_logger())

0 commit comments

Comments
 (0)