Google Vertex AI logging not happening for managed Jupyter notebook execution

Question

I have a python code inside google cloud managed Jupyter notebook. I am running this notebook by using using the 'Execute' button in notebook. This creates a job which can be seen in vertex-ai/training/custom-jobs. But the logs that I have in python code are not getting printed in the job logs. I have the Logs Writer permission for the notebook service account.

I am also getting this error in logs: ModuleNotFoundError: No module named 'bq_stats'

Below is the code that I have tried. I have also tried the code suggested in below link but it didn't work (dint print the log): https://cloud.google.com/logging/docs/setup/python

!pip install google-cloud-logging

import logging
from google.cloud import logging as gcp_logging
from google.cloud.logging.handlers import CloudLoggingHandler

# Initialize the Google Cloud Logging client
client = gcp_logging.Client()

# Create a CloudLoggingHandler for sending logs to Google Cloud
cloud_logger = CloudLoggingHandler(client)

# Get the root logger
logger = logging.getLogger() 

# Set the logging level to INFO (or another preferred level)
logger.setLevel(logging.INFO)

# Add the CloudLoggingHandler to the logger
logger.addHandler(cloud_logger)

# Additionally, create a StreamHandler for logging to stdout
stream_handler = logging.StreamHandler()
logger.addHandler(stream_handler)

# Log a message, which should now appear both in Google Cloud Logging and stdout
logger.info("This should be logged")

# Ensure all logs are flushed and sent to Google Cloud Logging
cloud_logger.flush()