airflow.providers.google.cloud.hooks.ray

This module contains a Google Cloud Ray Job hook.

Attributes

VERTEX_RAY_DOMAIN

Classes

RayJobHook

Hook for Jobs APIs.

Module Contents

airflow.providers.google.cloud.hooks.ray.VERTEX_RAY_DOMAIN = 'aiplatform-training.googleusercontent.com'[source]
class airflow.providers.google.cloud.hooks.ray.RayJobHook(gcp_conn_id='google_cloud_default', impersonation_chain=None, **kwargs)[source]

Bases: airflow.providers.google.common.hooks.base_google.GoogleBaseHook

Hook for Jobs APIs.

get_client(address)[source]

Create a client for submitting and interacting with jobs on a remote cluster.

Parameters:

address (str) – Either (1) the address of the Ray cluster, or (2) the HTTP address of the dashboard server on the head node, e.g. “http://<head-node-ip>:8265”. In case (1) it must be specified as an address that can be passed to ray.init(), e.g. a Ray Client address (ray://<head_node_host>:10001), or “auto”, or “localhost:<port>”.

serialize_job_obj(job_obj)[source]

Serialize JobDetails to a plain dict.

submit_job(entrypoint, cluster_address, runtime_env=None, metadata=None, submission_id=None, entrypoint_num_cpus=None, entrypoint_num_gpus=None, entrypoint_memory=None, entrypoint_resources=None)[source]

Submit and execute Job on Ray cluster.

When a job is submitted, it runs once to completion or failure. Retries or different runs with different parameters should be handled by the submitter. Jobs are bound to the lifetime of a Ray cluster, so if the cluster goes down, all running jobs on that cluster will be terminated.

Parameters:

  • entrypoint (str) – Required. The shell command to run for this job.

  • cluster_address (str) – Required. Either (1) the address of the Ray cluster, or (2) the HTTP address of the dashboard server on the head node, e.g. “http://<head-node-ip>:8265”. In case (1) it must be specified as an address that can be passed to ray.init(), e.g. a Ray Client address (ray://<head_node_host>:10001), or “auto”, or “localhost:<port>”.

  • submission_id (str | None) – A unique ID for this job.

  • runtime_env (dict[str, Any] | None) – The runtime environment to install and run this job in.

  • metadata (dict[str, str] | None) – Arbitrary data to store along with this job.

  • entrypoint_num_cpus (int | float | None) – The quantity of CPU cores to reserve for the execution of the entrypoint command, separately from any tasks or actors launched by it. Defaults to 0.

  • entrypoint_num_gpus (int | float | None) – The quantity of GPUs to reserve for the execution of the entrypoint command, separately from any tasks or actors launched by it. Defaults to 0.

  • entrypoint_memory (int | None) – The quantity of memory to reserve for the execution of the entrypoint command, separately from any tasks or actors launched by it. Defaults to 0.

  • entrypoint_resources (dict[str, float] | None) – The quantity of custom resources to reserve for the execution of the entrypoint command, separately from any tasks or actors launched by it.

stop_job(job_id, cluster_address)[source]

Stop Job on Ray cluster.

Parameters:

  • job_id (str) – Required. The job ID or submission ID for the job to be stopped.

  • cluster_address (str) – Required. Either (1) the address of the Ray cluster, or (2) the HTTP address of the dashboard server on the head node, e.g. “http://<head-node-ip>:8265”. In case (1) it must be specified as an address that can be passed to ray.init(), e.g. a Ray Client address (ray://<head_node_host>:10001), or “auto”, or “localhost:<port>”.

Returns:

True if the job was stopped, otherwise False.

Return type:

bool

delete_job(job_id, cluster_address)[source]

Delete Job on Ray cluster in a terminal state and all of its associated data.

If the job is not already in a terminal state, raises an error. This does not delete the job logs from disk. Submitting a job with the same submission ID as a previously deleted job is not supported and may lead to unexpected behavior.

Parameters:

  • job_id (str) – Required. The job ID or submission ID for the job to be deleted.

  • cluster_address (str) – Required. Either (1) the address of the Ray cluster, or (2) the HTTP address of the dashboard server on the head node, e.g. “http://<head-node-ip>:8265”. In case (1) it must be specified as an address that can be passed to ray.init(), e.g. a Ray Client address (ray://<head_node_host>:10001), or “auto”, or “localhost:<port>”.

Returns:

True if the job was deleted, otherwise False.

Return type:

bool

get_job_info(job_id, cluster_address)[source]

Get the latest status and other information associated with a Job on Ray cluster.

Parameters:

  • job_id (str) – Required. The job ID or submission ID for the job to be retrieved.

  • cluster_address (str) – Required. Either (1) the address of the Ray cluster, or (2) the HTTP address of the dashboard server on the head node, e.g. “http://<head-node-ip>:8265”. In case (1) it must be specified as an address that can be passed to ray.init(), e.g. a Ray Client address (ray://<head_node_host>:10001), or “auto”, or “localhost:<port>”.

Returns:

The JobDetails for the job.

Return type:

ray.dashboard.modules.job.pydantic_models.JobDetails

list_jobs(cluster_address)[source]

List all jobs along with their status and other information.

Lists all jobs that have ever run on the cluster, including jobs that are currently running and jobs that are no longer running.

Parameters:

cluster_address (str) – Required. Either (1) the address of the Ray cluster, or (2) the HTTP address of the dashboard server on the head node, e.g. “http://<head-node-ip>:8265”. In case (1) it must be specified as an address that can be passed to ray.init(), e.g. a Ray Client address (ray://<head_node_host>:10001), or “auto”, or “localhost:<port>”.

get_job_status(job_id, cluster_address)[source]

Get the most recent status of a Job on Ray cluster.

Parameters:

  • job_id (str) – Required. The job ID or submission ID for the job to be retrieved.

  • cluster_address (str) – Required. Either (1) the address of the Ray cluster, or (2) the HTTP address of the dashboard server on the head node, e.g. “http://<head-node-ip>:8265”. In case (1) it must be specified as an address that can be passed to ray.init(), e.g. a Ray Client address (ray://<head_node_host>:10001), or “auto”, or “localhost:<port>”.

Returns:

The JobStatus of the job.

Return type:

ray.dashboard.modules.job.common.JobStatus

get_job_logs(job_id, cluster_address)[source]

Get all logs produced by a Job on Ray cluster.

Parameters:

  • job_id (str) – Required. The job ID or submission ID for the job to be retrieved.

  • cluster_address (str) – Required. Either (1) the address of the Ray cluster, or (2) the HTTP address of the dashboard server on the head node, e.g. “http://<head-node-ip>:8265”. In case (1) it must be specified as an address that can be passed to ray.init(), e.g. a Ray Client address (ray://<head_node_host>:10001), or “auto”, or “localhost:<port>”.

Returns:

A string containing the full logs of the job.

Return type:

str

Was this entry helpful?