apache_beam.ml.inference package

A package with various modules for running inferences and predictions on models. This package contains support for popular frameworks as well as an interface for adding unsupported frameworks.

Note: on top of the frameworks captured in submodules below, Beam also has a supported TensorFlow model handler via the tfx-bsl library. See https://beam.apache.org/documentation/ml/about-ml/#tensorflow for more information on using TensorFlow in Beam.

Submodules

• apache_beam.ml.inference.base module
  • PredictionResult
  • RateLimitExceeded
  • ModelMetadata
    • ModelMetadata.model_id
    • ModelMetadata.model_name
  • RunInferenceDLQ
    • RunInferenceDLQ.failed_inferences
    • RunInferenceDLQ.failed_preprocessing
    • RunInferenceDLQ.failed_postprocessing
  • KeyModelPathMapping
    • KeyModelPathMapping.keys
    • KeyModelPathMapping.update_path
    • KeyModelPathMapping.model_id
  • ModelHandler
    • ModelHandler.load_model()
    • ModelHandler.run_inference()
    • ModelHandler.get_num_bytes()
    • ModelHandler.get_metrics_namespace()
    • ModelHandler.get_resource_hints()
    • ModelHandler.batch_elements_kwargs()
    • ModelHandler.validate_inference_args()
    • ModelHandler.update_model_path()
    • ModelHandler.update_model_paths()
    • ModelHandler.get_preprocess_fns()
    • ModelHandler.get_postprocess_fns()
    • ModelHandler.should_skip_batching()
    • ModelHandler.set_environment_vars()
    • ModelHandler.with_preprocess_fn()
    • ModelHandler.with_postprocess_fn()
    • ModelHandler.with_no_batching()
    • ModelHandler.share_model_across_processes()
    • ModelHandler.model_copies()
    • ModelHandler.override_metrics()
    • ModelHandler.should_garbage_collect_on_timeout()
  • RemoteModelHandler
    • RemoteModelHandler.create_client()
    • RemoteModelHandler.load_model()
    • RemoteModelHandler.retry_on_exception()
    • RemoteModelHandler.run_inference()
    • RemoteModelHandler.request()
  • KeyModelMapping
  • KeyedModelHandler
    • KeyedModelHandler.load_model()
    • KeyedModelHandler.run_inference()
    • KeyedModelHandler.get_num_bytes()
    • KeyedModelHandler.get_metrics_namespace()
    • KeyedModelHandler.get_resource_hints()
    • KeyedModelHandler.batch_elements_kwargs()
    • KeyedModelHandler.validate_inference_args()
    • KeyedModelHandler.update_model_paths()
    • KeyedModelHandler.update_model_path()
    • KeyedModelHandler.share_model_across_processes()
    • KeyedModelHandler.model_copies()
    • KeyedModelHandler.override_metrics()
    • KeyedModelHandler.should_garbage_collect_on_timeout()
  • MaybeKeyedModelHandler
    • MaybeKeyedModelHandler.load_model()
    • MaybeKeyedModelHandler.run_inference()
    • MaybeKeyedModelHandler.get_num_bytes()
    • MaybeKeyedModelHandler.get_metrics_namespace()
    • MaybeKeyedModelHandler.get_resource_hints()
    • MaybeKeyedModelHandler.batch_elements_kwargs()
    • MaybeKeyedModelHandler.validate_inference_args()
    • MaybeKeyedModelHandler.update_model_path()
    • MaybeKeyedModelHandler.get_preprocess_fns()
    • MaybeKeyedModelHandler.get_postprocess_fns()
    • MaybeKeyedModelHandler.should_skip_batching()
    • MaybeKeyedModelHandler.share_model_across_processes()
    • MaybeKeyedModelHandler.model_copies()
  • OOMProtectedFn
  • RunInference
    • RunInference.annotations()
    • RunInference.from_callable()
    • RunInference.expand()
    • RunInference.with_exception_handling()
  • load_model_status()
• apache_beam.ml.inference.gemini_inference module
  • generate_from_string()
  • generate_image_from_strings_and_images()
  • GeminiModelHandler
    • GeminiModelHandler.batch_elements_kwargs()
    • GeminiModelHandler.create_client()
    • GeminiModelHandler.request()
• apache_beam.ml.inference.huggingface_inference module
  • HuggingFaceModelHandlerTensor
    • HuggingFaceModelHandlerTensor.load_model()
    • HuggingFaceModelHandlerTensor.run_inference()
    • HuggingFaceModelHandlerTensor.update_model_path()
    • HuggingFaceModelHandlerTensor.get_num_bytes()
    • HuggingFaceModelHandlerTensor.get_metrics_namespace()
  • HuggingFaceModelHandlerKeyedTensor
    • HuggingFaceModelHandlerKeyedTensor.load_model()
    • HuggingFaceModelHandlerKeyedTensor.run_inference()
    • HuggingFaceModelHandlerKeyedTensor.update_model_path()
    • HuggingFaceModelHandlerKeyedTensor.get_num_bytes()
    • HuggingFaceModelHandlerKeyedTensor.get_metrics_namespace()
  • HuggingFacePipelineModelHandler
    • HuggingFacePipelineModelHandler.load_model()
    • HuggingFacePipelineModelHandler.run_inference()
    • HuggingFacePipelineModelHandler.update_model_path()
    • HuggingFacePipelineModelHandler.get_num_bytes()
    • HuggingFacePipelineModelHandler.get_metrics_namespace()
• apache_beam.ml.inference.model_manager module
  • GPUMonitor
    • GPUMonitor.fallback_memory_mb
    • GPUMonitor.poll_interval
    • GPUMonitor.peak_window_seconds
    • GPUMonitor.start()
    • GPUMonitor.stop()
    • GPUMonitor.reset_peak()
    • GPUMonitor.get_stats()
    • GPUMonitor.refresh()
  • ResourceEstimator
    • ResourceEstimator.logging_info()
    • ResourceEstimator.is_unknown()
    • ResourceEstimator.get_estimate()
    • ResourceEstimator.set_initial_estimate()
    • ResourceEstimator.add_observation()
  • QueueTicket
  • ModelManager
    • ModelManager.logging_info()
    • ModelManager.all_models()
    • ModelManager.try_enter_isolation_mode()
    • ModelManager.should_spawn_model()
    • ModelManager.acquire_model()
    • ModelManager.release_model()
    • ModelManager.shutdown()
• apache_beam.ml.inference.onnx_inference module
  • OnnxModelHandlerNumpy
    • OnnxModelHandlerNumpy.load_model()
    • OnnxModelHandlerNumpy.run_inference()
    • OnnxModelHandlerNumpy.get_num_bytes()
    • OnnxModelHandlerNumpy.get_metrics_namespace()
• apache_beam.ml.inference.pytorch_inference module
  • PytorchModelHandlerTensor
    • PytorchModelHandlerTensor.load_model()
    • PytorchModelHandlerTensor.update_model_path()
    • PytorchModelHandlerTensor.run_inference()
    • PytorchModelHandlerTensor.get_num_bytes()
    • PytorchModelHandlerTensor.get_metrics_namespace()
  • PytorchModelHandlerKeyedTensor
    • PytorchModelHandlerKeyedTensor.load_model()
    • PytorchModelHandlerKeyedTensor.update_model_path()
    • PytorchModelHandlerKeyedTensor.run_inference()
    • PytorchModelHandlerKeyedTensor.get_num_bytes()
    • PytorchModelHandlerKeyedTensor.get_metrics_namespace()
• apache_beam.ml.inference.sklearn_inference module
  • SklearnModelHandlerNumpy
    • SklearnModelHandlerNumpy.load_model()
    • SklearnModelHandlerNumpy.update_model_path()
    • SklearnModelHandlerNumpy.run_inference()
    • SklearnModelHandlerNumpy.get_num_bytes()
    • SklearnModelHandlerNumpy.get_metrics_namespace()
  • SklearnModelHandlerPandas
    • SklearnModelHandlerPandas.load_model()
    • SklearnModelHandlerPandas.update_model_path()
    • SklearnModelHandlerPandas.run_inference()
    • SklearnModelHandlerPandas.get_num_bytes()
    • SklearnModelHandlerPandas.get_metrics_namespace()
• apache_beam.ml.inference.tensorflow_inference module
  • TFModelHandlerNumpy
    • TFModelHandlerNumpy.load_model()
    • TFModelHandlerNumpy.update_model_path()
    • TFModelHandlerNumpy.run_inference()
    • TFModelHandlerNumpy.get_num_bytes()
    • TFModelHandlerNumpy.get_metrics_namespace()
  • TFModelHandlerTensor
    • TFModelHandlerTensor.load_model()
    • TFModelHandlerTensor.update_model_path()
    • TFModelHandlerTensor.run_inference()
    • TFModelHandlerTensor.get_num_bytes()
    • TFModelHandlerTensor.get_metrics_namespace()
• apache_beam.ml.inference.tensorrt_inference module
  • TensorRTEngine
    • TensorRTEngine.get_engine_attrs()
  • TensorRTEngineHandlerNumPy
    • TensorRTEngineHandlerNumPy.load_model()
    • TensorRTEngineHandlerNumPy.load_onnx()
    • TensorRTEngineHandlerNumPy.build_engine()
    • TensorRTEngineHandlerNumPy.run_inference()
    • TensorRTEngineHandlerNumPy.get_num_bytes()
    • TensorRTEngineHandlerNumPy.get_metrics_namespace()
    • TensorRTEngineHandlerNumPy.validate_inference_args()
• apache_beam.ml.inference.utils module
  • WatchFilePattern
    • WatchFilePattern.expand()
• apache_beam.ml.inference.vertex_ai_inference module
  • VertexAIModelHandlerJSON
    • VertexAIModelHandlerJSON.create_client()
    • VertexAIModelHandlerJSON.request()
    • VertexAIModelHandlerJSON.batch_elements_kwargs()
• apache_beam.ml.inference.vllm_inference module
  • OpenAIChatMessage
    • OpenAIChatMessage.role
    • OpenAIChatMessage.content
  • VLLMCompletionsModelHandler
    • VLLMCompletionsModelHandler.load_model()
    • VLLMCompletionsModelHandler.run_inference()
    • VLLMCompletionsModelHandler.share_model_across_processes()
  • VLLMChatModelHandler
    • VLLMChatModelHandler.load_model()
    • VLLMChatModelHandler.run_inference()
    • VLLMChatModelHandler.share_model_across_processes()
• apache_beam.ml.inference.xgboost_inference module
  • XGBoostModelHandler
    • XGBoostModelHandler.load_model()
    • XGBoostModelHandler.get_metrics_namespace()
  • XGBoostModelHandlerNumpy
    • XGBoostModelHandlerNumpy.run_inference()
    • XGBoostModelHandlerNumpy.get_num_bytes()
  • XGBoostModelHandlerPandas
    • XGBoostModelHandlerPandas.run_inference()
    • XGBoostModelHandlerPandas.get_num_bytes()
  • XGBoostModelHandlerSciPy
    • XGBoostModelHandlerSciPy.run_inference()
    • XGBoostModelHandlerSciPy.get_num_bytes()
  • XGBoostModelHandlerDatatable
    • XGBoostModelHandlerDatatable.run_inference()
    • XGBoostModelHandlerDatatable.get_num_bytes()