# Path: ml_pipeline_gen/parsers.py
# def parse_yaml(path):
#     """Parses the given config file."""
#     with open(path, "r") as f:
#         doc = f.read()
#     return yaml.load(doc, Loader=yaml.FullLoader)


import abc
import collections.abc
import datetime as dt
import json
import os
import pathlib
import shutil
import subprocess
import time
import jinja2 as jinja
import tensorflow.compat.v1 as tf
from googleapiclient import discovery
from googleapiclient import errors
from tensorflow.python.tools import saved_model_utils
from ml_pipeline_gen.parsers import parse_yaml
            tune: train with hyperparameter tuning if true.
            blocking: true if the function should exit only once the job
                completes.
            wait_interval: if blocking, how often the job state should be
                checked.

        Returns:
            job_id: a CAIP job id.
        """
        now = dt.datetime.now().strftime("%Y%m%d_%H%M%S")
        job_id = "train_{}_{}".format(self.model["name"], now)
        package_uri = self.upload_trainer_dist()
        jobs_client = self.ml_client.projects().jobs()
        body = {
            "jobId": job_id,
            "trainingInput": {
                "scaleTier": self.scale_tier,
                "packageUris": [package_uri],
                "pythonModule": "trainer.task",
                "args": [
                    "--model_dir", self.get_model_dir(),
                ],
                "jobDir": self.get_job_dir(),
                "region": self.region,
                "runtimeVersion": self.runtime_version,
                "pythonVersion": self.python_version,
            },
        }
        # TODO(smhosein): should we handle custom scale tiers?
        if tune:

# Path: ml_pipeline_gen/parsers.py
# def parse_yaml(path):
#     """Parses the given config file."""
#     with open(path, "r") as f:
#         doc = f.read()
#     return yaml.load(doc, Loader=yaml.FullLoader)


# Path: ml_pipeline_gen/models.py
import abc
import collections.abc
import datetime as dt
import json
import os
import pathlib
import shutil
import subprocess
import time
import jinja2 as jinja
import tensorflow.compat.v1 as tf
from googleapiclient import discovery
from googleapiclient import errors
from tensorflow.python.tools import saved_model_utils
from ml_pipeline_gen.parsers import parse_yaml

            tune: train with hyperparameter tuning if true.
            blocking: true if the function should exit only once the job
                completes.
            wait_interval: if blocking, how often the job state should be
                checked.

        Returns:
            job_id: a CAIP job id.
        """
        now = dt.datetime.now().strftime("%Y%m%d_%H%M%S")
        job_id = "train_{}_{}".format(self.model["name"], now)
        package_uri = self.upload_trainer_dist()
        jobs_client = self.ml_client.projects().jobs()
        body = {
            "jobId": job_id,
            "trainingInput": {
                "scaleTier": self.scale_tier,
                "packageUris": [package_uri],
                "pythonModule": "trainer.task",
                "args": [
                    "--model_dir", self.get_model_dir(),
                ],
                "jobDir": self.get_job_dir(),
                "region": self.region,
                "runtimeVersion": self.runtime_version,
                "pythonVersion": self.python_version,
            },
        }
        # TODO(smhosein): should we handle custom scale tiers?
        if tune:

            hp_config = parse_yaml(self.model_params["hyperparam_config"])