Save trained models, checkpoints and tensorboard logs to Google Cloud Storage

[khanhphan1311]

Hi @Toni-SM,

Thank you for amazing works.

Currently, I train skrl agent using custom training on Google Vertex AI (i.e. google.cloud.aiplatform).
The save config with skrl is set with agent:

cfg["experiment"]["directory"] = f"{experiment_dir}"

And in the Agent class init:

        # experiment directory
        directory = self.cfg.get("experiment", {}).get("directory", "")
        experiment_name = self.cfg.get("experiment", {}).get("experiment_name", "")
        if not directory:
            directory = os.path.join(os.getcwd(), "runs")
        if not experiment_name:
            experiment_name = "{}_{}".format(datetime.datetime.now().strftime("%y-%m-%d_%H-%M-%S-%f"), self.__class__.__name__)
        self.experiment_dir = os.path.join(directory, experiment_name)

But this does not work in Google Cloud.
Documentation for saving model in aiplatform here: https://cloud.google.com/vertex-ai/docs/training/code-requirements#environment-variables
There will be 3 variables:
AIP_MODEL_DIR: a Cloud Storage URI of a directory intended for [saving model artifacts].
AIP_CHECKPOINT_DIR: a Cloud Storage URI of a directory intended for [saving checkpoints].
AIP_TENSORBOARD_LOG_DIR: a Cloud Storage URI of a directory intended for saving [TensorBoard].

I have 2 questions:

1. How to save models and checkpoints to GCS when training?
2. How to load tensorboard logs to Tensorboard instances in aiplatform to track experiment? It works on local machine but in cloud, I have no idea to track tb logs of training.

Currently, I create another method to save all files in agent.experiment_dir to GCS after done training:

    trainer = SequentialTrainer(cfg=cfg_trainer, env=env_train, agents=[agent])

    # Start training
    trainer.train()

    # Store model to GCS
    gcs_util = GCSUtility()
    model_dir = Path(agent.experiment_dir)
    gcs_util.upload_folder_to_bucket(args.gcs_bucket, model_dir)

where the method here:

    def upload_folder_to_bucket(self, bucket_name, source_folder_name):
        bucket = self.storage_client.bucket(bucket_name)

        for dir_path, dir_names, file_names in os.walk(source_folder_name):
            for file_name in file_names:
                # construct the full file path
                file_path = os.path.join(dir_path, file_name)

                # ensure the path is using the correct path separator
                blob_path = file_path.replace("\\", "/")

                blob = bucket.blob(blob_path)
                with open(file_path, 'rb') as my_file:
                    blob.upload_from_file(my_file)

This is work-around solution for Question 1.
But for Question 2 with Tensorboard, I don't have any ideas.

[Toni-SM]

Hi @khanhphan1311

Since the Vertex AI environment variables for CustomJob are configured as follow (reference: baseOutputDirectory):

AIP_MODEL_DIR = <baseOutputDirectory>/model/
AIP_CHECKPOINT_DIR = <baseOutputDirectory>/checkpoints/
AIP_TENSORBOARD_LOG_DIR = <baseOutputDirectory>/logs/

The next small changes can be applied to the skrl code to work with Google Cloud Storage.
These changes overwrite the configuration with Vertex AI environment variables.

• Replace:

  skrl/skrl/agents/torch/base.py

  Line 86 in 00a2fd3

  self.experiment_dir = os.path.join(directory, experiment_name)

  with

  self.experiment_dir = os.path.join(directory, experiment_name)
  if "AIP_CHECKPOINT_DIR" in os.environ:
      self.experiment_dir = os.path.join(*os.path.split(os.environ["AIP_CHECKPOINT_DIR"])[:-1])

• Replace:

  skrl/skrl/agents/torch/base.py

  Line 158 in 00a2fd3

  self.writer = SummaryWriter(log_dir=self.experiment_dir)

  with

  self.writer = SummaryWriter(log_dir=os.path.join(self.experiment_dir, "logs" if "AIP_TENSORBOARD_LOG_DIR" in os.environ else ""))

To view Vertex AI TensorBoard in the Google Cloud console follow the next link: https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-training#view_your_in_the