Package: klay-beam
Klay BeamHelpers for running massively parallel Apache Beam jobs on audio data.NOTE:This is Beta. Documentation is incomplete. Expect breaking changes
prior to v1.0.Processing large batches of audio data can be very time consuming. It is often
helpful to use many VM instances to apply transformations or extract features
from a large audio dataset. This package bundles a collection of utilities and
examples designed to ease the process of massively parallel audio jobs with
GCP Dataflow and Apache Beam.The core transformations include:File manipulations (for local filesystem and cloud storage)Audio File writing and readingFeature writing (and reading?)file name mutations: Moving, and extension mutationSkipCompletedAudio data resamplingAudio channel manipulationYou can useklay-beamto write and launch your own custom jobs that build on
top of these primitives. It is setup to support a wide variety of custom
dependencies and environments, including pinned versions of Python, Pytorch,
CUDA (for GPU support), and more.Running LocallyTypically you will want to write and test a job on local machine, before testing
and executing on a massive dataset. For example:# Create the environment and install klay_beamcondaenvcreate-fenvironment/py3.10-torch2.0.yml
condaactivateklay-beam-py3.10-torch2.0
pipinstall-e".[tests, code-style, type-check]"Then launch the example job:# Then launch the job. Running locally allows you to use --source_audio_path#  values paths on your local filesystem OR in gs://object-storage. To use gs://# directories, you must be authenticated with GCPpython-mklay_beam.run_example\--runnerDirect\--source_audio_suffix.mp3\--source_audio_path'/local/path/to/mp3s/'Running on GCP via DataflowIf your audio files are in cloud storage you can process them using GCP
Dataflow, which allows for massive parallel execution. This requires additional
setup, including:Activate Dataflow APICreate GCP service accountCreate GCP Cloud Storage bucketSetup GCP permissions for launching and executing jobsFinally, you need a specialized docker container that bundlesapache_beam,klay_beam, and any additional dependencies. SeeMakefilefor examples.Setup GCPTo get started, setup a GCP project by following this steps below, which were
adapted from theDataflow Quickstart Guide.# Manually set the following variablesGCP_PROJECT_ID=your-gcp-project# ID of the GCP project that will run jobsUSER_EMAIL=you@example.com# The email associated with your GCP accountDATAFLOW_BUCKET_NAME=your-bucket# Temp data storage bucket for beam workersGCP_SA_NAME=beam-worker# GCP service account name used by beam workers# Compute the full email of the service account used by beam workersGCP_SA_EMAIL=${GCP_SA_NAME}@${GCP_PROJECT_ID}.iam.gserviceaccount.com# Compute a valid temporary storage path for job workers to use at runtime. The# example below is just a proposal. You can use any cloud storage path, as long# the Beam workers are able to write temporary files to this path during job# execution.TEMP_GS_URL=gs://${DATAFLOW_BUCKET_NAME}/tmp/# Create and activate a GCP project. You can skip `gcloud projects create` if# you have an existing gcp project that you want to use.gcloudinit
gcloudprojectscreate${GCP_PROJECT_ID}gcloudconfigsetproject${GCP_PROJECT_ID}# Make sure that billing is enabled for your project. If billing is not not# enabled, use the GCP console to enable it.gcloudbetabillingprojectsdescribe${GCP_PROJECT_ID}gcloudservicesenabledataflowcompute_componentloggingstorage_componentstorage_apibigquerypubsubdatastore.googleapis.comcloudresourcemanager.googleapis.com
gcloudauthapplication-defaultlogin# Dataflow jobs need to write temporary data to cloud storage during job# execution. Create a bucket using the gsutil mb (make bucket) command. See# `gsutil help mb` for details.gsutilmb--autoclass-lUS-bongs://${DATAFLOW_BUCKET_NAME}# Create a service account which will be used by the worker nodesgcloudiamservice-accountscreate$GCP_SA_NAME--description="Service account used by Apache Beam workers"--display-name="Beam Worker"# Give the service account access it needsgcloudprojectsadd-iam-policy-binding${GCP_PROJECT_ID}--member="serviceAccount:${GCP_SA_EMAIL}"--role=roles/dataflow.admin
gcloudprojectsadd-iam-policy-binding${GCP_PROJECT_ID}--member="serviceAccount:${GCP_SA_EMAIL}"--role=roles/dataflow.worker
gcloudprojectsadd-iam-policy-binding${GCP_PROJECT_ID}--member="serviceAccount:${GCP_SA_EMAIL}"--role=roles/storage.objectAdmin# Note that the last command above will give the service account (and any users# who can impersonate the service account) full access to ALL buckets in the# project. If this is undesirable, you can use the Cloud Storage section of# console.cloud.google.com to give the service account access to ONLY specific# buckets. To do this, navigate to a bucket, and click the "permissions" button.## If you choose bucket level permissions, you must also grant:# - read+list access to buckets where source data is saved# - write access to buckets where result data will be persisted# To allow users to impersonate the service account, run the following command# which grants a user the `roles/iam.serviceAccountUser` (AKA "Service Account#  User") role, but only for a specific service account:gcloudiamservice-accountsadd-iam-policy-binding${GCP_SA_EMAIL}\--member="user:${USER_EMAIL}"\--role="roles/iam.serviceAccountUser"# Alternatively, if you want to grant the user access to impersonate ALL service# accounts, use this command instead:gcloudprojectsadd-iam-policy-binding${GCP_PROJECT_ID}\--member="user:${USER_EMAIL}"\--role=roles/iam.serviceAccountUserLaunch GCP Dataflow Job# Run remotely via GCP Dataflow. Should be executed in the `klay-beam` conda# environment to ensure Beam SDK, python, and dependency parity between the# local environment and Worker environments.# You will need the following configuration values from the setup (above)GCP_PROJECT_ID=<your-gcp-project>GCP_SA_EMAIL=<your-service-account>@<your-gcp-project>.iam.gserviceaccount.comTEMP_GS_URL=gs://<your-gs-bucket>/<your-writable-dir/># Additionally, you need a compatible beam container image, and an gs:// url# that contains the audio files you want to read. You must ensure that the# service account has read access to these audio files.## The klay-beam image that you select should match the version(s) you have# installed locally. This docker image in the example below uses:## - klay_beam 0.12.2# - python 3.9# - apache_beam 2.51.0# - pytorch 2.0KLAY_BEAM_CONTAINER='klaymusic/klay-beam:0.12.2-py3.9-beam2.51.0-torch2.0'AUDIO_URL='gs://<your-audio-bucket>/audio/'python-mklay_beam.run_example\--runnerDataflowRunner\--max_num_workers=128\--regionus-central1\--autoscaling_algorithmTHROUGHPUT_BASED\--service_account_email${GCP_SA_EMAIL}\--experiments=use_runner_v2\--sdk_container_image${KLAY_BEAM_CONTAINER}\--sdk_location=container\--temp_location${TEMP_GS_URL}\--project${GCP_PROJECT_ID}\--source_audio_suffix.mp3\--source_audio_path${AUDIO_URL}\--machine_typen1-standard-8\--job_name'example-job-000'Notes:When running on Dataflow you can use the--setup_fileoption to upload a
local package to the workers. For example, when running with--runner DataflowRunner,--setup_file=./your_job/setup.pywould causeyour_jobto be bundled as ansdistand installed on the worker nodes
replacing any existing installation ofyour_jobthat may be in the docker
container. Any missing pip dependencies specified inyour_job/pyproject.tomlwill also be installed at runtime.options for--autoscaling_algorithmareTHROUGHPUT_BASEDandNONECustom Docker Images on DataflowIf you are storing your docker images in a private repo, use the IAM section ofhttps://console.cloud.google.comand grant the "Artifact Registry Reader" role
to your Beam worker service account.DevelopmentQuick StartCreatecondaenvironment. Environments labeledlocalare likely to work on
linux albeit without cuda support:condaenvcreate-fenvironments/py310-torch.local.ymlTo create or update an environment:condaenvupdate-fenvironment/py310-torch.local.ymlDocker ContainersWhen you launch a Beam job with--runner DataflowRunnerthat job will run via
theGCP Dataflowservice. It is usually best to specify the Docker
container that will run on Beam worker nodes in GCP Compute (via the--sdk_container_imageflag). However, you cannot use any docker image.
Instead, you must prepare an image specifically for working with Dataflow.Pre-configured docker images with a variety of dependencies and python version
are available athub.docker.com/repository/docker/klaymusic/klay-beam.You can also use./docker-build.shto further customize your images.When running a job, the docker image you select will be run on all workers. When
running a Beam job on GCP Dataflow you may specify a local python package the be
bundled as a pythonsdistand installed in each docker container. To send a
local package to your worker nodes, supply the--setup_fileCLI flag when
launching your job (local packages mustsetup.pyfile for this to work).
Missing dependencies will be installed using pip. However, to save time, it is
ideal to bundle large dependencies (or non-pip dependencies such as ffmpeg 4)
in the docker container.Publishing to pip/PyPI and DockerHubEnsure:Tests passsrc/klay_beam/__init__.pyhas the desired versionAnd you are ready to publish to pip, create a git tagTag the release:gitcheckoutmain
gitpull
gittagv$(./get_version.sh)gitpushoriginv$(./get_version.sh)Push topypibranch to publish to pip (see../.github/workflows/publish-pypi.yaml):gitcheckoutpublish-pypi
gitmergemain--ff-only
gitpushOnce pip action completes successfully, push to DockerHub (see../.github/workflows/publish-docker-hub.yaml):gitcheckoutdocker-hub
gitmergemain--ff-only
gitpushCode QualityTestingWe usepytestfor testing, there's no coverage target at the moment but
essential functions and custom logic should definitely be tested. To run the
tests:maketestsCode StyleWe useflake8for linting andblackfor formatting.makecode-styleStatic TypingWe check static types usingmypy.maketype-checkDesign PatternsIn Apache Beam, aPipelineis a Directed Acyclic Graph.Each node in the graph is a data processing operation called aPTransformor "Parallel Transform".PTransformsaccept one or morePCollectionsas input, and output one
or morePCollectionsPipelines written withklay_beamtypically perform operations on audio data.
The example below reads all.wavfiles in a cloud storage bucket, calculates
their total length, and logs the names of any corrupted/unreadable files.withbeam.Pipeline(argv=pipeline_args,options=pipeline_options)asp:audio,failed,durations=(p# MatchFiles produces a PCollection of FileMetadata objects|beam_io.MatchFiles("gs://your-bucket/**.wav")# Prevent "fusion"|"Reshuffle">>beam.Reshuffle()# ReadMatches produces a PCollection of ReadableFile objects|beam_io.ReadMatches()|"Load Audio">>beam.ParDo(LoadWithTorchaudio()).with_outputs("failed","duration_seconds",main="audio"))(durations|"SumLengths">>beam.CombineGlobally(sum)|"LogDuration">>beam.Map(lambdax:logging.info("Total duration of loaded audio: "f"~={x:.3f}seconds "f"~={x/60:.3f}minutes "f"~={x/60/60:.3f}hours")))(failed|"Log Failed">>beam.Map(lambdax:logging.warning(x))|"Count">>beam.combiners.Count.Globally()|"Log Failed Count">>beam.Map(lambdax:logging.warning(f"Failed to decode{x}files")))Manyklay_beampipelines, including the example above, start with the
following sequence ofPTransforms.apache_beam.io.fileio.MatchFilesReturns a collection ofFileMetadatas.apache_beam.Reshuffle()apache_beam.io.fileio.ReadMatchesReturns a collection ofReadableFiles.An audio loader for accessing audio data as apytorchTensor ornumpyarray:apache_beam.ParDo(klay_beam.transforms.LoadWithTorchaudio())apache_beam.ParDo(klay_beam.transforms.LoadWithLibrosa())The audio loaders above returns a collection of tuples with shape(filename, audio_data, sample_rate),MatchFilesTheMatchFilesTransforms returns a PCollection ofapache_beam.io.filesystem.FileMetadatainstances, which have the following
properties (code):path: str
size_in_bytes: int
last_updated_in_seconds: floatNote that in GCP Cloud Storage, thelast_update_in_secondsproperty reflectsAutoClasschanges.Preventing FusionTransforms such asMatchFilesoutput PCollections with MANY elements relative
to the number of input elements. This called a "fan-out" transform. Large
fan-out transforms should pre proceeded by a Reshuffle when running on GCP
Dataflow. SeePreventing
Fusionin the Dataflow docs.ReadMatchesTheReadMatchestransform that returns a PCollection ofapache_beam.io.fileio.ReadableFileinstances
(code).
Note that the name of this transform is misleading. It does not actually "read"
the file. It is a lightweight wrapper around the Metadata thatProvides helper methods for accessing matched filesSkips over any matched directories (as per theskip_directories=Trueinitializer arg)Forwards thecompression_typeinitializer argument (default:None) to
helper methodsThe Transforms that follow ReadMatches (often aklay_beamaudio loader) should
expect elements to beReadableFileinstances, which have a.metadataproperty and 3 additional helper methods:metadata:apache_beam.io.filesystem.FileMetadataopen(self,mime_type='text/plain',compression_type=None)->io.BufferedReader# (for gs:// paths)read(self,mime_type='application/octet-stream')->FileLikeread_utf8(self)LoadWithTorchaudioLoadWithTorchaudio is a custombeam.DoFn, turned into a PTransform via thebeam.ParDohelper. See the source for implementation details. Generally,
custom functions have a few requirements that help them work well in on
distributed runners. They are:The function should be thread-compatibleThe function should be serializableRecommended: the function be idempotentFor details about these requirements, see the Apache Beam documentation:https://beam.apache.org/documentation/programming-guide/#requirements-for-writing-user-code-for-beam-transforms