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