From c01df6681e5f94d74dcc5fd85a9ea4af4b9dcea6 Mon Sep 17 00:00:00 2001 From: Yeesian Ng <ysian@google.com> Date: Tue, 5 Mar 2024 09:02:20 -0800 Subject: [PATCH] feat: Allow module_name to be specified with the base logger. PiperOrigin-RevId: 612860933 --- google/cloud/aiplatform/base.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/google/cloud/aiplatform/base.py b/google/cloud/aiplatform/base.py index 385d407e63..cb4c8ecdf4 100644 --- a/google/cloud/aiplatform/base.py +++ b/google/cloud/aiplatform/base.py @@ -92,6 +92,8 @@ def log_create_complete( cls: Type["VertexAiResourceNoun"], resource: proto.Message, variable_name: str, + *, + module_name: str = "aiplatform", ): """Logs create event is complete. @@ -102,17 +104,23 @@ def log_create_complete( Vertex AI Resource Noun class that is being created. resource (proto.Message): Vertex AI Resource proto.Message - variable_name (str): Name of variable to use for code snippet + variable_name (str): + Name of variable to use for code snippet. + module_name (str): + The module namespace under which the Vertex AI Resource Noun + is available. Defaults to `aiplatform`. """ self.info(f"{cls.__name__} created. Resource name: {resource.name}") self.info(f"To use this {cls.__name__} in another session:") - self.info(f"{variable_name} = aiplatform.{cls.__name__}('{resource.name}')") + self.info(f"{variable_name} = {module_name}.{cls.__name__}('{resource.name}')") def log_create_complete_with_getter( self, cls: Type["VertexAiResourceNoun"], resource: proto.Message, variable_name: str, + *, + module_name: str = "aiplatform", ): """Logs create event is complete. @@ -123,11 +131,16 @@ def log_create_complete_with_getter( Vertex AI Resource Noun class that is being created. resource (proto.Message): Vertex AI Resource proto.Message - variable_name (str): Name of variable to use for code snippet + variable_name (str): + Name of variable to use for code snippet. + module_name (str): + The module namespace under which the Vertex AI Resource Noun + is available. Defaults to `aiplatform`. """ self.info(f"{cls.__name__} created. Resource name: {resource.name}") self.info(f"To use this {cls.__name__} in another session:") - self.info(f"{variable_name} = aiplatform.{cls.__name__}.get('{resource.name}')") + usage_message = f"{module_name}.{cls.__name__}.get('{resource.name}')" + self.info(f"{variable_name} = {usage_message}") def log_delete_with_lro( self,