airflow.contrib.hooks.gcp_mlengine_hook

Module Contents

airflow.contrib.hooks.gcp_mlengine_hook._poll_with_exponential_delay(request, max_n, is_done_func, is_error_func)[source]
class airflow.contrib.hooks.gcp_mlengine_hook.MLEngineHook(gcp_conn_id='google_cloud_default', delegate_to=None)[source]

Bases: airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook

get_conn(self)[source]

Returns a Google MLEngine service object.

create_job(self, project_id, job, use_existing_job_fn=None)[source]

Launches a MLEngine job and wait for it to reach a terminal state.

Parameters
  • project_id (str) – The Google Cloud project id within which MLEngine job will be launched.

  • job (dict) –

    MLEngine Job object that should be provided to the MLEngine API, such as:

    {
      'jobId': 'my_job_id',
      'trainingInput': {
        'scaleTier': 'STANDARD_1',
        ...
      }
    }
    

  • use_existing_job_fn (function) – In case that a MLEngine job with the same job_id already exist, this method (if provided) will decide whether we should use this existing job, continue waiting for it to finish and returning the job object. It should accepts a MLEngine job object, and returns a boolean value indicating whether it is OK to reuse the existing job. If ‘use_existing_job_fn’ is not provided, we by default reuse the existing MLEngine job.

Returns

The MLEngine job object if the job successfully reach a terminal state (which might be FAILED or CANCELLED state).

Return type

dict

_get_job(self, project_id, job_id)[source]

Gets a MLEngine job based on the job name.

Returns

MLEngine job object if succeed.

Return type

dict

Raises:

googleapiclient.errors.HttpError: if HTTP error is returned from server

_wait_for_job_done(self, project_id, job_id, interval=30)[source]

Waits for the Job to reach a terminal state.

This method will periodically check the job state until the job reach a terminal state.

Raises:

googleapiclient.errors.HttpError: if HTTP error is returned when getting the job

create_version(self, project_id, model_name, version_spec)[source]

Creates the Version on Google Cloud ML Engine.

Returns the operation if the version was created successfully and raises an error otherwise.

set_default_version(self, project_id, model_name, version_name)[source]

Sets a version to be the default. Blocks until finished.

list_versions(self, project_id, model_name)[source]

Lists all available versions of a model. Blocks until finished.

delete_version(self, project_id, model_name, version_name)[source]

Deletes the given version of a model. Blocks until finished.

create_model(self, project_id, model)[source]

Create a Model. Blocks until finished.

get_model(self, project_id, model_name)[source]

Gets a Model. Blocks until finished.