API Reference

This page contains the auto-generated API reference for the public modules in the dataproc-ml library.

A subpackage for model inference handlers.

class google.cloud.dataproc_ml.inference.GenAiModelHandler

Bases: BaseModelHandler

A handler for running inference with Gemini models on Spark DataFrames.

This class extends BaseModelHandler to provide a convenient way to apply Google’s Gemini generative models to data in a distributed manner using Spark. It uses a builder pattern for configuration.

It automatically authenticates and discovers the project and location from the environment if not explicitly provided, making it seamless to use within Dataproc or other configured GCP environments.

Required Configuration:

  • Input specification via one of the following methods: - .prompt(str): A prompt template with one or more placeholders for

    input column/s (e.g., “Compare these two texts {col1} & {col2}”).

    • .input_cols(…) and .pre_processor(…): For more complex prompt construction logic.

Optional Configuration:

  • .project(str): Your Google Cloud project ID. If not set, it’s inferred from the environment.

  • .location(str): The GCP region for the Vertex AI API call. If not set, it’s inferred from the environment.

  • .model(str): The model to use (defaults to “gemini-2.5-flash”).

  • .output_col(str): The name of the prediction output column (defaults to “predictions”).

  • Other methods like generation_config,`max_concurrent_requests`, etc.

Example

>>> # Assumes project and location are discoverable from the environment
>>> result_df = (
...     GenAiModelHandler()
...     .prompt("What is the capital of {city} in single word?")
...     .transform(df)
... )
>>>
>>> # Explicitly setting all configurations
>>> result_df_explicit = (
...     GenAiModelHandler()
...     .project("my-gcp-project")
...     .location("us-central1")
...     .model("gemini-2.5-flash")
...     .prompt("What is the capital of {city} in single word?")
...     .output_col("predictions")
...     .generation_config(GenerationConfig(temperature=0.2))
...     .transform(df)
... )
generation_config(generation_config: GenerationConfig) → GenAiModelHandler

Sets the generation config for the model.

Parameters:

generation_config – The vertexai.generative_models.GenerationConfig object for the model.

Returns:

The handler instance for method chaining.

location(location: str) → GenAiModelHandler

Sets the Google Cloud location (region) for the Vertex AI API call.

If not provided, the location is inferred from the environment, which is useful when running on Dataproc or other GCP services.

Parameters:

location – The GCP location (e.g., “us-central1”).

Returns:

The handler instance for method chaining.

max_concurrent_requests(n: int) → GenAiModelHandler

Sets the maximum number of concurrent requests to the model API by each Python process.

Defaults to 5.

Parameters:

n – The maximum number of concurrent requests.

Returns:

The handler instance for method chaining.

model(model: str = 'gemini-2.5-flash', provider: ModelProvider = ModelProvider.GOOGLE) → GenAiModelHandler

Sets the Gemini model to be used for inference.

Parameters:

  • model – The name of the model. Defaults to “gemini-2.5-flash”.

  • provider – Provider of the model.

  • supported. (Currently only ModelProvider.GOOGLE is)

Returns:

The handler instance for method chaining.

project(project: str) → GenAiModelHandler

Sets the Google Cloud project for the Vertex AI API call.

If not provided, the project is inferred from the environment, which is useful when running on Dataproc or other GCP services.

Parameters:

project – The GCP project ID.

Returns:

The handler instance for method chaining.

prompt(prompt_template: str) → GenAiModelHandler

Configures the handler using a string template for the prompt.

This method parses a string template (e.g., “Summarize this: {text_column}”) to automatically identify the input column(s) of the dataframe and create the necessary vectorized pre-processor. It supports one or more placeholders.

Parameters:

prompt_template – A string with named placeholders, like “Capital of {city} in {country}?”.

Returns:

The handler instance for method chaining.

Raises:

ValueError – If the prompt template contains no placeholders.

retry_strategy(retry_strategy: BaseRetrying) → GenAiModelHandler

Sets a custom tenacity retry strategy for API calls.

This will override the default strategy for the selected provider.

Parameters:

retry_strategy – A tenacity retry object (e.g., tenacity.retry(stop=tenacity.stop_after_attempt(3))).

Returns:

The handler instance for method chaining.

class google.cloud.dataproc_ml.inference.PyTorchModelHandler

Bases: BaseModelHandler

A handler for running inference with PyTorch models on Spark DataFrames.

This handler supports two primary modes of operation: 1. Loading a full, saved PyTorch model object. 2. Loading a model from a state_dict (weights file), which requires

providing the model’s class architecture separately.

Required Configuration:

  • .model_path(str): The GCS path to the model file (.pt, .pth).

  • .input_cols(str, …): The column(s) in the DataFrame to use as input.

  • .set_model_architecture(…): This is required when loading from a state dict, but should not be used when loading a full model.

Optional Configuration:

  • .output_col(str): The name for the prediction output column (defaults to “predictions”).

  • .device(str): The device to run on (e.g., “cpu”, “cuda”). Defaults to “cuda” if available.

  • .pre_processor(callable): A function to preprocess input data.

  • .set_return_type(DataType): The Spark DataType of the output.

Example

Load a full model saved in GCS:

>>> result_df = (
...     PyTorchModelHandler()
...     .model_path("gs://test-bucket/test-model.pt")
...     .input_cols("input_col")
...     .output_col("prediction")
...     .transform(input_df)
... )

Load a model from a state dictionary saved in GCS:

>>> from torchvision import models
>>> model_class = models.resnet18
>>> model_kwargs = {"weights": None}
>>> result_df = (
...     PyTorchModelHandler()
...     .model_path("gs://my-bucket/resnet18_statedict.pt")
...     .set_model_architecture(model_class, **model_kwargs)
...     .input_cols("features")
...     .output_col("predictions")
...     .transform(input_df)
... )
device(device: str) → PyTorchModelHandler

Sets the device to load the PyTorch model on.

Parameters:

device – The device to use, either “cpu” or “cuda”.

Returns:

The handler instance for method chaining.

Raises:

ValueError – If the device is not “cpu” or “cuda”.

model_path(path: str) → PyTorchModelHandler

Sets the GCS path to the saved PyTorch model.

Parameters:

path – The GCS path to the model file (e.g., “gs://bucket/model.pt”).

Returns:

The handler instance for method chaining.

Raises:

ValueError – If the path does not start with “gs://”.

set_model_architecture(model_callable: Callable[[...], Module], *args, **kwargs) → PyTorchModelHandler

Sets the PyTorch model’s architecture using a class or a factory function.

This is required if you are loading a model saved as a state_dict.

Parameters:

  • model_callable – The model’s class (e.g., MyImageClassifier) or a factory function (e.g., torchvision.models.resnet18) that returns a model instance.

  • *args – Positional arguments for the model’s constructor or factory function.

  • **kwargs – Keyword arguments for the model’s constructor or factory function.

Returns:

The handler instance for method chaining.

class google.cloud.dataproc_ml.inference.TensorFlowModelHandler

Bases: BaseModelHandler

A handler for running inference with Tensorflow models on Spark DataFrames.

This handler supports two primary modes of operation: 1. Loading a full TensorFlow SavedModel. 2. Loading a model from a weights file (e.g., H5 or checkpoint), which

requires providing the model’s Keras class architecture separately.

Currently, only models returning a single output are supported. The handler casts the input to a tensor with dtype tf.float32.

Required Configuration:

  • .model_path(str): The GCS path to the model directory or file.

  • .input_cols(str, …): The column(s) in the DataFrame to use as input.

  • .set_model_architecture(…): This is required when loading from a weights file, but should not be used when loading a SavedModel.

Optional Configuration:

  • .output_col(str): The name for the prediction output column (defaults to “predictions”).

  • .pre_processor(callable): A function to preprocess input data.

  • .set_return_type(DataType): The Spark DataType of the output (defaults to ArrayType(FloatType())).

Example

Load a full model in SavedModel format:

>>> result_df = (
...     TensorFlowModelHandler()
...     .model_path("gs://test-bucket/test-model-saved-dir")
...     .input_cols("input_col")
...     .transform(input_df)
... )

Load a model from a checkpoint file:

>>> result_df = (
...     TensorFlowModelHandler()
...     .model_path("gs://test-bucket/test-model-checkpoint.h5")
...     .set_model_architecture(model_class, **model_kwargs)
...     .input_cols("input_col")
...     .transform(input_df)
... )
model_path(path: str) → TensorFlowModelHandler

Sets the GCS path to the saved TensorFlow model artifact.

Parameters:

path – The GCS path to the model directory or file.

Returns:

The handler instance for method chaining.

Raises:

ValueError – If the path is not a string or does not start with “gs://”.

set_model_architecture(model_class: Type[Model], *args, **kwargs) → TensorFlowModelHandler

Sets the TensorFlow Keras model’s architecture using its class.

This is required if you are loading a model from a weights file (e.g., a checkpoint).

Parameters:

  • model_class – The model’s class (e.g., tf.keras.applications.ResNet50).

  • *args – Positional arguments for the model’s constructor.

  • **kwargs – Keyword arguments for the model’s constructor.

Returns:

The handler instance for method chaining.

Raises:

TypeError – If model_class is not a callable.