From cd01fe056453915a5065fc3df3f90e035cfd79ca Mon Sep 17 00:00:00 2001 From: vinicvaz Date: Mon, 12 Feb 2024 11:23:47 -0300 Subject: [PATCH] update var name and add return storage usage to ops --- src/domino/base_piece.py | 8 ++++---- src/domino/custom_operators/docker_operator.py | 4 +++- src/domino/custom_operators/k8s_operator.py | 1 + 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/src/domino/base_piece.py b/src/domino/base_piece.py index 68732fb3..55f68b88 100644 --- a/src/domino/base_piece.py +++ b/src/domino/base_piece.py @@ -60,7 +60,7 @@ def __init__( self.logger = get_configured_logger(f"{self.__class__.__name__ }-{self.task_id}") self.display_result = None - self._storage_size_in_bytes = 0 + self._shared_storage_usage_in_bytes = 0 def start_logger(self): """ @@ -302,10 +302,10 @@ def run_piece_function( # Push XCom xcom_obj = self.format_xcom(output_obj=output_obj) shared_storage_base_path = f"{self.workflow_shared_storage_path}/{self.task_id}" - self._storage_size_in_bytes = self._get_folder_size(shared_storage_base_path) - xcom_obj['_storage_size_in_bytes'] = self._storage_size_in_bytes + self._shared_storage_usage_in_bytes = self._get_folder_size(shared_storage_base_path) + xcom_obj['_storage_size_in_bytes'] = self._shared_storage_usage_in_bytes self.push_xcom(xcom_obj=xcom_obj) - self.logger.info(f"Piece used {self._storage_size_in_bytes} bytes of storage.") + self.logger.info(f"Piece used {self._shared_storage_usage_in_bytes} bytes of storage.") self.logger.info("End cut point for logger 48c94577-0225-4c3f-87c0-8add3f4e6d4b") diff --git a/src/domino/custom_operators/docker_operator.py b/src/domino/custom_operators/docker_operator.py index 4b538547..17548760 100644 --- a/src/domino/custom_operators/docker_operator.py +++ b/src/domino/custom_operators/docker_operator.py @@ -190,4 +190,6 @@ def execute(self, context: Context) -> Optional[str]: self.domino_client = DominoBackendRestClient(base_url="http://domino-rest:8000/") # env var format = {"name": "value"} self._prepare_execute_environment(context=context) - return super().execute(context=context) + result = super().execute(context=context) + self._shared_storage_usage_in_bytes = result['_shared_storage_usage_in_bytes'] + return result diff --git a/src/domino/custom_operators/k8s_operator.py b/src/domino/custom_operators/k8s_operator.py index f48e248c..128fe4e8 100644 --- a/src/domino/custom_operators/k8s_operator.py +++ b/src/domino/custom_operators/k8s_operator.py @@ -510,6 +510,7 @@ def execute(self, context: Context): ti.xcom_push(key='pod_name', value=self.pod.metadata.name) ti.xcom_push(key='pod_namespace', value=self.pod.metadata.namespace) if self.do_xcom_push: + self._shared_storage_usage_in_bytes = result['_shared_storage_usage_in_bytes'] return result