unable to specify master_type in MLEngineTrainingOperator

问题

I am using airflow to schedule a pipeline that will result in training a scikitlearn model with ai platform. I use this DAG to train it

    with models.DAG(JOB_NAME,
                    schedule_interval=None,
                    default_args=default_args) as dag:

        # Tasks definition
        training_op = MLEngineTrainingOperator(
            task_id='submit_job_for_training',
            project_id=PROJECT,
            job_id=job_id,
            package_uris=[os.path.join(TRAINER_BIN)],
            training_python_module=TRAINER_MODULE,
            runtime_version=RUNTIME_VERSION,
            region='europe-west1',
            training_args=[
                '--base-dir={}'.format(BASE_DIR),
                '--event-date=20200212',
            ],
            python_version='3.5')
        training_op

The training package loads the desired csv files and train a RandomForestClassifier on it.

This works fine until the number and the size of the files increase. Then I get this error:

ERROR - The replica master 0 ran out-of-memory and exited with a non-zero status of 9(SIGKILL). To find out more about why your job exited please check the logs:

The total size of the files is around 4 Gb. I dont know what is the default machine used but is seems not enough. Hoping this would solve the memory consumption issue I tried to change the parameter n_jobs of the classifier from -1 to 1, with no more luck.

Looking at the code of MLEngineTrainingOperator and the documentation I added a custom scale_tier and a master_type n1-highmem-8, 8 CPUs and 52GB of RAM , like this:

with models.DAG(JOB_NAME,
                schedule_interval=None,
                default_args=default_args) as dag:

    # Tasks definition
    training_op = MLEngineTrainingOperator(
        task_id='submit_job_for_training',
        project_id=PROJECT,
        job_id=job_id,
        package_uris=[os.path.join(TRAINER_BIN)],
        training_python_module=TRAINER_MODULE,
        runtime_version=RUNTIME_VERSION,
        region='europe-west1',
        master_type="n1-highmem-8",
        scale_tier="custom",
        training_args=[
            '--base-dir={}'.format(BASE_DIR),
            '--event-date=20200116',
        ],
        python_version='3.5')
    training_op

This resulted in an other error:

ERROR - <HttpError 400 when requesting https://ml.googleapis.com/v1/projects/MY_PROJECT/jobs?alt=json returned "Field: master_type Error: Master type must be specified for the CUSTOM scale tier.">

I don't know what is wrong but it appears that is not the way to do that.

EDIT: Using command line I manage to launch the job:

gcloud ai-platform jobs submit training training_job_name --packages=gs://path/to/package/package.tar.gz --python-version=3.5 --region=europe-west1 --runtime-version=1.14 --module-name=trainer.train --scale-tier=CUSTOM --master-machine-type=n1-highmem-16

However i would like to do this in airflow.

Any help would be much appreciated.

EDIT: My environment used an old version of apache airflow, 1.10.3 where the master_type argument was not present. Updating the version to 1.10.6 solved this issue

回答1:

My environment used an old version of apache airflow, 1.10.3 where the master_type argument was not present. Updating the version to 1.10.6 solved this issue

来源：https://stackoverflow.com/questions/60194310/unable-to-specify-master-type-in-mlenginetrainingoperator

标签

out-of-memory

airflow

google-cloud-ml