Blog

Coiled Stability Update

2026-07-08T00:00:00+00:00

Coiled experienced two outages in the last few weeks; several production resources began operating too close to their limits, and our monitoring did not surface the risk early enough. The result was multiple production outages. This level of stability does not meet Coiled’s standard for a service customers rely on, so we want to share what happened and what we have done to address it.

Database CPU Saturation

On May 6, the production database became the bottleneck when the Aurora writer sustained near-max CPU utilization. During that period, database-backed platform operations became slow or unreliable, including cluster startup. Each cluster startup includes a check of current account usage and quota to ensure users don’t go past their desired spend limit. That check performed a CPU-intensive scan to sum historical usage, and as startup volume increased, both the number of checks and the amount of historical usage data grew. Together, that put too much load on the database writer.

This has been fixed by moving usage calculations to incremental running totals. Instead of scanning historical usage on every cluster startup, quota checks now use a stored checkpoint and only process newer usage events.

RDS CPUUtilization panel in the AWS console.

The impact is visible in RDS CPU utilization for the production Aurora writer. Before the deployment on July 8, CPU was highly variable with frequent spikes; the deployment occurred after noon, and utilization was lower and more stable afterward. Improved database monitoring and alerting is also now in place so database pressure is visible earlier.

Application Memory Saturation

On June 29, the application infrastructure serving platform traffic became non-responsive due to memory saturation, causing platform requests to fail. Increased platform traffic had driven memory utilization higher across the app stack, pushing the service too close to its operating limits.

This has been addressed with the following interventions:

increasing application serving capacity and memory headroom
rolling application logs more aggressively after disk pressure was found to be contributing to instability
tuning app worker counts
removing unnecessary application dependencies

Application memory utilization in Grafana.

Grafana metrics show the impact of the resize: before the incident, application memory utilization was steadily climbing and reached roughly 90%, leaving little operating headroom. After application serving capacity was increased later on June 29, memory utilization dropped sharply and stabilized around 25-40%, leaving substantially more room for traffic spikes.

This incident also exposed a monitoring gap. During our move from Datadog to Grafana, host-level resource alerts had not been fully re-established, so memory and disk pressure were not surfaced early enough to operating engineers. Alerting on host saturation is essential for a service like Coiled, and host-level Grafana monitoring and alerting is now in place for metrics like memory and disk utilization.

Silver Lining

These failures were due to load increasing beyond our expectations and plans. Usage has roughly doubled since the beginning of the year. While this doesn’t excuse outages, it does signal that the company itself is healthy and highly motivated to keep the platform well supported and running smoothly. This is our top priority.

Closing

The main lesson is that saturation becomes a production risk before a resource reaches 100%. We have added capacity, reduced avoidable database load, and put monitoring in place so these pressures are visible earlier.

We know Coiled is a service that customers depend on to run important workloads. Supporting those workloads reliably is our highest priority, and we will continue to invest in the capacity, monitoring, and operational practices needed to meet that bar.

MLflow on Coiled with Sidecars and Filestores

2025-10-20T00:00:00+00:00

Oct 20, 2025

4m read

Coiled lets you run ephemeral clusters in the cloud. We introduce sidecars to run arbitrary containers alongside your workload, and filestores to sync local data to the cloud and persist data across runs. We use them to run an MLflow server concurrently with model training.

MLflow is a widely used open source tool for ML experiment tracking, model packaging, registry management, and deployment. Some people host their own server, but many use a hosted solution. Our users asked us about running MLflow on Coiled, and, using some new features, we finally found a way that aligns with Coiled’s principle of only paying for what you use. It’s idiosyncratic, but it works quite well.

Running MLflow as a Sidecar

It’s always been possible to run MLflow on Coiled using coiled run, but this required starting the server before training, obtaining the server’s URL, and passing it to the training code—too many moving parts.[1]

Instead, we introduce sidecars, a feature that lets you specify one or more Docker containers to run on the scheduler and/or workers during jobs. With them, you can run the MLflow server alongside the training code. We have no intention of going full Docker Compose on you, but you may recognize the format if you squint at a sidecar spec YAML file:

mlflow:
  image: astral/uv:debian-slim
  command: uvx mlflow server --host 0.0.0.0 --port 5000 
  ports:
    - 5000

There’s a name for the sidecar, the image to use, the command to run, and any ports to expose. To launch MLflow, we use a small uv container and uvx instead of using the official MLflow image, like we did in the marimo post.

You can also define sidecars directly in your Python code. Here’s an example of adding a sidecar to a Dask cluster:

cluster = coiled.Cluster(
    ...,
    scheduler_sidecars=[
        {
            "name": "mlflow",
            "image": "astral/uv:debian-slim",
            "command": "uvx mlflow server --host 0.0.0.0 --port 5000",
            "ports": [5000]
        }
    ]
)

Persisting Data with Filestores

To persist data across jobs, we use filestores, a new feature that enables:

Copying data between your local machine and Coiled-managed VMs in the cloud,
Persisting it between multiple runs, and
Transferring data between steps in a processing pipeline.

Here’s a quick example of using a filestore with coiled batch to copy local files from local-directory to the cloud VM where my_script.py runs:

coiled batch run --sync local-directory --wait my_script.py

The files will be available in the workers’ working directory, and any outputs that the code writes to the working directory on the cloud VM will be stored in a filestore, available for download to the local machine or to other Coiled-managed cloud VMs. All tasks in a batch job share the filestore, and you can sync back outputs from all the tasks together to your local machine with:

coiled file download --into output-data

In the example above, we add the --wait flag to wait until work completes, and then sync new or modified files back to the local machine, instead of the default non-blocking behavior for coiled batch run.

The sidecar spec lets you specify filestores to attach. Data in named filestores are downloaded to the file:///scratch/filestores directory on the VMs. Here, we attach a filestore named mlflow-data to persist the MLflow backend data:

mlflow:
  image: astral/uv:debian-slim
  command: uvx --python 3.13 mlflow server --host 0.0.0.0 --port 5000 
    --backend-store-uri file:///scratch/filestores/mlflow-data
  ports:
    - 5000
  filestores:
    - name: mlflow-data

Training Code and Launching the Cluster

With Coiled Batch, the training code doesn’t need to know anything about Coiled. That way, it can also run locally. Yet, we find it convenient to set a few # COILED comments in the header that Coiled will pick up, like:

# COILED env MLFLOW_TRACKING_URI=http://coiled-scheduler:5000
# COILED env MLFLOW_TRACKING_USERNAME=alex
# COILED gpu true
# COILED map-over-values 0.1, 0.05, 0.01

These comments:

Set the MLflow server’s URL as the scheduler’s address, now available as http://coiled-scheduler on the workers (it’s still available as the COILED_BATCH_SCHEDULER_ADDRESS environment variable too),
Set the user name so it’s logged with the experiment results,
Request VMs with a GPU attached,
And specify the list of learning rates to try. Coiled will launch a VM for each value.

The model we train in the example is a straightforward PyTorch classifier on the FashionMNIST dataset and doesn’t contain any Coiled-specific code. The only notable thing is that it accepts the loss value as an argument, so we can train the model locally and on the cloud without modifying the code. You can run it locally with:

python train.py 0.1

Or in the cloud, with the MLflow sidecar alongside, with:

coiled batch run --scheduler-sidecar-spec mlflow-sidecar.yaml \
   -- train.py \$COILED_BATCH_TASK_INPUT

While the cluster is running, you can access the MLflow server using the button on the cluster’s page:

Once the cluster has stopped, you can download the filestore data to your computer with:

coiled file download --into mlflow-data

And start the MLflow server to browse the results:

mlflow server \
        --backend-store-uri mlflow-data \
        --artifacts-destination mlflow-data

Check out the complete MLflow example to learn how to browse MLflow data remotely and split the metadata and model data into separate filestores.

What Else Can You Do With Sidecars and Filestores?

Other than running MLflow, we know some users plan on using sidecars to run a proxy on each worker to make authorized, encrypted connections to databases.

As for filestores, they address a common request to make it easy to upload local files to clusters when using Coiled Batch. We think you’ll like them. Check out the filestore docs to learn more, and let us know what you think about them.

How to run pandas on GPUs with Coiled

2025-09-05T00:00:00+00:00

Sept. 5, 2025

6 min read

Sarah Johnson

Pandas is the go-to library for data manipulation and analysis in Python, but it struggles with large datasets due to its single-threaded, CPU-bound design. As data grows in size and complexity, analysts and data scientists often face slow processing times and memory bottlenecks that interrupt analytical flow and limit business responsiveness.

Traditional pandas is single-threaded and CPU-bound, which makes it great for prototyping but increasingly inefficient for large-scale data work. By offloading operations to parallel GPU cores, you can accelerate pandas-like workflows without rewriting your entire codebase. Accessing GPUs can be challenging, though. Most people don’t have NVIDIA GPUs locally, so they need cloud GPUs. This typically involves complex cloud setup, IAM configuration, Docker images, infrastructure maintenance, and cost management.

With tools like RAPIDS cuDF and Coiled, it’s possible to scale these workloads easily and access GPU acceleration with minimal code changes.

In this article, you’ll learn how to:

Accelerate pandas operations using cuDF with minimal code changes
Easily access cloud GPUs with Coiled
Achieve dramatic performance improvements on real datasets
Determine whether GPU-acceleration might be a good fit for your workflow

Why GPU-accelerated pandas delivers faster performance

Many data analysis tasks on large datasets are well-suited for GPU acceleration. Embarrassingly parallel operations like filtering, element-wise mappings, and reductions can be run in parallel on a GPU rather than sequentially, resulting in dramatically shorter computation times.

GPU-accelerated pandas operations can achieve performance improvements of 10-150x compared to CPU-only processing. Real-world benchmarks demonstrate these dramatic speedups:

H2O.ai Database-like Ops Benchmark: 150x faster processing on 5GB datasets with complex joins and group-by operations
NYC ride-share analysis: 8.9x faster execution time on 64.8 million records

By leveraging GPU capabilities, data scientists and analysts can:

Process larger datasets without hitting memory limitations
Run transformations and calculations significantly faster
Iterate more quickly on data exploration and model development

How Coiled simplifies GPU access

Despite the performance advantages of GPUs, accessing cloud GPU resources often requires complex infrastructure setup, specialized cloud configuration, and significant DevOps expertise. Coiled eliminates these barriers by making cloud GPUs accessible to anyone with basic Python skills, removing the need for specialized hardware investments or cloud infrastructure management.

The Coiled platform removes common barriers to GPU computing, including:

Hardware acquisition costs
Complex cloud infrastructure configuration
Driver compatibility issues
DevOps expertise requirements
Resource provisioning and scaling complexity

Example: How often do New Yorker’s tip?

To see how to use Coiled for GPU-accelerated pandas in practice, we’ll look at a subset of Uber and Lyft rides in New York City over the last several years and perform some simple calculations.

Install Coiled

Next, we’ll install Coiled.

$ pip install coiled
$ coiled login

This will redirect us to the Coiled website to authenticate our computer. It will also save our Coiled token on our computer.

Create a GPU-enabled software environment

By default, Coiled automatically copies our local environment to remote VMs. However, if we do not have an NVIDIA GPU available locally, then we need to create a software environment or use the RAPIDS Docker container (see the RAPIDS installation guide).

To define a software environment, we define the dependencies in a conda environment YAML file:

rapids-env.yaml

channels:
  - rapidsai
  - conda-forge
  - nvidia
dependencies:
  - python=3.12
  - rapids>=25.1
  - cuda-version>=12.0,<=12.9
  - s3fs
  - jupyterlab

And then create a Coiled software environment:

coiled env create --name rapids --gpu-enabled --conda rapids-env.yaml

Launch a GPU-enabled Jupyter Notebook

Next, we start a JupyterLab instance on a GPU-enabled VM on the cloud:

coiled notebook start --gpu --software rapids

Under the hood, Coiled handles:

Provisioning a cloud VM with GPU hardware. In this case, a g4dn.xlarge (Tesla T4) instance on AWS.
Downloading the software environment onto the VM
Tearing everything down once you’re done (no idle costs)

We could also include additional arguments like:

--sync to live-sync local files, including Jupyter Notebooks, to the cloud machine. Any changes made in the notebook running remotely will be synchronized back to our computer. We need to install Mutagen for that.
--vm-type to request a specific VM type.
--region to specify a region. We find GPUs are often easier to get in AWS region us-west-2.

See the Coiled notebooks documentation for more details.

Running GPU-accelerated `cudf.pandas`

NVIDIA RAPIDS offers GPU acceleration for data science workloads with minimal code changes. The cudf.pandas accelerator enables instant GPU execution of pandas operations by loading the extension with: %load_ext cudf.pandas.

We adapted this example from the RAPIDS documentation. We’ll look at a subset of the NYC taxi dataset (~60 million rows) and calculate how often New Yorkers tip.

%load_ext cudf.pandas
import pandas as pd
import s3fs

path_files = []
for i in range(660,715):
    path_files.append(pd.read_parquet(f"s3://coiled-data/uber/part.{i}.parquet", storage_options={"anon": True}))
df = pd.concat(path_files, ignore_index=True)

df["tipped"] = df["tips"] != 0
df["tip_frac"] = df["tips"] / (df["base_passenger_fare"] + df["tolls"] + df["bcf"] + df["sales_tax"] + df["congestion_surcharge"].fillna(0) + df["airport_fee"].fillna(0))
df["service"] = df["hvfhs_license_num"].map({
    "HV0003": "uber",
    "HV0005": "lyft",
    "HV0002": "juno",
    "HV0004": "via",
})
df = df.drop(columns="hvfhs_license_num")

df["tipped"].mean()

We can monitor GPU utilization from the Coiled dashboard as the code runs (and afterwards too):

Run a Python script on a cloud GPU

While Jupyter Notebooks are helpful for interactive workflows, there are a number of ways we can use Coiled to access GPUs. In particular, Coiled Batch is an easy way to run any script on the cloud. It is especially useful for scaling large, independent tasks (ie, embarrassingly parallel workloads) across hundreds of VMs.

We could run the same analysis as above by adding the following metadata to the top of the script:

my_script.py

# COILED gpu True
# COILED software rapids
# COILED region us-east-2

import pandas as pd
import s3fs

path_files = []
for i in range(660,715):
    path_files.append(pd.read_parquet(f"s3://coiled-data/uber/part.{i}.parquet", storage_options={"anon": True}))
df = pd.concat(path_files, ignore_index=True)

df["tipped"] = df["tips"] != 0
df["tip_frac"] = df["tips"] / (df["base_passenger_fare"] + df["tolls"] + df["bcf"] + df["sales_tax"] + df["congestion_surcharge"].fillna(0) + df["airport_fee"].fillna(0))
df["service"] = df["hvfhs_license_num"].map({
    "HV0003": "uber",
    "HV0005": "lyft",
    "HV0002": "juno",
    "HV0004": "via",
})
df = df.drop(columns="hvfhs_license_num")

df["tipped"].mean()

and then launch the script with the following command, using the RAPIDS CLI to use cudf.pandas:

coiled batch run python -m cudf.pandas my_script.py

Under the hood Coiled will:

Spin up appropriate machines as defined by # COILED comments
Download the specified software environment
Run the script
Shut down the machines

In this example, we’re running a single Python script, but Coiled Batch especially excels at easily scaling large embarrassingly parallel jobs, without the scheduler overhead of tools like Dask. Some common use cases include:

Running an arbitrary bash script 100 times in parallel
Training a PyTorch model on a GPU
Multi-node GPU LLM fine-tuning

When to use GPU acceleration

While GPUs can offer impressive speedups, they’re not always the right choice, and there’s some nuance to using them well. Some general guidelines include:

1. Minimize memory transfers

Avoid moving data back and forth between CPU and GPU repeatedly
The PCIe bus transfer time can dominate your computation time
For small datasets with simple operations, this transfer overhead can make GPU processing slower than CPU

2. Maximize operational intensity

Choose problems that do lots of computation on relatively small amounts of data
Good fit: Matrix multiplication (O(n³) operations on O(n²) data)
Poor fit: Simple string matching (O(n) operations on O(n) data)
Use roofline analysis: memory-bound operations may not benefit from GPU acceleration

3. Don’t launch lots of small kernels

Avoid operations that create thousands of tiny CUDA kernels
Each kernel launch has overhead - many small operations can be slower than fewer large ones
Let optimized libraries like cuDF handle the kernel management

Kaashif Hymabaccus gave a detailed talk on this topic at PyCon US.

FAQs

Can you run pandas on a GPU?

Traditional pandas cannot run directly on a GPU, but RAPIDS cuDF offers a nearly identical API that executes on GPU hardware.

What is the fastest way to run pandas on GPUs?

The fastest way to run pandas on GPUs is by using RAPIDS cuDF with the cudf.pandas accelerator. It requires minimal code changes and can achieve up to 150x speedup on analytics workloads. You can scale this further using Dask-cuDF for multi-GPU processing and manage GPU clusters easily with Coiled’s cloud platform.

How do I set up Coiled for GPU-accelerated pandas workflows?

To use Coiled for GPU workloads, install Coiled via pip/conda/uv and launch a GPU-enabled cluster. Coiled handles all cloud provisioning, GPU drivers, and scaling. See the Coiled documentation on GPUs to learn more.

When should I use GPU acceleration for pandas workflows?

Consider using GPU acceleration for:

Large datasets with high computational intensity: There is some overhead with copying data to the GPU, so your data must be sufficiently large to offset this cost.
Matrix operations and linear algebra: Operations that scale as O(n²) or O(n³) with input size
Parallel aggregations: Group-by operations on large datasets with many groups (where data shuffling isn’t required)
Batch processing: You can process multiple operations together rather than one at a time

When should I stick with a CPU for pandas workflows?

GPU acceleration may actually hurt performance when:

Small datasets: Transfer overhead dominates computation time
Simple operations: String searches, basic filtering, or operations with low computational intensity
Memory-bound problems: Operations limited by memory bandwidth rather than compute capacity
Frequent CPU-GPU transfers: Workflows that require moving data back and forth repeatedly

Building a scientific Python stack with Coiled

2025-08-26T00:00:00+00:00

Aug 26, 2025

9 min read

At the heart of most scientific computing stacks are a handful of core libraries that provide general-purpose numerical methods, domain-specific analyses, and the machinery for machine learning.

NumPy and SciPy provide fundamental array operations, linear algebra, and scientific routines. Pandas and Xarray handle structured/tabular and labeled multi-dimensional data, respectively. Matplotlib and Plotly enable visualization. Scikit-learn, PyTorch, and TensorFlow power machine learning and deep learning workflows. Dask scales computations from a single laptop to a distributed cluster.

In many cases, custom libraries or domain-specific frameworks are also part of the mix, especially for specialized fields like geospatial analysis, genomics, or computational physics.

Traditionally, these tools are used on local machines, possibly leveraging multi-core CPUs or single-GPU setups. As dataset sizes grow or as computational needs (such as parameter sweeps, simulations, or parallel data processing) increase, teams quickly reach the limits of what a single piece of hardware can deliver.

The challenge: Scaling beyond a workstation

Scaling up scientific Python workflows introduces technical barriers that often distract from the actual scientific or analytical task. Installing and synchronizing software environments across different machines creates consistency challenges. Managing cluster orchestration—starting, stopping, and monitoring resources—requires specialized expertise. Teams must troubleshoot network, storage, and reliability issues while dealing with cost controls and resource optimization to avoid unexpected bills.

For a data science team in biomedical research, e.g., these challenges might manifest when analyzing terabytes of RNA sequencing data. A local workstation might handle preliminary analyses, but running full pipelines, bootstrapped permutations, or deep learning model training on such massive datasets is impractical without access to scalable compute resources.

The transition to distributed computing is often daunting. Traditional solutions require extensive setup through frameworks like MPI or Kubernetes and frequent code rewrites to adapt to new APIs or cluster environments. As a result, many teams either over-invest in local infrastructure with high up-front costs or underutilize the cloud due to operational friction.

Coiled: A platform for scalable Python computing on the cloud

Coiled addresses these obstacles by providing a platform that lets teams run their Python code—unchanged—on the cloud, at any scale. The platform enables frictionless development where Python developers can continue using their existing tooling and libraries, with no need to rewrite code for new environments.

Coiled provides efficient, on-demand cloud computing by provisioning the right resources in the right cloud regions, running jobs, and tearing down resources automatically. Environment synchronization ensures whatever is installed on your development machine is mirrored in the cloud cluster, maintaining consistency and minimizing dependency errors.

Integrated dashboards and auto-shutdown prevent runaway costs and provide clarity on performance through built-in observability and cost controls.

Consider a research group in climate science building an ensemble model to forecast temperature patterns at high spatial resolution. Their workflow starts with basic array processing and data cleaning but quickly escalates to bootstrapped uncertainty estimation and model validation across decades of global data. With Coiled, they simply update their workflow to request cloud resources, and everything else remains familiar. The cluster scales up, the code runs, and results are returned—without the team needing to assemble custom Docker images or coordinate with IT to set up cloud networking.

Integration with the scientific Python stack

Coiled fits into the existing ecosystem for scientific Python by leveraging key libraries. Through Dask integration, Coiled enables parallel and distributed computing with minimal code changes. Complex dataframes, arrays, and machine learning computations can be chunked and performed in parallel, using the same high-level APIs as local development.

Beyond distributed Dask clusters, Coiled Batch enables parallel processing of any code—whether that’s running hundreds of Python simulations with different parameters or leveraging Julia for high-performance scientific computing. This flexibility means teams can parallelize Monte Carlo simulations, parameter sweeps, or any embarrassingly parallel workload across cloud resources without being limited to specific frameworks.

Custom environments allow researchers to add dependencies or switch between projects without delay. For example, a genomics team needing a new release of Biopython only needs to install it locally; Coiled syncs the environment automatically when the job is launched on the cloud.

Coiled works with existing Python packaging, workflow orchestration tools like Prefect or Dagster, and any Python-compatible editor or notebook environment. Teams can mix and match open-source tools without vendor lock-in.

Consider the Lawrence Berkeley National Lab using Xarray to analyze satellite imagery data. As the area of analysis grows from local watersheds to entire continents, memory limitations become severe. Rather than rewriting their codebase, the team used Coiled to provision a distributed cluster, leveraging cloud resources to handle continent-scale analyses.

Example: Mineral exploration at scale

KoBold Metals, a company dedicated to discovering new battery metals, depends on evaluating vast regions for signs of subsurface mineral deposits using remote sensing, geophysics, and AI models.

Early in a project, they analyze limited geochemical samples and low-resolution magnetic maps. But as interest grows in a region, the data volume expands: airborne geophysics, electromagnetic surveys, drilling logs, and hyperspectral imagery all need to be synthesized at fine spatial resolution.

On a single machine, these tasks are not feasible. Workflows such as computing spatial derivatives via Fourier transforms on gigapixel magnetic map rasters cannot complete in a reasonable timeframe or within available memory.

Using Coiled, KoBold shifts its analysis to a distributed Dask cluster, using their existing Python codebase. The computation, which previously failed or took days, now completes in hours across the cloud, allowing the exploration team to rapidly test hundreds of geological hypotheses. This approach increases scientific rigor and shortens time-to-decision, letting the business allocate exploration budgets where data suggests the best prospects.

As KoBold scales its operations, the technical team does not need to build or maintain cloud infrastructure. Scientists and engineers can install the latest analysis libraries as needed, confident that their development environment will be matched in each distributed run.

Observability and optimization

Scientific computing demands not only raw power but also transparency. When data loads or model training jobs fail, researchers need immediate, actionable feedback. Coiled provides detailed dashboards showing cluster health, progress, and cost, with direct links to Python-level logs and errors.

A data science team training deep learning models on environmental sensor data, using Optuna and Dask to parallelize hyperparameter tuning, might notice that memory usage spikes in certain epochs or that certain data partitions take significantly longer to process. Using Coiled’s observability features, they identify bottlenecks and optimize both their models and cluster resource allocation.

Cost control and cloud efficiency

Cloud resources can be cost-effective when used conscientiously, or can become unexpectedly expensive through misuse or oversight. Common pitfalls include leaving machines running after computations are complete, unnecessary data transfer between cloud regions, or selecting oversized machine types for small jobs.

Coiled is designed to encourage responsible usage. Clusters auto-shutdown when idle, leaving no orphaned resources. Dashboards make cloud spend visible and actionable. Jobs and resources are configured to match the scale of the task, avoiding both underutilization and overprovisioning.

For a biomedical research lab processing hundreds of thousands of medical images nightly, this means they can scale up to meet demand but also automatically scale down—maintaining only minimal persistent resources and controlling costs in line with real activity.

Focus on science, not infrastructure

The primary advantage of building a scientific Python stack with Coiled is the shift of focus back to scientific or analytical value. Many teams report significant reductions in the time engineers spend managing infrastructure, debugging environments, or adapting workflows to unfamiliar distributed systems.

Domain experts can iterate on analysis code, experiments, or modeling approaches without waiting for IT or DevOps support. New projects or collaborations can begin immediately, without lengthy setup or onboarding. Teams can experiment more freely, attempting analyses or rescaling projects that would previously have been impractical due to technical hurdles.

Example: Mapping tropical forests at planetary scale

Space Intelligence, a 60-person company, tackles one of the most computationally demanding geospatial challenges today: creating detailed maps of the world’s tropical forests to support their protection and restoration. To map a country like Brazil, they process hundreds of terabytes of multi-spectral satellite data at 10-meter resolution across multiple time periods.

Their workflow combines data acquisition, dimension reduction, machine learning classification, and post-processing—all built on the Python geospatial stack including Xarray, Zarr, Cloud-Optimized GeoTIFFs, and Scikit-learn, orchestrated through Dagster.

Initially, the team tried local processing on laptops, then moved to AWS Lambda functions, and eventually built their own Kubernetes-based Dask infrastructure. However, their DIY approach suffered from significant inefficiencies. Up to half their compute usage was wasted due to cluster fragmentation, with memorable incidents like keeping 1,000 CPUs running for three days due to a single long-running job.

After transitioning to Coiled, Space Intelligence found the setup “phenomenally easy” and can now process continental-scale data using up to 300 workers. The platform’s simple VM-based architecture eliminated the fragmentation issues they experienced with Kubernetes, while integrated monitoring provided the high-quality metrics their team needed.

Most importantly, Coiled’s expert support extended beyond just the platform itself, helping optimize their broader Python geospatial stack and debug performance issues across the entire ecosystem. This shift allowed the team to focus on their core mission: building detailed forest maps that power conservation efforts and carbon credit programs in the fight against climate change.

Business and scientific impact

The capabilities provided by Coiled have downstream implications beyond individual project acceleration. Companies can prototype and move to production faster, testing new hypotheses or products in days instead of weeks or months. Engineering and research time is freed up to focus on core innovation. Access to scalable computing means organizations are no longer constrained by prior infrastructure investments or organizational silos.

In highly competitive domains—whether mineral exploration, pharmaceuticals, or climate analytics—this capability translates to actionable insights, smarter resource allocation, and speed to impact in the market.

Conclusion

Coiled makes large-scale, distributed scientific computing feel as natural as local development. By abstracting complexity, ensuring environment consistency, and embedding intelligent cost controls, teams can step directly from ideas to execution at any scale.

As the world’s most pressing challenges demand larger and more complex analyses, the ability to build scalable scientific Python stacks without friction becomes essential. Coiled provides this foundation, ensuring computational power is available where it’s needed most.

FAQs

What is scientific Python?

Scientific Python refers to a collection of Python libraries and tools specifically designed for scientific computing, data analysis, and visualization. At its core are libraries like NumPy and SciPy for numerical computations, Pandas and Xarray for data manipulation, Matplotlib and Plotly for visualization, and libraries like Scikit-learn, PyTorch, and TensorFlow for machine learning. These tools collectively enable researchers, data scientists, and analysts to perform complex scientific computations efficiently.

How is Python used in science?

Python is used extensively across various scientific disciplines. Scientists use it for data analysis, simulation, modeling, and visualization. For example, climate scientists use Python to build multi-model ensembles for temperature forecasting, biomedical researchers analyze RNA sequencing data, and geospatial analysts process satellite imagery. Python’s flexibility allows it to scale from simple analyses on local workstations to complex distributed computations on cloud platforms using tools like Dask and Coiled, making it suitable for both small experiments and large-scale scientific projects.

Does NASA use Python?

Yes, there are a number of different teams at NASA, like the National Snow and Ice Data Center, that extensively use Python. Python plays a crucial role in processing and analyzing the vast amounts of data collected from space missions, telescopes, and satellites. NASA scientists use Python libraries for tasks ranging from orbital mechanics calculations to image processing of planetary and deep space observations. Python’s accessibility and robust scientific computing capabilities make it valuable for NASA’s research teams working on complex problems in astrophysics, planetary science, and space mission planning.

Marimo on Coiled

2025-08-21T00:00:00+00:00

Aug 21, 2025

5m read

Nat Tabris and Alexandre Chabot-Leclerc

Coiled makes it easy to take anything you can run locally and run it in the cloud. That way, you have access to more machines, bigger machines, or machines that are close to your cloud data.

Many applications have an architecture where you run a local server, and then use the application in your web browser. Coiled supports this type of interactive use case with coiled run --interactive.

In this post, we’ll walk through how to use this to run marimo on a cloud VM.

For non-interactive work, or for work that you want to run in parallel across many cloud VMs, we recommend using coiled batch run.

Running in Docker on Your Local Computer

First, we’ll see one way to run marimo locally using Docker.

You could use a Docker image that already has marimo installed, but we’ll use a small uv container and uvx to install marimo on the fly. It’s a great way to get started quickly, and it gives you a lot of flexibility. Also, these uv progress bars are addictive.

We can run a marimo notebook server locally via Docker using a uv container like this:

docker run --rm -it -p 2718:2718 \
  astral/uv:debian-slim \
  uvx marimo new --port 2718 --headless --host 0.0.0.0

This command will print a URL to access your marimo notebook in a web browser. Any computations we execute in that notebook will run on our local computer inside the Docker container that’s running marimo.

When we run marimo, we use --port 2718 to tell it which port to use for the web server; this needs to match the port we expose on the Docker container. We use --headless to tell marimo not to try automatically opening a browser window (which won’t work from inside Docker). Finally, we use --host 0.0.0.0 so that marimo doesn’t need to know the IP address that it’s serving on (which will be different inside and outside Docker).

Running in the Cloud

Now that we know how to run marimo locally in Docker, it’s easy to run the same code (or anything else) in the cloud using Coiled.

Why run marimo in the cloud? When you’re running on a cloud VM, any compute you execute in the notebook will run on that cloud VM, not your local machine. The cloud gives you easy access to powerful hardware—big CPUs, lots of memory, or GPUs—and makes it easy to run computations close to data stored in the cloud, which can significantly reduce time and cost for moving data between the cloud and your local machine.

In fact, marimo recently launched molab, a free, hosted notebook workspace that’s fantastic for quickly trying things out without any setup. If you need more flexibility—running on your own cloud, choosing any instance type, or using GPUs—Coiled can provide it.

As before, we’ll use the astral/uv:debian-slim container with uvx to install marimo on the fly.

Instead of docker run, we’ll use coiled run, with just a few small tweaks to the arguments. Here’s the full command:

coiled run -it --port 2718 \
  --container astral/uv:debian-slim \
  -- uvx marimo new --port 2718 --headless --host 0.0.0.0 \
    --proxy \$COILED_CLUSTER_HOSTNAME:2718

There’s no need for the --rm flag that we used with Docker to clean up the container on exit. On Coiled, when you exit marimo (either in the browser or by hitting Control-C in your terminal), Coiled will automatically shut down the VM and clean up.

For the port, we specify a single one with --port 2718. Docker needs to know which internal port to map to which external port, but Coiled only needs to know which port to open in the VM’s firewall.

We’re explicitly specifying to use a container with the --container option because, unlike Docker, coiled run supports other ways to ship your software environment to the cloud.

For the marimo command, we add --proxy \$COILED_CLUSTER_HOSTNAME:2718 so that marimo prints the correct URL for our cloud VM.

That’s it! Once Coiled prepares your cloud VM, you’ll get a link to use Marimo in your web browser:

Launching marimo with Coiled from a terminal.

To shut things down, you can either use the shutdown button in the notebook or hit Control-C in your terminal. Once the notebook exits, Coiled will automatically shut down the VM and clean up everything in the cloud.

Sync Files

To make the experience even smoother, we can sync our local files with the VM where marimo is running with the --sync flag. Coiled will sync files in the current directory with the directory where marimo is running on the VM. That way, the files you edit on the cluster will still be available locally when you shut down the VM. So here’s how we’d sync our local files:

coiled run -it --port 2718 --sync \
  --container astral/uv:debian-slim \
  -- uvx marimo new --port 2718 --headless --host 0.0.0.0 \
    --proxy \$COILED_CLUSTER_HOSTNAME:2718 

Install Packages Automatically

We could use Coiled’s support for automatic package synchronization to install our local Python packages onto the cloud VM—it’s very convenient if you have a local-only package or an editable install[1]—but it’s not necessary here. Marimo automatically detects the packages required and offers to install them. Because uv is so fast, it’s fine to install them on the fly. (If you’re looking for those uv progress bars, they’ll be in your terminal, not in the marimo UI.)

Run marimo on a VM With a GPU

Coiled lets you pick any VM type your cloud provider offers, which means it’s easy to run marimo on a VM with a GPU. You can specify a VM type by name or specify --gpu to get the reasonable default machine with a GPU (g4dn.xlarge for AWS right now). Coiled (and uv in this case) will install the appropriate drivers and libraries to match your hardware.

coiled run -it --port 2718 --sync --gpu \
  --container astral/uv:debian-slim \
  -- uvx marimo new --port 2718 --headless --host 0.0.0.0 \
    --proxy \$COILED_CLUSTER_HOSTNAME:2718

Launching with the --gpu option gives you a VM with CUDA-compatible GPU.

Wrapping Up

We used coiled run, as the replacement for docker run, to launch marimo in the cloud using a uv container. Coiled’s capabilities make the cloud feel like an extension of our local development environment. Combining coiled run with uv, is a powerful pattern to execute any Python application on the cloud, like OpenLLM.

Is AWS Batch right for data professionals?

2025-08-19T00:00:00+00:00

Aug 19, 2025

8 min read

AWS Batch is a powerful service for batch computing workloads in the cloud. It offers sophisticated job scheduling, automatic resource provisioning, and seamless scaling capabilities that make it an excellent choice for organizations with dedicated cloud infrastructure teams.

However, if you’re a data scientist, quantitative analyst, or research scientist who wants to focus on analysis rather than infrastructure management, you might find AWS Batch requires more cloud operations expertise than you’d prefer to develop and maintain.

AWS Batch excels when you have experienced DevOps professionals managing your cloud infrastructure. But for data professionals who want to spend their time on models, analysis, and insights rather than VPC configurations and IAM policies, the operational overhead can be substantial. Let’s explore what this means in practice.

Navigating infrastructure setup and configuration

AWS Batch provides comprehensive control over your compute infrastructure, which is excellent for organizations with dedicated cloud teams. However, even with Fargate’s managed infrastructure, setting up your first job involves several AWS services that may be unfamiliar to data professionals:

Creating and configuring VPCs and subnets
Setting up security groups and network access controls
Defining IAM roles and policies
Configuring compute environments, job queues, and job definitions
Specifying detailed resource requirements and networking parameters

For data scientists accustomed to running analyses locally, this infrastructure-first approach represents a significant shift in workflow. While this level of control enables powerful use cases, it can feel like a steep learning curve when you just want to parallelize a data processing task.

Understanding networking requirements

AWS Batch’s networking model is designed for strict security requirements, but this can create unexpected challenges for data professionals. A common example: jobs may fail with timeout errors like failed to resolve ref docker.io/library/amazonlinux:latest: dial tcp 44.208.254.194:443: i/o timeout.

The root cause is often a networking configuration detail: Fargate tasks require "assignPublicIp":"ENABLED" to access the internet for pulling Docker images. While this security-first approach makes sense in enterprise environments, it can be surprising for data scientists who expect their containers to “just work” like they do in local Docker environments. Additionally, fixing such configuration issues often requires recreating multiple AWS resources, making iteration slower than many data workflows expect.

Managing complex IAM permissions

AWS Batch implements security through detailed IAM permissions, which provides excellent control for organizations with security teams. However, this creates a complex permission model where different components (compute environments, job definitions, and jobs) each require specific IAM configurations. For example:

Reading data from S3 requires configuring task roles, trust policies, and specific permissions like AmazonS3ReadOnlyAccess
Using GPU instances requires switching from Fargate to EC2, which introduces additional IAM requirements like instance profiles and ECS instance roles

For data professionals accustomed to file-based permissions or simpler cloud services, like AWS Lambda, this granular security model can feel like learning a specialized domain that’s separate from their core data expertise.

Debugging across multiple AWS services

AWS Batch’s distributed architecture means debugging information is spread across multiple AWS services, each with their own logging interfaces:

CloudWatch Logs for application output
EC2 system logs for infrastructure issues
Container logs for runtime problems
AWS Batch service logs for job orchestration

There’s no centralized, easy-access log aggregation. This distributed logging model aligns with AWS’s microservices architecture and provides comprehensive observability for DevOps teams familiar with AWS tooling. However, for data professionals accustomed to seeing all their program output in one terminal or notebook, this can feel fragmented. Job failures at different stages (provisioning, container startup, or execution) require checking different log sources, which can slow down the iterative debugging process that’s common in data work.

Scaling and GPU considerations

AWS Batch’s scaling behavior and GPU support are optimized for large-scale, production batch workloads, but may not align with typical data science workflows:

Cold-start provisioning for EC2-based environments can take 5-10 minutes or more, which is might be reasonable for long-running batch jobs but can feel slow during iterative development
GPU workloads require EC2 compute environments rather than Fargate, necessitating additional infrastructure configuration
AWS service quotas (like GPU instance limits) can cause jobs to remain in RUNNABLE status without clear error messages, requiring familiarity with AWS quota management

These characteristics work well for planned, production batch workloads but can create friction during the experimental, iterative phases common in data science projects.

Container-first approach vs. script-based workflows

AWS Batch follows a container-first approach. Rather than directly submitting scripts, you embed commands in job definitions or package them into Docker images. For more complex workflows, the typical pattern involves uploading scripts to S3 and configuring job definitions to retrieve and execute them.

This containerized approach provides excellent reproducibility and environment isolation, which is valuable for production workloads. However, data professionals often prefer the immediacy of script-based workflows and being able to quickly modify a Python file and re-run it. The container-first model introduces additional steps (building images or managing S3 uploads) that can slow down the rapid iteration cycles common in exploratory data analysis.

For data professionals operating without a dedicated cloud infrastructure or devops team, this can be a prohibitive bottleneck.

DevOps integration requirements

AWS Batch integrates well with AWS-native tools like CloudFormation and AWS CLI, making it a natural fit for teams already invested in the AWS ecosystem. However, integration with popular third-party tools can require additional effort:

Terraform support exists but may require custom scripting for complex configurations
Automated deployment of job definitions works best with AWS-native CI/CD tools, though integration with other systems is possible with additional setup

For organizations with established AWS DevOps teams, these integrations are manageable. However, for data teams using tools like Dagster, Prefect, GitHub Actions, or other non-AWS tooling may find the integration requires more DevOps expertise than they are able to maintain.

When AWS Batch works well

AWS Batch excels in the right organizational context. Teams find success when they have:

Dedicated DevOps or cloud infrastructure teams to handle setup and ongoing maintenance
Established AWS expertise and infrastructure-as-code practices
Already fully embedded in the AWS ecosystem
Long-running, production batch workloads that benefit from AWS’s configurability
Complex security and compliance requirements that benefit from fine-grained IAM controls

Organizations with these characteristics often appreciate AWS Batch’s comprehensive feature set and deep AWS integration.

Alternatives for data-focused teams

For data professionals who want to focus on analysis rather than infrastructure, several alternatives offer simpler approaches:

Kubernetes jobs provide more transparent control and ecosystem integration, though they still require container expertise
AWS Lambda works well for lightweight, short-duration batch tasks, though with execution time and memory limitations
Data-focused platforms like Coiled are designed specifically for data science workflows, offering the scale of cloud computing with the simplicity of local development

These alternatives prioritize different trade-offs. While AWS Batch optimizes for enterprise control and AWS integration, data-focused platforms like Coiled optimize for researcher productivity and rapid iteration. See our blog post on Choosing an AWS Batch alternative for a detailed comparison.

Batch jobs with Coiled

If you’re a data professional looking for cloud-scale computing without infrastructure complexity, Coiled offers a different approach. Designed specifically for data science workflows, Coiled automates the provisioning of EC2 instances in your AWS account, synchronizes your Python packages, and handles cleanup when you’re done. It bridges the gap between local development and cloud scale, letting you go from interactive notebook testing to distributed computing without building Docker containers or managing Kubernetes.

Importantly, you don’t have to sacrifice security or production-ready scale with this approach. Coiled works equally well for iterative development and large-scale production jobs, meaning your entire team can use one tool throughout the development lifecycle. While researchers can “package sync away” during exploration and prototyping, the same platform seamlessly handles production workloads with Docker containers and enterprise security controls. This unified approach eliminates the friction of switching between development and production tools, allowing teams to move much faster from experimentation to deployment.

With Coiled, you can:

Launch quickly: Spin up VMs in minutes without configuring VPCs or IAM policies
Scale transparently: Adjust resources automatically with clear feedback on provisioning status
Save with Spot instances: Use cost-effective Spot instances with automatic replacement handling
Access GPUs simply: Enable GPU acceleration without switching compute environments or managing quotas
Debug intuitively: View centralized logs and performance metrics in familiar interfaces

To run a batch job, add #COILED comments to your script to specify the cloud resources you want:

Bash Script

Spin up ten cloud VMs with 32 GB of memory to run their own echo command.

my_script.sh

#!/bin/bash

#COILED ntasks 10
#COILED memory 32GB
#COILED container ubuntu:latest

echo Hello from $COILED_BATCH_TASK_ID

Then launch your script with coiled batch run:

$ coiled batch run my_script.sh

COILED_BATCH_TASK_ID is an identifier unique to each task which runs from “0”, “1”, “2”, …, “9” in this case.

Python Script

Use # COILED comments directly in Python scripts. Drop the container directive to use package sync to automatically replicate your Python environment on remote VMs.

my_script.py

# COILED n-tasks     10
# COILED memory      8 GiB
# COILED region      us-east-2

import os

print(f"Hello from {os.environ['COILED_BATCH_TASK_ID']}")

and then launch your script with coiled batch run:

$ coiled batch run my_script.py

It’s easy to get started with Coiled:

$ pip install coiled
$ coiled quickstart

FAQs

What is AWS Batch used for?

AWS Batch is a managed service for running batch computing workloads in the cloud. It automatically provisions compute resources, queues jobs, and scales up or down based on demand. It’s often used for data processing, scientific computing, ML training, and other large-scale, parallelizable workloads.

Why might AWS Batch feel complex for data professionals?

AWS Batch is designed as an enterprise-grade service that requires configuring multiple AWS components (VPCs, subnets, security groups, IAM roles, and job definitions) before running jobs. This comprehensive approach provides control and security, but can feel like a significant learning investment for data professionals who want to focus on their analysis rather than infrastructure management.

Does AWS Batch support GPUs?

Yes, AWS Batch supports GPUs through EC2 compute environments. GPU support requires using EC2 rather than Fargate, which involves additional configuration of IAM roles, instance profiles, and ensuring adequate GPU service quotas are in place.

How do I submit scripts to AWS Batch?

AWS Batch does not natively accept standalone scripts. Commands must be embedded in job definitions as JSON or included in container images. For multi-line scripts, the typical workflow is to upload the script to S3, create a job role with S3 read access, and configure your job to pull and run the script.

What are common AWS Batch considerations for data teams?

Data professionals often encounter these learning areas when adopting AWS Batch:

Networking configurations that affect container internet access
IAM permission models that differ from file-based systems
AWS service quotas that can affect job scheduling
Cold-start provisioning times that may impact iterative workflows
Distributed logging across multiple AWS services

What are alternatives to AWS Batch?

Different alternatives optimize for different use cases:

Kubernetes jobs offer more transparent control and ecosystem integration
AWS Lambda or ECS provide simpler container-based execution models
Data-focused platforms like Coiled specialize in Python/data science workflows, emphasizing ease of use and rapid iteration over enterprise infrastructure controls

Choosing an AWS Batch alternative

2025-08-05T00:00:00+00:00

Aug 05, 2025

6 min read

AWS Batch is powerful, but not always the right fit.

AWS Batch is great when you need a fully managed job orchestration system for running containerized workloads at scale. It shines in production environments where DevOps teams are already managing infrastructure, and where strict job queueing are critical to large-scale batch pipelines.

But, for many data practitioners, it’s overkill. Many people come across AWS Batch when they’ve hit the limits of AWS Lambda and quickly find it’s not nearly as user-friendly. Others are looking for an easy way to run their script on the cloud, but then got stuck configuring job queues and troubleshooting opaque errors in CloudWatch. Then, once you’ve figured out the error, it can take several minutes to rebuild Docker images and try again.

In this post, we’ll cover some of the challenges of AWS Batch and provide some alternatives.

What makes AWS Batch challenging?

AWS Batch is powerful, but difficult to use for a number of reasons, including:

Complex setup: You need job definitions, compute environments, and container specs, even for simple jobs.
Opaque Debugging: When jobs fail, logs are spread across CloudWatch with minimal helpful context.
Hard to use Spot: AWS Batch can get expensive quickly, and to save on costs Amazon recommends using EC2 Spot. This requires additional management to do well, though, including picking the right allocation strategy, configuring subnets for multiple availability zones, making sure your job isn’t running for too long, and ensuring you have automated retries set up.
Quota limits: One of the more common job stuck in RUNNABLE status culprits is if you don’t have the right quota for the resources you want to use (this could be GPUs or Spot or even an instance type with more memory).

In the sections below we’ll go into getting started with AWS Batch on Fargate and on EC2.

AWS Batch with Fargate

You’ll first need to decide whether to use AWS Batch with Fargate or EC2. Amazon claims Fargate is their easier, serverless option. There is still a lot of setup required, though, including:

Find default VPC, subnets
Identify security groups
Create required IAM roles
Create a Fargate compute environment with public subnets
Create a job queue linked to your compute environment
Register a job definition

All before you can even submit your job. Once you’re ready to actually submit your job, you can’t actually submit a script. Instead, you’ll need to do one of the following:

Put your script as a JSON argument to Batch (this gets pretty awkward if your script is more than a couple lines)
Build it into the Docker image (OK for production, but slow for development)
Upload your script to S3 (and create and attach the correct IAM roles and access)

Which is a lot of overhead, especially if all you want is to run a script on the cloud. After all that, you also can’t use Fargate with GPUs, so you might instead try EC2.

AWS Batch with EC2 for GPU access

If you want to use AWS for GPUs, you’ll need to use EC2. That means you’ll first have to go through the EC2 setup, which includes:

Finding (or creating) your VPC, with subnets and security groups
Using (or creating) VPC infrastructure (you may need to create this if you want to use a different region from the default)
Creating the EC2 compute environment with your GPU instances

Then, you’ll be able to start setting up AWS Batch:

Create Batch EC2 compute environment with GPU instances
Create GPU job queue
Register GPU job definition (with your image, GPU resource requirement, and any additional memory your job will need)

After that, you’ll be able to try to run your script.

What to look for in an AWS Batch alternative?

All of that setup can be a lot of overhead if all you need to do is run a script on the cloud. AWS Fargate and AWS Lambda are both easier to use compared to AWS Batch, but come with their own tradeoffs like not being able to use GPUs and container image size limits.

Feature	AWS Batch	AWS Fargate	AWS Lambda	Coiled
Max execution time	No limit	No limit	15 minutes	No limit
Setup complexity	High	Medium	Low	Low
Container image size limit	No limit (w/ EC2)	20GB (200GB w/ extra config)	10GB	No limit
Containerization required	Yes	Yes	No	No (package sync or Docker)
Built-in parallelism	✅ Yes	❌ No	❌ No	✅ Yes (Dask, Coiled Batch, Coiled Functions)
Interactive development	❌ No	❌ No	❌ No	✅ Yes
Spot instances	✅ Yes (requires extra config)	✅ Yes	❌ No	✅ Yes
GPU support	✅ Yes (w/ EC2)	❌ No	❌ No	✅ Yes
Scale-to-zero	✅ Yes	❌ No	✅ Yes	✅ Yes

Coiled is an easier alternative to these AWS tools, making it easy to go from your laptop to the cloud.

Coiled as an AWS Batch alternative: An easy way to scale scripts on the cloud

Coiled is an easier alternative to running your workflows on the cloud. It automates the provisioning of EC2 instances in your AWS account, automatically synchronizes your Python packages, and turns things off when you’re done. It’s easy to go from testing interactively in a notebook to scaling out to cluster of VMs, without needing to build Docker containers or manage Kubernetes. You can easily take advantage of heavily discounted Spot instances (most computations end up costing ~$0.10 / TB).

It makes everything so much easier. We can always use the fastest machines. We don’t have to worry about infrastructure. In production we use Docker, but for the researchers we let them just package sync away and not worry about it. It’s really easy to parallelize things that might have previously taken two or three days.

Nelson Griffiths, Head of Engineering at Double River Investments

There are a few ways to use Coiled, but Coiled Batch is the most natural AWS Batch alternative.

Coiled Batch for single jobs to large-scale parallelism

Coiled Batch is an easy way to run any script on the cloud and is especially useful for scaling large, independent tasks (ie, embarrassingly parallel workloads) across hundreds of VMs. Unlike other Coiled APIs, it doesn’t rely on Dask for parallelism, which means you can run any code (not just Python).

You can use Coiled Batch to:

Run an arbitrary bash script 100 times in parallel
Reproject thousands of satellite images with GDAL for $1
Train a PyTorch model on a GPU for $0.12
Stress-testing S3 concurrency
Multi-node GPU LLM fine-tuning

To run a batch job, add #COILED comments to your script to specify the cloud resources you want:

Bash Script

Spin up ten cloud VMs with 32 GB of memory to run their own echo command.

my_script.sh

#!/bin/bash

#COILED ntasks 10
#COILED memory 32GB
#COILED container ubuntu:latest

echo Hello from $COILED_BATCH_TASK_ID

Then launch your script with coiled batch run:

$ coiled batch run my_script.sh

COILED_BATCH_TASK_ID is an identifier unique to each task which runs from “0”, “1”, “2”, …, “9” in this case.

Python Script

Use #COILED comments directly in Python scripts. Drop the container directive to use package sync to automatically replicate your Python environment on remote VMs.

my_script.py

#COILED n-tasks     10
#COILED memory      8 GiB
#COILED region      us-east-2

import os

print(f"Hello from {os.environ['COILED_BATCH_TASK_ID']}")

and then launch your script with coiled batch run:

$ coiled batch run my_script.py

Under the hood Coiled will:

Inspect your script
Spin up appropriate machines as defined by #COILED comments
Download your software onto them or use the container you specified
Run your script
Shut down the machines

Next steps

It’s easy to get started with Coiled:

$ pip install coiled
$ coiled quickstart

Learn more:

Long-running Lambda workloads: Challenges and alternatives

2025-07-16T00:00:00+00:00

Jul 16, 2025

6 min read

AWS Lambda excels at handling short-lived, event-driven tasks, though it quickly runs into limits when your workload needs more than 15 minutes, more memory, or specialized hardware like GPUs. For data practitioners working with large datasets or long-running processes, these constraints can turn a powerful tool into a frustrating bottleneck. In this article we’ll explore:

Common AWS tools for running long-duration tasks such as Step Functions, AWS Batch, and running containerized workflows with Fargate.
How alternative options, like Coiled, offer a simpler, more scalable option for running Python in the cloud without the overhead of maintaining cloud infrastructure or managing K8s.

The challenge of long-running Lambda tasks

AWS Lambda is designed for short-lived, event-driven functions. And for these for small, event-driven tasks it works really well. This includes things like responding to API Gateway requests, parsing log files, and other short-running Python scripts. Because it’s so lightweight and easier to use than other tools in the AWS ecosystem, it’s not uncommon for data practitioners to find themselves in an anti-pattern where they begin to use AWS Lambda for jobs it wasn’t really designed for. This includes things like:

Training long-running machine learning models
ETL pipelines that churn through large volumes of data
Simulations and heavy computations that demand memory, CPU, and time
Any other jobs that exceed the 15-minute timeout, when you’ve likely hit the AWS Lambda Task timed out after 900 seconds error message

At this point, it’s probably time to use something else.

Better alternatives to AWS Lambda for long-running tasks

In this post, we’ll go through a number of other options for long-running tasks on AWS:

AWS Fargate for running containerized workloads
AWS Step Functions for chaining together steps that involve a mix of short- and long-running processes (including retry logic, branching, waiting).
AWS Batch for long-running, compute-heavy batch jobs.
Coiled as an alternative to any of the above AWS services for easily running Python (or any language, really) on the cloud

Fargate for serverless containers

Fargate lets you run containers on demand. Similar to AWS Lambda, Fargate abstracts away infrastructure management. You define the CPU, memory, and runtime, and AWS spins up the infrastructure. It’s great for when you’ve already containerized your workloads and they need more than 15 minutes to complete.

Pros:

No servers to manage
Supports tasks of arbitrary length
Integrates well with ECS, EKS, and AWS Batch
Option to scale up/down

Cons:

Slightly more setup than Lambda
Cold starts can still be an issue for short tasks
If you’re not already using Docker, this can add some additional complexity
No GPU support (you’d need to use EC2)

For example, let’s say a data team needs to process a large hourly batch of raw JSON logs, apply transformation logic, and write to a data lake. The entire workflow runs in ~40 minutes. Instead of trying to chunk the job into multiple Lambdas, they could use Fargate to run a single container.

Step Functions for workflow orchestration

Step Functions let you coordinate multiple AWS tools together when logic, retries, and branching are needed.

Pros:

Visual state machine interface
Built-in retries, waits, error handling
Works with both short and long steps

Cons:

Limited to AWS ecosystem
Costs can climb for high-frequency workflows (pricing is based on state transitions)

For example, let’s say you have a data pipeline that extracts data from an API, cleans and joins multiple sources, runs a forecasting model, and stores the results. You could use Step Functions to orchestrate this across Lambdas for light steps and Fargate tasks for heavy lifting.

AWS Batch for large-scale compute jobs

AWS Batch is designed for submitting, scheduling, and executing a large volume of jobs. It’s often used for compute-heavy or long-running workloads like scientific simulations or model training. It’s especially useful for array jobs, or when you need to run a bunch of independent jobs in parallel with the same definition. It runs on top of Fargate or EC2. Unlike most other AWS services, you don’t pay an extra cost for using Batch.

Pros:

Fully managed job scheduling and provisioning
Use Spot to reduce cloud costs (but need to manually manage with retries)
Powerful tool for embarrassingly parallel workloads

Cons:

More complex to set up than Lambda or Fargate
Inefficient for low-latency tasks (partially due to scheduler overhead)
Long cold-starts for large container images

A research group might use AWS Batch to run daily simulations that each take 3–5 hours. They can submit hundreds of jobs at once, automatically scaled across hundreds of VMs, to run these in parallel.

Using Coiled instead of AWS services for long-running Lambda workarounds

While AWS tools like Fargate and Batch are powerful, they often introduce friction for data practitioners. You typically need to write Dockerfiles, manage IAM roles, configure queues, and set up infrastructure. For teams used to working in Jupyter notebooks or Python scripts, this overhead can be a major blocker.

Coiled is an easier alternative to running Python on the cloud. It automates the provisioning of EC2 instances in your AWS account, automatically synchronizes your Python packages, and turns things off when you’re done. It’s easy to go from testing interactively in a notebook to scaling out to cluster of VMs, without needing to build Docker containers or manage Kubernetes.

Feature	AWS Lambda	AWS Fargate	AWS Batch	Coiled
Max execution time	15 minutes	No limit	No limit	No limit
Setup complexity	Low	Medium	High	Low
Container image size limit	10GB	20GB (200GB w/ extra config)	No limit (w/ EC2)	No limit
Containerization required	No	Yes	Yes	No (package sync or Docker)
Built-in parallelism	❌ No	❌ No	✅ Yes	✅ Yes (Dask, Coiled Batch, Coiled Functions)
Interactive development	❌ No	❌ No	❌ No	✅ Yes
Spot instances	❌ No	✅ Yes	✅ Yes	✅ Yes
GPU support	❌ No	❌ No	✅ Yes (w/ EC2)	✅ Yes
Scale-to-zero	✅ Yes	❌ No	✅ Yes	✅ Yes

In the following sections we’ll go through some ways you can use Coiled as an alternative to other AWS services.

Coiled serverless Python functions

If you have a Python function you’d like to run on the cloud, you can use the @coiled.function decorator. Behind the scenes, Coiled will spin up an EC2 instance in your AWS account, run your function, and then turn things off when you’re done.

More details in Coiled Functions docs.

Data processing

import coiled
import pandas as pd

@coiled.function(
    region="us-west-2",             # Run close to data
)
def process(filename):
    df = pd.read_parquet(filename)  # Read S3 data from EC2
    df = df[df.name == "Alice"]     # Filter data on cloud
    return df                       # Return subset

result = process("s3://my-bucket/data.parquet")  # Runs remotely
print(result)                                    # Runs locally

Model training

import coiled

@coiled.function(
    vm_type="g6.2xlarge",           # Run on a GPU instance
)
def train():
    import torch
    device = torch.device("cuda")
    ...
    return model

model = train()                    # Runs remotely

Parallelism

import coiled
import pandas as pd

@coiled.function(region="us-west-2")  # Run close to the data
def process(filename):
    output_filename = filename[:-4] + ".parquet"
    df = pd.read_csv(filename)
    df.to_parquet(output_filename)
    return output_filename

# result = process("s3://my-bucket/data.parquet")  # one file

results = process.map(filenames)   # many files in parallel
for filename in results:
    print("Finished", filename)

Coiled Batch for single jobs to large-scale parallelism

Coiled Batch is an easy way to run any script on the cloud, and is especially useful for scaling large, independent tasks (ie, embarrassingly parallel workloads) across hundreds of VMs. You can use Coiled Batch to:

Run an arbitrary bash script 100 times in parallel
Train a PyTorch model on a single GPU
Distributed model training across multiple VMs
Reproject thousands of satellite images with GDAL

To run a batch job, add #COILED comments to your script to specify the cloud resources you want:

Bash Script

Spin up ten cloud VMs with 32 GB of memory to run their own echo command.

my_script.sh

#!/bin/bash

#COILED ntasks 10
#COILED memory 32GB
#COILED container ubuntu:latest

echo Hello from $COILED_BATCH_TASK_ID

Then launch your script with coiled batch run:

$ coiled batch run my_script.sh

COILED_BATCH_TASK_ID is an identifier unique to each task which runs from “0”, “1”, “2”, …, “9” in this case.

Python Script

Use #COILED comments directly in Python scripts. Drop the container directive to use package sync to automatically replicate your Python environment on remote VMs.

my_script.py

#COILED n-tasks     10
#COILED memory      8 GiB
#COILED region      us-east-2

import os

print(f"Hello from {os.environ['COILED_BATCH_TASK_ID']}")

and then launch your script with coiled batch run:

$ coiled batch run my_script.py

Under the hood Coiled will:

Inspect your script
Spin up appropriate machines as defined by #COILED comments
Download your software onto them or use the container you specified
Run your script
Shut down the machines

Next steps

It’s easy to get started with Coiled:

$ pip install coiled
$ coiled quickstart

Learn more:

uv and Coiled for easy Python scripts on the cloud

2025-06-30T00:00:00+00:00

Jun 30, 2025

3 min read

James Bourbeau

The Python packaging world is (being generous) … vast. One of the more delightful recent developments is uv, a fast Python package and environment manager. uv has the features you’d expect from a Python package manager (project-specific environments, lock files, etc.), but the core reasons folks like uv so much is that it’s easy to use and really, really fast.

This post shows how to run Python scripts on the cloud using uv and Coiled, an easy-to-use, UX-focused cloud platform that provides a similarly straightforward approach to running Python scripts on any cloud hardware.

uv for managing software

Here’s a script that uses pandas to load a Parquet data file hosted in a public bucket on S3 and print the first few rows:

process.py

import pandas as pd

df = pd.read_parquet(
    "s3://coiled-data/uber/part.0.parquet",         
    storage_options={"anon": True},
)
print(df.head())

We’ll use this concrete example throughout the rest of the post, but know that the script contents could be any Python code.

Running this script requires pandas, pyarrow, and s3fs. uv has a pretty slick way for embedding these dependencies directly in the script using PEP 723 script metadata with the uv add --script command.

$ uv add --script process.py pandas pyarrow s3fs

which adds these comments with the specified dependencies to the script

process.py

# /// script
# requires-python = ">=3.12"
# dependencies = [
#   "pandas",
#   "pyarrow",
#   "s3fs",
# ]
# ///

import pandas as pd

df = pd.read_parquet(
    "s3://coiled-data/uber/part.0.parquet",         
    storage_options={"anon": True},
)
print(df.head())

We then use uv run to run the script

$ uv run process.py

uv automatically creates a virtual environment, installs the dependencies, and then runs process.py in that environment. Here’s the output we get:

hvfhs_license_num dispatching_base_num originating_base_num    request_datetime   on_scene_datetime  ... shared_request_flag shared_match_flag  access_a_ride_flag  wav_request_flag  wav_match_flag
__null_dask_index__                                                                                                      ...
18979859                       HV0003               B02875               B02875 2019-05-26 23:29:35 2019-05-26 23:30:33  ...                   N                 N                 NaN                 N             NaN
18979860                       HV0003               B02875               B02875 2019-05-26 23:56:48 2019-05-26 23:57:03  ...                   N                 N                 NaN                 N             NaN
18979861                       HV0003               B02765               B02765 2019-05-26 23:56:35 2019-05-26 23:56:41  ...                   N                 N                 NaN                 N             NaN
18979862                       HV0003               B02682               B02682 2019-05-26 22:52:34 2019-05-26 23:10:38  ...                   Y                 N                 NaN                 N             NaN
18979863                       HV0002               B03035               B03035 2019-05-26 23:16:34 1970-01-01 00:00:00  ...                   N                 N                   N                 N             NaN

[5 rows x 24 columns]

What’s nice about this is we didn’t have to think about managing local virtual environments ourselves and the dependencies needed are included directly in the script which makes things nicely self-contained.

That’s really all we need for running scripts locally. However, there are times when we need resources beyond what’s available on our local workstation, like:

Processing large amounts of cloud-hosted data
Needing accelerated hardware like GPUs or a big machine with more memory
Running the same script with hundreds or thousands of different inputs, in parallel

In these situations running scripts directly on cloud hardware is often a good solution.

Coiled for managing hardware

Similar to how uv makes it super straightforward to handle Python dependency management, Coiled makes it easy to handle running code on cloud hardware (if you’re not already familiar with how unergonomic cloud APIs can be, see this example on using uv with AWS Lambda).

Let’s start by installing the coiled Python package and authenticating whatever machine you’re running on with Coiled using the coiled login CLI (you’ll be prompted to create a Coiled account if you don’t already have one – and it’s totally free to start using Coiled)

$ uvx coiled login

To have Coiled run our script on a VM on AWS, we’ll add these two comments:

process.py

# COILED container ghcr.io/astral-sh/uv:debian-slim
# COILED region us-east-2

# /// script
# requires-python = ">=3.12"
# dependencies = [
#   "pandas",
#   "pyarrow",
#   "s3fs",
# ]
# ///

import pandas as pd

df = pd.read_parquet(
    "s3://coiled-data/uber/part.0.parquet",         
    storage_options={"anon": True},
)
print(df.head())

They tell Coiled to use the official uv Docker image when running our script (this makes sure uv is installed) and to run the script in the us-east-2 region on AWS (where this data file happens to live) to avoid any data egress. There are many other options we could have specified here like VM instance type (the default it an m6i.xlarge instance on AWS), whether to use spot instance, etc. See the Coiled Batch docs for more details.

Finally we use the coiled batch run CLI to run our existing uv run command on a cloud VM.

$ uvx coiled batch run \
    uv run process.py

The same exact thing that happened locally before now happens on a cloud VM on AWS, only this time the script is faster because we didn’t have to transfer any data from S3 to a local laptop.

View of Coiled UI with our script running on AWS.

Conclusion

uv is great. It’s a really pleasant experience for managing dependencies when running Python scripts. It also works well with Coiled when you want to run scripts on cloud hardware.

Together uv and Coiled make for a powerful, ergonomic cloud Python experience.

How to run a Python script on EC2

2025-06-25T00:00:00+00:00

Jun 25, 2025

5 min read

There are a lot of ways to run Python on AWS. You might use Lambda for event-driven tasks, Glue for managing data pipelines, or SageMaker for machine learning workloads. Each of these tools offers helpful abstractions, but also comes with constraints. For example:

Lambda is fast and simple, but capped at 15-minute execution time.
EMR can scale data processing across a cluster, but has limited support for custom libraries.
SageMaker is powerful for ML, but can be heavyweight for general-purpose Python scripts.

That’s why a lot of engineers will instead use EC2, which is highly configurable. You get you raw virtual machines anywhere in the world in about 30 seconds, fully under your control, for just pennies.

While AWS EC2 offers remarkable flexibility, it’s pretty low-level and not the most intuitive. Many data scientists and engineers avoid using it directly because it can feel tedious or overly complex. Worse, it’s often inefficient (like spinning up a machine and accidentally leaving it on all week). Coiled aims to address these issues, without sacrificing that flexibility.

Coiled makes it easy to run Python scripts on EC2 and scale out to run your jobs in parallel. It provisions machines on your behalf, replicates your local Python environment, and shuts everything down when the job finishes. You get the versatility of EC2, without the constraints of other AWS tools.

This post walks through two ways to run Python on EC2:

The manual route (for full control).
A Python-native method using Coiled.

Option 1: Run Python on EC2 manually

This process gives you full control over how and where your Python script runs in AWS, but it’s a bit involved. Here’s how to do it, correctly and safely.

Step 1: Launch an EC2 instance

Go to the EC2 Console: https://console.aws.amazon.com/ec2/
Launch Instance: Click “Launch Instance”.
Choose an AMI: Pick something like Ubuntu or Amazon Linux.
Select Instance Type: t3.medium is a good starting point. Use larger instance types for bigger workloads (Vantage has a nice EC2 comparison table).
Create a Key Pair (RSA): Download the .pem file (keep it private, you’ll need this for SSH access).
Configure Network Settings:
- Open SSH (port 22) in the security group only to your IP (not 0.0.0.0/0).
- Optionally, enable HTTP (port 80) or HTTPS (443) if your script serves a web app.
Launch the Instance

Step 2: Connect to your EC2 instance

Once your EC2 instance is running, click “Connect” to see the instructions for how to connect to your instance.

Ensure the .pem file has the right permissions:

chmod 400 your-key.pem

Ubuntu

ssh -i your-key.pem ubuntu@<your-ec2-public-ip>

Amazon Linux

ssh -i your-key.pem ec2-user@<your-ec2-public-ip>

Step 3: Set up your Python environment

Once you’re connected to the instance:

Ubuntu

The Ubuntu AMI should already have Python 3 installed, but if not, you’ll need to install it first.

# Update the system
sudo apt update && sudo apt upgrade -y

# Install Python
sudo apt install python3 python3-venv -y

# Set up a virtual environment
python3 -m venv env
source env/bin/activate

Amazon Linux

The Amazon Linux AMI should already have Python 3 installed, but if not, you’ll need to install it first.

# Update the system
sudo yum update -y

# Install Python
sudo yum install python3 -y

# Set up a virtual environment
python3 -m venv env
source env/bin/activate

Step 4: Transfer your Python script

From your local machine, use scp to upload your Python files to EC2:

Ubuntu

scp -i your-key.pem script.py ubuntu@<your-ec2-ip>:~

Amazon Linux

scp -i your-key.pem script.py ec2-user@<your-ec2-ip>:~

You can also copy folders or zipped environments if needed.

Step 5: Run your Python script

You should now be able to see the Python script you uploaded on your EC2 instance. You can list the files in your working directory to check with ls -l. Then, you can run the script:

python3 script.py

Step 6: Turn off your instance

Once you’re done, if you don’t turn off your EC2 instance, AWS will keep charging you. You can do this from the AWS console. It’s worth noting that if you have an attached EBS volume, AWS will continue to charge you for storage unless you terminate the instance.

AWS EC2 best practices

Use security groups to restrict SSH to your IP to control who can access your instances
Terminate EC2 instances when you’re done to avoid a surprise bill
Use S3 for data storage to decouple your data from ephemeral VMs
Since you’ll need to SSH into your remote VM, it’s helpful to at least know some basic shell commands

This manual approach offers full control but can be cumbersome, especially for scaling or handling complex dependencies.

Option 2: Use Coiled to run Python on EC2

Coiled is a Python-native platform designed to simplify running Python workloads on the cloud. It allows you to easily manage virtual machines with minimal code. Coiled automates the provisioning of EC2 instances in your AWS account, synchronizes your environment, executes your code, and ensures resources are efficiently managed. Coiled:

Eliminates the need complex cloud console navigation.
Replicates your local Python environment and packages.
Supports any hardware available on EC2 (GPUs, ARM, Spot, etc.)
Automatically shuts down instances after execution to save costs.
Offers support for custom VPCs, subnets, firewalls, and private IPs.

Step 1: Install Coiled

pip install coiled

Step 2: Configure AWS access

Coiled requires limited IAM permissions to provision resources in your AWS account. You can set this up using the Coiled CLI:

coiled setup aws

Step 3: Run a Python script on AWS with Coiled

my_script.py

#COILED n-tasks     10
#COILED memory      8 GiB
#COILED region      us-east-2

import os

print(f"Hello from {os.environ['COILED_BATCH_TASK_ID']}")

And then launch your script with coiled batch run:

$ coiled batch run my_script.py

Coiled makes it easy to scale your Python scripts to run on as many VMs as you need. In this example, we ran my_script.py in parallel on 10 EC2 instances. Under the hood Coiled will:

Inspect your script
Spin up the appropriate machines as defined by the #COILED comments
Download your software onto them or use the container you specified
Run your script, automatically scaling out as needed with #COILED n-tasks
Shut down the machines

Using batch jobs is one easy way to run a simple Python script, but there are a number of other options:

Serverless Functions

Like AWS Lambda, but better for long-running jobs or if you need GPUs. Decorate a function to run it in the cloud. Scale out with the .map method.

import coiled
import pandas as pd

@coiled.function(region="us-west-2")  # Run close to the data
def process(filename):
    output_filename = filename[:-4] + ".parquet"
    df = pd.read_csv(filename)
    df.to_parquet(output_filename)
    return output_filename

# result = process("s3://my-bucket/data.parquet")  # one file

results = process.map(filenames)   # many files in parallel
for filename in results:
    print("Finished", filename)

Dask Clusters

Like AWS EMR, but easier to use (plus, Dask is often faster than Spark).

import coiled
cluster = coiled.Cluster(
    n_workers=100,
)
client = cluster.get_client()

# Use Dask + Pandas together
import dask.dataframe as dd

df = dd.read_parquet("s3://bucket/lots-of-data.parquet")
df.groupby("name").amount.sum().compute()

Jupyter Notebooks

Run Jupyter on large cloud VMs. Synchronize files back to your local hard drive.

coiled notebook start --sync --vm-type m6i.16xlarge

When would you want to use EC2 manually vs. Coiled to run Python?

Coiled supports a lot, including custom VPCs and firewall rules, but there are some cases where manually running Python on EC2 might be better:

Custom OS or low-level system configuration
Custom AMI creation
Highly regulated cloud environments with air-gapped requirements
Advanced networking (beyond custom VPCs), like hybrid clouds

Next steps

It’s easy to get started with Coiled:

$ pip install coiled
$ coiled quickstart

Learn more:

Coiled: A GPU-accelerated serverless alternative to AWS Lambda

2025-06-11T00:00:00+00:00

Jun 11, 2025

4 min read

AWS Lambda is widely used for event-driven workloads that require quick, lightweight execution without infrastructure management. But when it comes to tasks that demand GPU acceleration, like training deep learning models, running inference on images, or performing parallel scientific simulations, Lambda falls short.

With no GPU support, a 15-minute execution limit, and container image size limits, AWS Lambda isn’t designed for compute-heavy or memory-intensive applications.

In this post, we’ll explore how Coiled provides a serverless Python experience with full GPU support, offering a flexible, Lambda-like alternative for GPU-accelerated computing in your own AWS account.

Why Lambda alone falls short

AWS Lambda was specifically designed for short-lived compute tasks. And for these for small, event-driven tasks it works really well. This includes things like:

Responding to API Gateway requests
Parsing log files
Short-running Python scripts

But it falls short for high-performance compute due to a number of limitations.

Feature	AWS Lambda	Coiled
GPU support	❌ No	✅ Yes
Spot instances	❌ No	✅ Yes
Max execution time	15 minutes	No limit, configurable
Container image size	Capped at 10 GB, uncompressed	No limit, package sync or Docker
Scale-to-zero	✅ Yes (by default)	✅ Yes (by default)
Memory	Capped at 10 GB	No limit, configurable
vCPU	Proportional to memory, not directly configurable	No limit, configurable

AWS Lambda imposes several architectural constraints that limit its suitability for GPU-based or high-performance workloads. Memory and vCPU are tied together, maxing out at 10 GB and proportional CPU allocation, with no GPU support. Container images must remain under 10 GB uncompressed, which is often too small for GPU-accelerated Python workloads that bundle frameworks like PyTorch, TensorFlow, CUDA, and cuDNN (not to mention any pre-trained model weights or scientific dependencies). By contrast, Coiled offers additional flexibility: define your software environment using conda or pip, use Docker images of any size, and deploy clusters with GPU-backed instances on demand, all without container size limits or infrastructure overhead.

Coiled: A serverless GPU alternative

With Coiled, you can run Python workloads on-demand in your AWS account with optional GPU support. Coiled spins up EC2 instances as you need them, replicating your local Python environment on remote VMs. With Coiled, you can:

Spin up GPU-backed clusters (e.g., g5.xlarge)
Run Torch, TensorFlow, or custom CUDA workloads
Use Spot instances for cost efficiency

Use cases for GPUs + Coiled

Deep learning inference at scale
Image/video processing pipelines
Scientific simulations (e.g., genomics)
ETL pipelines with GPU libraries
Parallel compute tasks in finance or physics

Step-by-step guide: Running GPU-accelerated workloads on Coiled

Install Coiled

Next, we’ll install Coiled. We’ll also include PyTorch, since this example will go through training a PyTorch model on a GPU.

pip

$ pip install torch torchvision coiled

conda

$ conda create -n env python=3.11 pytorch torchvision coiled
$ conda activate env

Train PyTorch model

Next, we’ll define our PyTorch model:

import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 1024, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(1024, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

Now that we have our model, let’s train it on our dataset on the hardware of our choosing. Fortunately, PyTorch makes it straightforward to utilize different types of hardware that are available locally.

Below is a train function that takes a PyTorch device as input (defaults to "cpu"), loads the CIFAR10 dataset, trains our model on the dataset on the specified hardware, and returns the trained model.

import torch
import torchvision
import torch.optim as optim
import torchvision.transforms as transforms


def train(device="cpu"):
    # Select hardware to run on
    device = torch.device(device)

    model = Net()
    model = model.to(device)

    # Load CIFAR10 dataset
    trainset = torchvision.datasets.CIFAR10(
        root='./data', train=True, download=True, transform=transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))]),
    )
    trainloader = torch.utils.data.DataLoader(
        trainset, batch_size=256, shuffle=True,
    )
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

    for epoch in range(10):
        print(f"Epoch {epoch + 1}")
        model.train()
        for batch in trainloader:
            # Move training data to device
            inputs, labels = batch[0].to(device), batch[1].to(device)

            optimizer.zero_grad()

            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

    return model.to("cpu")

Run on a cloud GPU

Our train function can now be run on any hardware that’s available locally. To train our model on a GPU, we’ll use Coiled Functions to lightly annotate our existing train function to run on a GPU-enabled cloud VM.

import coiled

@coiled.function(
    vm_type="g5.xlarge",   # NVIDIA A10 GPU instance
    region="us-west-2",
)
def train(device="cpu"):
    # Same training code as before
    ...

Now when the train function is run, Coiled will automatically handle provisioning a cloud VM (a g5.xlarge instance on AWS in this case), installing the same software that’s installed locally on the cloud VM, running our train function on the VM, and returning the result back locally.

model = train(device="cuda")  # Train model on cloud GPU

This takes ~4.5 minutes to run (including VM spinup time) and costs ~$0.08 (running locally would have taken ~1.5 hours).

GPU utilization throughout model training. Utilization is typically high throughout, meaning we’re utilizing the available hardware well.

How does Coiled compare to other serverless tools?

Tools like AWS Lambda, Google Cloud Functions, and Azure Functions are optimized for lightweight, fast-response tasks (such as web endpoints or trigger-based coordination), but they are less suited for heavy data processing or GPU-accelerated workloads. We also find these cloud provider-specific tools leave a lot to be desired in terms of developer experience and often require cloud provider-specific knowledge. See the full comparison for more details.

Tools like Modal and Coiled on the other hand, prioritize ease of use, making it easy for anyone to use GPUs on the cloud. Modal, for example, offers cold start times around one second and is an excellent choice for latency-sensitive applications. However, this comes with tradeoffs, like higher cost per compute unit (often 4–5x raw EC2 cost) and less flexibility in hardware and environment setup. Coiled functions, on the other hand, are cheaper at scale (raw EC2 instances are spun up as needed) and offer more security and flexibility (everything runs in your cloud account). For a deeper dive, read the full comparison on Coiled vs. Modal.

In short:

If you need sub-second response times for short-lived web requests, tools like Modal are a good fit.
If you’re doing large-scale, resource-intensive workloads, and have more specific security requirements, Coiled is a better fit.

Next steps

Serverless functions aren’t the only way to run GPU-accelerated workloads with Coiled. You can also use:

Batch jobs: Like AWS Batch, but a lot easier to use.
Jupyter notebooks: Run a Jupyter Notebook on AWS/GCP/Azure.
Distributed compute: Speed up and scale your Python workflows with Dask clusters.

Python is annoying to run on EMR, but you have options

2025-05-08T00:00:00+00:00

May 7, 2025

6m read

Amazon EMR (Elastic MapReduce) is a powerful managed cluster platform that helps organizations process large amounts of data. While AWS EMR offers significant performance benefits, its cost structure is complex and and many teams fail to consider the additional cost of developer time and energy.

In this article, we’ll cover:

The main drivers of EMR cost
Why Python teams face additional challenges with EMR
Coiled as a Python-native alternative for data processing workloads

Though EMR can run distributed data processing jobs with a number of open source libraries, this post will primarily focus on running Apache Spark on EMR, as that is the most widely used.

Architecture diagram of running an Apache Spark cluster on AWS EMR on EKS from the AWS Big Data Blog.

AWS EMR costs broken down

There are multiple factors to consider when assessing the prospective cost of AWS EMR for your analytics and data workloads. While AWS EMR uses a consumption-based pricing model, you’ll also need to factor in EC2 costs, storage, data transfer, and a number of additional hidden costs.

AWS EMR pricing model

AWS EMR follows a “pay-as-you-go” consumption-based pricing model where users pay for the resources they consume. Pricing varies based on the deployment mode and instance type selected.

Primary EMR deployment options include:

EMR on EC2: Users pay for both the EC2 instances and EMR service fees.
EMR on EKS (Elastic Kubernetes Service): Costs are based on Kubernetes resources used, including EC2 instances or Fargate compute charges.
EMR Serverless: Users are billed based on the number of compute units and memory usage per second.

EC2 instance costs

Since EMR clusters run on Amazon EC2 instances, compute costs are a significant portion of EMR expenses. Instance options include:

On-demand instances: Pay for compute capacity by the second, with no long-term commitments.
Burstable instances: Perfect for workloads like web servers which have low CPU utilization most of the time, but occasionally fully utilize the CPU.
Spot instances: Significant cost savings (up to 90%) by using spare AWS capacity, but they can be interrupted.
ARM instances: Using ARM (Graviton) instances can reduce costs for certain workloads, in some cases up to 50%.
Reserved instances: Provide discounts in exchange for committing to EC2 capacity for one to three years.

Storage and data transfer costs

Storage and data transfer also contribute significantly to AWS EMR costs:

Amazon S3 storage: Many EMR workloads rely on Amazon S3 for data storage. Costs depend on the storage tier (standard, infrequent access, or glacier).
Amazon EBS (Elastic Block Store): If EC2 instances use attached EBS volumes, storage fees apply based on provisioned capacity.
Data transfer fees: Transferring data within the same AWS region is free, but inter-region and external data transfers incur additional charges.

EMR is also hard to use

When evaluating AWS EMR, most organizations focus primarily on direct compute and storage costs, though there are significant hidden costs to consider:

1. Cluster management and DevOps overhead

Unlike fully serverless solutions, AWS EMR requires ongoing management and tuning to maintain optimal performance. Key challenges include:

Provisioning and scaling clusters
Cluster lifecycle management
Monitoring and logging complexity

2. Engineering costs for Python-based data science teams

For teams working primarily in Python, EMR introduces additional engineering overhead:

Tuning Spark performance (efficient memory allocation, partitioning, etc. require Spark expertise)
Inefficient debugging and development cycles (Spark clusters are notoriously time-consuming to debug)
PySpark serialization and rewriting code for Spark compatibility

3. Vendor lock-in risks

Committing to AWS EMR can create dependencies on AWS-specific features and services, making it difficult to migrate workloads to other platforms in the future.

Coiled: A Python-native alternative to AWS EMR

For teams working in Python, AWS EMR’s reliance on Spark and Hadoop adds unnecessary overhead, making development slower and more expensive. If you haven’t committed to EMR yet or want easier alternatives for new workloads, Coiled offers an alternative.

Why choose Coiled over AWS EMR?

Coiled provides a modern, Python-native alternative that’s easy to use, highly scalable, and cost-effective.

1. Python-native

Coiled is designed for Python data developers. For distributed computing Coiled leverages Dask, an open source parallel computing library. Dask feels a lot like pandas, making it easy for teams to scale their existing pandas code with minimal modifications.

EMR, by comparison, commonly relies on Spark for distributed computing. For data teams using primarily Python, migrating to Spark is notoriously challenging.

A lot of the Spark internals are run on the JVM, so it’s actually not that easy to debug. You spend a lot of time on Stack Overflow trying to figure out what did I do wrong, why am I getting this traceback. It was painful.

Sébastien Arnaud, CDO/CTO at Steppingblocks, on why they switched from Spark to Dask

Since Dask is just Python, you don’t have to waste time debugging JVM errors. Additionally, you can easily see metrics like CPU and memory utilization in the Coiled UI or detailed task-level progress with the Dask dashboard.

2. Zero infrastructure management

Coiled eliminates cluster management overhead by automatically handling provisioning, scaling, and maintenance of AWS resources. This means you no longer need to:

Wait for Docker images or build custom AMIs. With Coiled, software is automatically replicated on your cloud VMs with package sync.
Rightsize Kubernetes. Coiled relies on a simpler architecture, automatically spinning up raw EC2 instances when you need them in ~ a minute.
Worry about VMs left on over the weekend. Coiled automatically turn things off when your computation is done and sends out alerts if things have been running longer than usual.

For teams without a dedicated cloud infrastructure team, this frees up key engineers from infrastructure overhead. Instead of spending their time figuring out why EKS isn’t working, they can focus on their area of expertise.

3. Cost savings

The cloud can be quite cheap, ~$0.10 / TiB, if you’ve configured everything correctly. Coiled makes it easy to use the cloud at scale efficiently:

Auto-scaling based on workload demands
Trivial to use Spot, ARM
Read data in-region to avoid egress fees
Set team- and user-specific resource limits
Eliminate dead costs with auto-shutdown and clean-up once you’re done

There’s also a generous free tier of 500 CPU hours per month.

import coiled

cluster = coiled.Cluster(
    n_workers=[10, 100],               # Autoscale as needed
    spot_policy="spot_with_fallback",  # Use spot when available
    arm=True,                          # Use ARM
    region="us-west-2",                # Same region as your data
) 

4. No vendor lock-in

Coiled works with AWS, GCP, or Azure. Access your data wherever it’s stored (Snowflake, S3, Azure Blob, etc.) and use the same libraries you’re already using (Polars, PyTorch, Xarray, even GDAL).

S3 + Polars

Read in Parquet files from S3 with your preferred DataFrame library (pandas, Polars, DuckDB, Dask, …) with Coiled’s serverless functions.

import polars as pl
import coiled

@coiled.function(memory="512 GiB")
def run_query():
    df = (
        pl.scan_parquet("s3://my-bucket/*")
        .filter(pl.col("balance") > 0)
        .group_by("account")
        .agg(pl.all().sum())
    )
    return q.collect

Snowflake + Dask

With dask_snowflake, you can read from a Snowflake table into a Dask DataFrame.

import coiled
import dask_snowflake

cluster = coiled.Cluster(n_workers=50)
client = cluster.get_client()

example_query = """
    SELECT *
    FROM SNOWFLAKE_SAMPLE_DATA.TPCH_SF1.CUSTOMER;
"""
df = read_snowflake(
    query=example_query,
    connection_kwargs={
        "user": "...",
        "password": "...",
        "account": "...",
    },
)

When to still use EMR?

EMR still has its place, especially in certain Spark-heavy environments. Here are some reasons why you might stick with EMR:

You’re deeply invested in Apache Spark

EMR is optimized for Spark, and if your team already has Spark pipelines, libraries, and infrastructure, migrating everything to Dask might be time consuming.

You need tight AWS service integration

If you’re already embedded into the AWS ecosystem and using tools like AWS Glue and Step Functions, EMR fits in naturally.

Your team already has EMR expertise

If your data team already knows how to wrangle EMR clusters, it’s hard to justify switching unless there’s a real gain in productivity or cost.

Getting started with Coiled

It’s easy to get started with Coiled:

$ pip install coiled
$ coiled quickstart

Learn more:

Reducing Memory Pressure for Xarray + Dask Workloads

2025-03-17T00:00:00+00:00

Mar 17, 2025

3m read

Patrick Hoefler

Historically, Dask scheduled thousands of data-loading tasks at once, overloading the system. The new queuing mechanism prevents this, making Xarray workloads smoother and more reliable.

As of Dask 2025.1.0, Dask-backed Xarray workloads are more memory-efficient. Processing satellite images from Microsoft Planetary Computer, for example, is much smoother.

The Challenge: Overeager Scheduling of Data-Loading Tasks

Dask’s scheduler aims to distribute workloads efficiently, but data-loading tasks—such as reading from Zarr, NetCDF, or remote storage—have historically been scheduled too aggressively.

Previously, Dask would immediately schedule all available data-loading tasks, causing workers to fetch large amounts of data before they were ready to process it. This led to excessive memory use, frequent spilling to disk, and even crashes.

Dask introduced root task queuing in 2022 to mitigate this issue by identifying and scheduling these tasks more cautiously. However, the detection mechanism was fragile and often failed in cases like:

Opening multiple Zarr stores
Loading many NetCDF files
Accessing data from Microsoft Planetary Computer

This meant that for many Xarray workloads, the memory footprint was unpredictable, leading to inefficient resource usage.

Example of the Problem

In the following example, we open multiple Zarr stores and calculate moving averages:

import dask.array as da

zarr_paths = ["/tmp/1.zarr", "/tmp/2.zarr", "/tmp/3.zarr", "/tmp/4.zarr", "/tmp/5.zarr"]

arrs = [da.from_zarr(p) for p in zarr_paths]
arr = da.concatenate(arrs, axis=0)
weights = da.ones_like(arr, chunks=arr.chunks)
avg = da.sum(arr * weights, axis=(1, 2)) / da.sum(weights, axis=(1, 2))
avg.compute()

Tip

Run the following snippet to re-create the test dataset used in the example above:

import dask.array as da

zarr_paths = ["/tmp/1.zarr", "/tmp/2.zarr", "/tmp/3.zarr", "/tmp/4.zarr", "/tmp/5.zarr"]
for path in zarr_paths:
    arr = da.random.random((365, 1280, 2560), chunks=(10, -1, -1))
    da.to_zarr(arr, path)

We can see that Dask is quickly spilling to disk, because it’s overeagerly loading the source data.

Prior to Dask version 2025.1.0, loading multiple Zarr stores with Xarray caused memory issues. This is apparent in the Dask dashboard, with grey bars in the cluster memory and workers memory graphs indicating Dask is spilling to disk.

Before the recent improvements, Dask would load all chunks into memory at once, slowing down workloads and even crashing due to memory issues. This was especially troublesome for workloads creating multiple datasets from different Zarr stores, like in the example above, or when loading large NetCDF files.

The Solution: More Robust Root Task Detection

Dask 2025.01.0 introduces a more reliable mechanism for detecting and scheduling data-loading tasks. Dask explicitly marks data-loading tasks so the scheduler can manage them efficiently. Now, data is only loaded when it can be processed, keeping memory usage stable and preventing unnecessary disk writes.

With the fixes implemented in Dask 2025.1.0, the same example takes ~30 seconds to complete vs. over a minute. Data-loading tasks are queued rather than being scheduled immediately. There’s much less memory pressure on the cluster and Dask is no longer spilling to disk.

Running the same example as above, memory on the cluster remains fairly constant and reliably low enough that we don’t run into spilling issues. This particular example also runs twice as fast since Dask is no longer constantly writing to and reading from disk.

Real World Implications

This change has a number of real-world implications for Xarray users and users of related Python libraries. Most workloads opening NetCDF files will run into this issue one way or another at some point when scaling to larger workloads. Additionally, we’ve also seen this regularly occurring with users of satellite data through odc-stac. odc-stac implements a great interface to load satellite images, but the graph construction is prone to run into this problem as well.

Dask dashboard while processing satellite images loaded via odc-stac. Data-loading tasks in the bottom right are queued rather than being scheduled immediately. Previously, Dask would schedule tens of thousands of tasks at once, overloading the scheduler.

This improved queuing mechanism removes these issues and provides a smoother and more reliable user experience when working with large-scale datasets.

There is no user-side change required. The new behavior is picked up automatically after upgrading Dask.

conda

conda update dask xarray

pip

pip install --upgrade "dask[complete]" xarray

For other recent improvements to using Xarray with Dask, you might consider:

Coiled 2024 in Review

2025-01-02T00:00:00+00:00

Jan 2, 2025

12m read

Matthew Rocklin

Hello Users, Teammates, and General Well-Wishers!

It’s that time of year where companies issue year-end summaries, looking back on the past twelve months, acclaiming success, and forecasting incredible growth for the next year. I thought I’d do something similar for Coiled; it’s been quite a year for us.

Tone here will be a bit more grounded/realistic than is typical for founders, which readers of my writing may find familiar. Let’s talk through where we are as a company, what we’ve been doing, and what we need to accomplish.

How did we start 2024?

Coiled spent its first couple years (2020-2022) as a confused mess (we call this “the beforetimes” or “dark ages” internally). 2023 was the year when we got our shit together and built a decent product. We did this by shutting down all other parts of the company (sales, marketing, HR, …) and focusing on product development, linking customer feedback and product direction directly to engineers. This worked, and by the end of 2023, we had a delightful product and no ability to sell it.

Here is an excerpt from an e-mail to the company summarizing 2023:

Today we have a good product. It’s a product that I think we’re all proud of. Kudos to us 🎉!

We also have customers (not many yet) and a funnel that actually seems to convert eyeballs into users into dollars (although not well enough yet). The company, today, seems to actually work. Given where we started the year I’m impressed. Thank you all for your efforts in this.

So in 2024 our biggest question changed:

[x] ~~Can we build a great product?~~
[ ] Can we sell this thing? 🤔

In 2024 we shifted our focus to figure out sales and marketing.

Sales

Where we started

Advice I got at this time:

Advice: Great job on product. Time to hire a proven sales/marketing leader!
Me: Hell no.

I still felt burnt from the Coiled Dark Ages and was suspicious of Sales/Marketing VP types who had all the answers. My experience was that experienced players come with more confidence than context in our space, and mostly cause noise and trouble. My stance at the time was that early stage startups don’t work like other companies, and we’d be better off figuring out things for ourselves. While intrepid, that thought process mostly didn’t work 🙂

We started 2024 focusing on a self-serve motion.

User acquisition: Content marketing and strategic placement of examples in highly trafficked OSS pages (turns out people click on examples more than on explicit mentions of SaaS products).
User conversion: People would increase usage naturally until they exceeded the free tier and then put in a credit card

This worked, but not well enough.

What we did and how it worked

In particular we were better at building interest and success than at closing deals, so we decided to hire a couple of sales reps. Here’s a blog post about this at the time: Hiring Sales.

Finding reps that were good at operating in a chaotic environment with high uncertainty was hard. Most sales reps want, you know …

A well-defined thing they can sell
An ideal customer profile
Knowledge about why that customer typically wants to buy it
A defined pricing scheme

We had none of those things. Finding people who are excited about uncertainty is hard but possible. Eventually we found a couple and they started in June. One survived our chaos and today we’re thrilled to have Austin Brown in the family. Austin was previously at FlatFile, and before that an entrepreneur himself with cold brew coffee. (we really value both entrepreneurial experience and experience outside of tech).

Today we sell things. We had to figure a lot of stuff out (for example, see Pivoting away from Usage Based Pricing), and onboarding sales into a product-focused organization was tough, but today the company feels more balanced and general optimism is a lot higher than it has been.

Our typical users / platform customers today fit into a few buckets

A bunch of free users (students, independent contractors, scientists, etc.) who we’ll probably never get money out of (but who we still feel good about serving, as it’s mostly free for us)
Some pay-as-you-go customers paying us ~$1k/month
Small teams paying around $25k-$50k / year
Heavy users paying $100-200k / year
A few larger enterprises paying more

It’s a good mix. We have great customers who give excellent feedback and from whom we’ve learned a great deal. Our customers force us toward success, and we care about making them happy. For the most part I’m proud to think that we’ve achieved that.

What we have to do in 2025

Today I still worry about sales, but more in a day-to-day operations way, rather than in a “will we ever $@%ing make this work” way. The transition from existential dread to quotidian grind is a delightful change.

There’s still plenty to do to mechanize sales. We spend non-trivial effort educating customers on the value we provide, how to use us well, and on how best to work our way through larger enterprise procurement cycles. This is all hard but straightforward work and I’m confident we have the people in place who can do it.

2025 in sales is about execution, tuning, and eventually expansion. I’d love for us to be in a place where we could hire another AE in 6 months and grow beyond that in 12.

Marketing

Our marketing for most of 2024 was kinda meh. This is because when we hired sales we slowed down marketing; there was just too much work to onboard and align the new function in the company. New user flow didn’t stop, it just didn’t grow as fast as would have been ideal. We’ve refocused on this for the last couple months and are seeing glimmers of hope.

Generally though, it turns out that marketing is hard, especially today.

I mostly grew Dask’s awareness through four mechanisms:

A personal blog
Twitter
Relationships with other PyData/SciPy core developers
Speaking at Python conferences

That was ten years ago, and the same tricks don’t seem to work today. The world has become a lot noisier since.

We’re trying lots of things and generally getting smarter about marketing, but this is still a big unknown for us. Figuring this out here is a key objective for us for 2025.

Dask

Where we started

Two years ago Dask was well-loved but had stability issues. It would break from time to time, which eroded confidence among larger users. Those were mostly fixed in 2023.

In 2024, Dask was stable, but still difficult to use at large scale. This isn’t unusual, ask anyone who has driven Spark at the 10+TiB scale and you’ll hear similar stories. Still though, as a company we target lots of small groups without deep distributed computing expertise. We want to stretch the boundaries of allowing increasingly novice users to attempt increasingly complex workloads. In 2024 we wanted to make Dask workloads smooth at the 10-100 TiB scale.

This is especially hard in the Dask/PyData space because performance doesn’t just depend on Dask, but rather on the interplay of dozens of different libraries. We weren’t developing Dask, we had to develop the ecosystem.

What we did

To address this we built, ran, and optimized large-scale benchmark suites that targeted full lifecycle workloads.

Dataframes

We started with TPC-H, a standard benchmark suite for tabular/dataframe/SQL computations, and compared ourselves to DuckDB, Polars, and Spark between scales of 10 GiB to 10 TiB. Over the course of several months and with many exciting new features (query optimization, parquet rewrite, shuffle rewrite, copy-on-write pandas memory semantics) we climbed to a point where Dask Dataframe was on-par across a variety of scales and exceeding the others at large scale. For more information, see DataFrames at Scale Comparison: TPC-H.

This is great, and our dataframe customers are generally much happier on the new system today. We also spend far less time supporting them. Dask dataframe today “just works” far more often than it used to and is technically well positioned to displace Spark.

However, we didn’t do a good job marketing this success (the dataframe space is noisy right now, and we were distracted by integrating the Sales function) and so we missed an opportunity here. This might be something for us to come back to in 2025.

We had a technical win here, but didn’t convert that into a business win.

Geo

After wrapping up the dataframe technical effort the team switched to geospatial workloads (a non-trivial fraction of our current customer base deals with geospatial data). Unfortunately, there was no standard benchmark in this space, so we started by building one. Here are some blogs on the subject:

This technical work is ongoing and includes work across Dask, Xarray, Zarr, GeoPandas, S3FS, and other libraries. We’ve climbed from “works reliably on 10 GiB” to “works reliably on 100 GiB with improvements on 1 TiB” but still have more work to do as we climb to the 10-100 TiB scale.

Geo is a fun place for us at our current stage; it’s a good place for us to figure ourselves out. Geo is a tighter community where word travels fast, and without as much noise. It’s been a fun vertical for the team to focus on. It’s still early, but uptake and commercials look good.

What we need to do in 2025

The Dask team at Coiled is delivering broad technical wins fairly reliably. This is good for everyone. Repeating for Geo what we did for Dataframes is a good focus for the first couple quarters of the year.

However, we need to do a better job of connecting this technical success into commercial success. For Geo we’re being more mindful of this than we were for dataframes.

I hope that once the geo effort is feeling complete we can return to the topic of dataframes and stride into that conflict more purposefully, and with more maturity.

Cloud Platform

The product is good. Usage is up. Service is stable. People have a great time.

We’re thinking about where to go from here.

Usage is up

Here’s a chart showing usage, grouped by organization, over the last twelve months (each block is a different org).

Compute hours managed by Coiled over time, grouped by organization. Growth is steady. There are a few large organizations and many many small ones.

More importantly though, we now charge for more than just usage (see Pivoting away from Usage Based Pricing) and so revenue is no longer tied to this chart (it’s quite a bit better).

What’s next?

Coiled, originally conceived as a Dask SaaS product, accidentally became a great system to run general purpose computing tasks. It turns out that to launch Dask clusters, you first need to figure out how to launch general Python jobs in the cloud, and we did that well, arguably better than many cloud services.

This presents some open opportunities we’re considering:

More Dask UX

More charts, more automated warnings, automated hardware selection, etc..
Jupyter Notebooks

Almost every Jupyter Notebook startup that has gone down this route has failed (excepting Hex?), but who knows, we might fare better. The number of users and institutions interested in this feature (despite the flood of startups offering it) is unnervingly high.
Batch jobs

We recently released a feature letting people run arbitrary programs (not just Dask or even Python) on Coiled, competing with tools like AWS/Azure Batch. This took a couple days to put together and has had surprisingly positive response in a short time.
Prefect serverless

Prefect’s serverless push work pool ends up being a surprisingly nice match for Coiled and would be of use to a non-trivial fraction of our current customers.

Probably we could do something with Dagster/Airflow as well.

I’m torn between focusing on fine-tuning our already strong Dask experience, or broadening out to “general-purpose easy-to-use large-scale computing”. Usually I like solving focused problems (like Dask), and view the general purpose problem to be more the purview of the clouds themselves. However, after engaging with their offerings (AWS Batch, Lambda, Sagemaker, etc.) my sense is that they’re not serving the Python data community well, mostly because of a developer experience that doesn’t fit well with the taste of most data developers. (See post Cloud Computing is Broken)

The way we answer this question is likely by observing users. Today our users use Coiled for all sorts of things we didn’t expect, including:

Running geophysical inversion codes to detect where to mine for minerals
Starting up Jupyter notebooks for interactive GPU computation
Augmenting workflow managers like Prefect and Dagster with large machines
Bulk conversion of millions of geotiff images
Benchmarking GPU codes on all possible GPU types available on the cloud

It ends up being easier to misuse Coiled beyond its intended purpose than to use the traditional tools created by the commercial clouds. This atypical usage isn’t currently enough of our revenue to pivot the company today, but it’s an enticing possibility for the future, and one that we’ll likely explore with baby steps in 2025.

Funding

Ever since the Dark Ages we’ve run a lean shop and so our runway is both generous and growing now that we’ve unlocked sales. I have no interest today in taking on additional outside capital.

Internally though, we should probably start to hire. I think that we have a better idea of how to effectively put our existing capital to use and it’s time to (very gently) push on the gas.

Personal

In my e-mail to the company at the end of 2023 I wrote the following:

Personally I find this job to be challenging (I don’t think of myself as being naturally suited to this role) but there have been times in the last few months when I’ve actually found myself enjoying it. Thank you all for that.

The trend away from “Leading a startup feels like chewing glass” and towards “I could maybe enjoy this job some day” has continued. I find myself significantly more calm today than I was a year ago. Much of this has to do with everyone in the company growing themselves and learning how to work together more effectively.

Working with smart people who share a common mission makes all the difference.

Summary

We’re doing ok. In 2024 we unblocked sales and accumulated enough additional exposure to user pain that we have some ideas that might be good.

In 2025 if all we do is execute and tune and mechanize then that’ll be fine. However, I’m excited about increasing creativity, especially in marketing and a broadening of our product offering.

Faster Xarray Quantile Computations with Dask

2024-12-17T00:00:00+00:00

Dec 17, 2024

2m read

Patrick Hoefler

There have been a number of engineering improvements to Dask Array like consistent chunksizes in Xarray rolling-constructs and improved efficiency in map_overlap. Notably, as of Dask version 2024.11.2, calculating quantiles is much faster and more reliable.

Dask dashboard of the new quantile implementation, which is ~20x faster for this microbenchmark.

Calculating Quantiles with Xarray

Calculating quantiles is a common operation for geospatial data. Quantiles show how a dataset is distributed over time, allowing you to identify trends, anomalies, and variation within specific grid cells or regions. These calculations are typically performed either for distinct groups within the dataset or across the dataset as a whole.

The Problem: `np.quantile` Can Be Slow

The data we often encounter typically has a relatively short time axis, consisting of only a few hundred to a few thousand values, while the latitude and longitude dimensions are extensive.

Previously, the quantile calculation for every coordinate was done on a pretty small array with the NumPy implementation np.quantile or np.nanquantile:

import numpy as np

np.quantile(np.random.randn(500), q=0.5)

Dask lacked an efficient efficient multidimensional implementation for calculating quantiles, which meant calling the one-dimensional NumPy implementation millions of times in Python. This is very slow and also blocks the GIL (looking forward to a free-threaded Python world 😅). This caused large slowdowns on workers with more than one thread and could lead to runtimes over 200s per chunk.

Running the following computation previously took over 3 minutes to complete:

import xarray as xr
import dask.array as da

arr = da.random.random((50, 3_000, 3_000), chunks=(-1, "auto", "auto"))

darr = xr.DataArray(
    arr, dims=["time", "x", "y"]
)

darr.quantile(dim="time", q=0.75).compute()

Computing quantiles used to be painfully slow. The computation takes ~3 minutes to complete.

The Solution: A New `dask.array.quantile`

As of dask=2024.11.2, we’ve added a high-level quantile API to Dask that uses top level NumPy functions to extract the quantile for each time slice. The operations are more expensive than an optimal quantile implementation, but the vectorized, multidimensional nature of each call makes it a lot faster in aggregate. And we don’t block the GIL anymore, so you can run with proper parallelism on your Dask workers.

Using the new implementation, this same microbenchmark takes ~9 seconds, a 20x speedup.

The new quantile implementation reduces runtime to ~1s per chunk, independent of the number of threads. This means we’re able to calculate quantiles hundreds of times faster than before! You can expect the speedups to scale with the size of your quantile axis.

You can try out the new implementation by upgrading to the latest versions of Dask and Xarray.

conda

conda update dask xarray

pip

pip install --upgrade "dask[complete]" xarray

For other recent improvements to using Xarray with Dask, you might consider:

Improving GroupBy.map with Dask and Xarray

2024-11-21T00:00:00+00:00

Nov 21, 2024

4m read

Patrick Hoefler

This post was originally published on the Xarray blog.

Running large-scale GroupBy-Map patterns with Xarray that are backed by Dask arrays is an essential part of a lot of typical geospatial workloads. Detrending is a very common operation where this pattern is needed.

In this post, we will explore how and why this caused so many pitfalls for Xarray users in the past and how we improved performance and scalability with a few changes to how Dask subselects data.

What is GroupBy.map?

GroupBy.map lets you apply a User Defined Function (UDF) that accepts and returns Xarray objects. The UDF will receive an Xarray object (either a Dataset or a DataArray) containing Dask arrays corresponding to one single group. Groupby.reduce is quite similar in that it applies a UDF, but in this case the UDF will receive the underlying Dask arrays, not Xarray objects.

The Application

Consider a typical workflow where you want to apply a detrending step. You may want to smooth out the data by removing the trends over time. This is a common operation in climate science and normally looks roughly like this:

def detrending_step(arr: DataArray) -> DataArray:
    # important: the rolling operation is applied within a group
    return arr - arr.rolling(time=30, min_periods=1).mean()

data.groupby("time.dayofyear").map(detrending_step)

We are grouping by the day of the year and then are calculating the rolling average over 30-year windows for a particular day.

Our example will run on a 1 TiB array, 64 years worth of data and the following structure:

The array isn’t overly huge and the chunks are reasonably sized.

The Problem

The general application seems straightforward. Group by the day of the year and apply a UDF to every group. There are a few pitfalls in this application that can make the result of this operation unusable. Our array is sorted by time, which means that we have to pick entries from many different areas in the array to create a single group (corresponding to a single day of the year). Picking the same day of every year is basically a slicing operation with a step size of 365.

Our example has a year worth of data in a single chunk along the time axis. The general problem exists for any workload where you have to access random entries of data. This particular access pattern means that we have to pick one value per chunk, which is pretty inefficient. The right side shows the individual groups that we are operating on.

One of the main issues with this pattern is that Dask will create a single output chunk per time entry, e.g. each group will consist of as many chunks as we have year.

This results in a huge increase in the number of chunks:

This simple operation increases the number of chunks from 5000 to close to 2 million. Each chunk only has a few hundred kilobytes of data. This is pretty bad!

Dask computations generally scale along the number of chunks you have. Increasing the chunks by such a large factor is catastrophic. Each follow-up operation, as simple as a-b will create 2 million additional tasks.

The only workaround for users was to rechunk to something more sensible afterward, but it still keeps the incredibly expensive indexing operation in the graph.

Note this is the underlying problem that is solved by flox for aggregations like .mean() using parallel-native algorithms to avoid the expense of indexing out each group.

Improvements to the Data Selection algorithm

The method of how Dask selected the data was objectively pretty bad. A rewrite of the underlying algorithm enabled us to achieve a much more robust result. The new algorithm is a lot smarter about how to pick values from each individual chunk, but most importantly, it will try to preserve the input chunksize as closely as possible.

For our initial example, it will put every group into a single chunk. This means that we will end up with the number of chunks along the time axis being equal to the number of groups, i.e. 365.

The algorithm reduces the number of chunks from 2 million to roughly 30 thousand, which is a huge improvement and a scale that Dask can easily handle. The graph is now much smaller, and the follow-up operations will run a lot faster as well.

This improvement will help every operation that we listed above and make the scale a lot more reliably than before. The algorithm is used very widely across Dask and Xarray and thus, influences many methods.

What’s next?

Xarray selects one group at a time for groupby(...).map(...), i.e. this requires one operation per group. This will hurt scalability if the dataset has a very large number of groups, because the computation will create a very expensive graph. There is currently an effort to implement alternative APIs that are shuffle-based to circumvent that problem. A current PR is available here.

The fragmentation of the output chunks by indexing is something that will hurt every workflow that is selecting data in a random pattern. This also includes:

.sel if you aren’t using slices explicitly
.isel
.sortby
groupby(...).quantile()
and many more.

We expect all of these workloads to be substantially improved now.

Additionally, Dask improved a lot of things related to either increasing chunksizes or fragmentation of chunks over the cycle of a workload with more improvements to come. This will help a lot of users to get better and more reliable performance.

SLURM-Style Job Arrays on the Cloud with Coiled

2024-11-19T00:00:00+00:00

Nov 19, 2024

3m read

Matthew Rocklin

Update Jan. 9th, 2025: We’ve added documentation for Coiled Batch.

SLURM and other job schedulers have made it easy to run scripts on HPC systems for decades. Job scripts are straightforward and solve a common problem in a way that is familiar to researchers who have had exposure to HPC systems. Heck, I personally used them as a graduate student long before I knew anything about parallel computing.

Unfortunately, the cloud doesn’t offer anything similar with the same level of ease. (It’s oddly difficult to do this on the cloud). We wanted to fix that, and so just released this new feature with Coiled reinventing the SLURM-style job script for the cloud.

Disclaimer: we just pushed this out last week and it’s not polished yet, doesn’t have docs, etc.. It felt easy enough though that we felt safe talking about it early in hope of getting more feedback.

Example: Hello world!

We’ll give lots more examples below but for now let’s do a trivial Hello world example parallel job script

$ coiled batch run hello.sh
Running job array of 10 tasks on 10 VMs...
Cluster Details: https://cloud.coiled.io/clusters/661309?account=dask-engineering

To check status, run:
  coiled batch status 661309

$ coiled batch status 661309
                    Batch Jobs for Cluster 661309 (stopped)
                               Command: hello.sh
┏━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━━┓
┃          ┃ Assigned   ┃       ┃ Start     ┃           ┃          ┃           ┃
┃ Array ID ┃ To         ┃ State ┃ Time      ┃ Stop Time ┃ Duration ┃ Exit Code ┃
┡━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━━┩
│ 0        │ 10.0.46.1… │ done  │ 15:59:46  │ 15:59:46  │ 0s       │ 0         │
│ 1        │ 10.0.46.1… │ done  │ 15:59:46  │ 15:59:46  │ 0s       │ 0         │
│ 2        │ 10.0.46.1… │ done  │ 15:59:46  │ 15:59:46  │ 0s       │ 0         │
│ 3        │ 10.0.46.1… │ done  │ 15:59:46  │ 15:59:47  │ 0s       │ 0         │
│ 4        │ 10.0.46.1… │ done  │ 15:59:47  │ 15:59:47  │ 0s       │ 0         │
│ 5        │ 10.0.46.1… │ done  │ 15:59:47  │ 15:59:47  │ 0s       │ 0         │
│ 6        │ 10.0.46.1… │ done  │ 15:59:47  │ 15:59:47  │ 0s       │ 0         │
│ 7        │ 10.0.46.1… │ done  │ 15:59:48  │ 15:59:48  │ 0s       │ 0         │
│ 8        │ 10.0.37.50 │ done  │ 15:59:48  │ 15:59:48  │ 0s       │ 0         │
│ 9        │ 10.0.46.1… │ done  │ 15:59:48  │ 15:59:48  │ 0s       │ 0         │
└──────────┴────────────┴───────┴───────────┴───────────┴──────────┴───────────┘

And we can easily go and get logs from all the machines at once

$ coiled logs 661309 | grep Hello
2024-11-19 15:59:46.131000 Hello from 0 !
2024-11-19 15:59:46.419000 Hello from 1 !
2024-11-19 15:59:46.743000 Hello from 2 !
2024-11-19 15:59:47.019000 Hello from 3 !
2024-11-19 15:59:47.324000 Hello from 4 !
2024-11-19 15:59:47.587000 Hello from 5 !
2024-11-19 15:59:47.862000 Hello from 6 !
2024-11-19 15:59:48.256000 Hello from 7 !

What Does This Do?

We …

Inspect your script
Spin up appropriate machines as defined by the comments (this takes about 30s)
Download your software onto them (maybe another 30-60s)

(in this case a Docker image, but Coiled’s environment synchronization works here too)
Run your script
Shut down the machines.

Your script can do anything. It doesn’t have to be Dask related, or even Python related (indeed the echo program used above predates the Python language).

This is very simple, but also very powerful. Let’s look at a few more examples:

Cloud Job Script Examples

Hello (bash)

Our example from above, for completeness’ sake.

#!/bin/bash
#COILED n-tasks     10
#COILED memory      8 GiB
#COILED container   ubuntu:latest
#COILED region      us-east-2

echo Hello from $COILED_BATCH_TASK_ID

Hello (python)

We can submit python scripts too

#COILED n-tasks     10
#COILED memory      8 GiB
#COILED region      us-east-2

import os

print("Hello from", os.environ["COILED_BATCH_TASK_ID"])

Here we also drop the container directive, and rely on Coiled’s environment synchronization to copy all of our Python libraries to the remote machines automatically.

Process Many Files

A common pattern is to list many files, get the ith file, and run some command on that file.

#COILED n-tasks     1000
#COILED memory      8 GiB
#COILED region      us-east-2

import os
import s3fs

s3 = s3fs.S3FileSystem()

# Get the i'th file to process
task_id = int(os.environ["COILED_BATCH_TASK_ID"])
filenames = s3.ls("my-bucket")
filename = filenames[task_id]

# Process that file
def process(filename):
    result = ...
    return result

result = process(filename)

# Store result
s3.put(...)  # put result back in S3 somewhere

GPUs

We can ask for any hardware on any cloud, including a collection of GPUs machines

#COILED n-tasks     100
#COILED vm-type     g5.xlarge
#COILED region      us-east-2

import torch

model = ...

for epoch in range(50):
    model.train()
    ...

When you use a GPU machine with Coiled’s environment sync feature, it will automatically install GPU equivalents to local packages like pytorch, along with all the necessary NVIDIA CUDA drivers. So you can develop locally on your CPU or Apple hardware:

$ python train.py               # Uses your local CPU

And when you’re ready to run at full scale you can launch your script to run on many GPU machines:

$ coiled batch run train.py     # Uses many cloud GPUs

Video

Status and Logs

Just like on HPC, it’s easy to get the status of all jobs in your Coiled workspace. This also includes the approximate cloud cost for each job (we’ve seen people be curious about this in the past).

$ coiled batch status
┏━━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━┓
┃ Cluster ID ┃  State  ┃ Tasks Done ┃           Submitted ┃            Finished ┃ Approx Cloud Cost ┃ User Command        ┃
┡━━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━┩
│ 661309     │  done   │   10 /  10 │            15:58:50 │            15:59:48 │             $0.01 │ hello.sh            │
│ 659567     │  done   │   10 /  10 │ 2024-11-18 17:25:08 │ 2024-11-18 17:26:07 │             $0.01 │ hello.sh            │
│ 658315     │  done   │   10 /  10 │ 2024-11-17 16:24:27 │ 2024-11-17 16:25:20 │             $0.01 │ hello.sh            │
│ 658292     │  done   │   20 /  20 │ 2024-11-17 15:28:44 │ 2024-11-17 15:29:39 │             $0.02 │ hello.sh            │
│ 658288     │  done   │    0 /  20 │ 2024-11-17 15:24:32 │                     │             $0.02 │ hello.sh            │
│ 657615     │  done   │    0 /  10 │ 2024-11-16 20:32:46 │                     │             $0.01 │ hello.sh            │
│ 656571     │  done   │   10 /  10 │ 2024-11-15 14:32:35 │ 2024-11-15 14:33:35 │             $0.01 │ hello.sh            │
│ 655848     │  done   │   20 /  20 │ 2024-11-14 23:56:06 │ 2024-11-14 23:58:01 │             $0.01 │ echo 'Hello, world' │
└────────────┴─────────┴────────────┴─────────────────────┴─────────────────────┴───────────────────┴─────────────────────┘

Shoutout to Python’s rich library for the nice formatting 🙂

And while you can always look at jobs on the Cloud UI, we also provide a convenient way to get logs given a cluster ID (this is true of all Coiled jobs)

$ coiled logs 661309 | grep Hello
2024-11-19 15:59:46.131000 Hello from 0 !
2024-11-19 15:59:46.419000 Hello from 1 !
2024-11-19 15:59:46.743000 Hello from 2 !
2024-11-19 15:59:47.019000 Hello from 3 !
2024-11-19 15:59:47.324000 Hello from 4 !
2024-11-19 15:59:47.587000 Hello from 5 !
2024-11-19 15:59:47.862000 Hello from 6 !
2024-11-19 15:59:48.256000 Hello from 7 !
2024-11-19 15:59:48.315000 Hello from 8 !
2024-11-19 15:59:48.508000 Hello from 9 !

Making Common Things Easy

This kind of embarrassingly parallel workload is super common.

We saw Python users struggling with the services provided by cloud providers like AWS Batch, GCP Cloud Run, or Azure Batch and thought that this little utility on top of Coiled might make it easier for these users get use out of the cloud.

We hope you like it! As always, if you want to try out Coiled, you can get started at coiled.io/start.

Airflow, Dask, & Coiled: Adding Big Data Processing to Your Cloud Toolkit

2024-11-12T00:00:00+00:00

Nov 12, 2024

6m read

Stephen Schneider, Franco Bosetti

Summary

In this article, we showcase how Coiled technology enabled our data engineering and analytics team at Siemens to tackle complex data processing challenges. For an internal customer, we needed to analyze a large dataset of training records for an employee learning platform. By seamlessly integrating with our existing Python code and allowing us to deploy scalable Dask clusters in the cloud, Coiled reduced our ETL execution time by over 80%. This optimization not only benefited our internal Siemens customers, but also expanded our own data engineering portfolio for future projects.

We run Airflow locally and use Coiled to easily start a remote Dask cluster on AWS.

Background

As a data engineering and analytics team that works mainly in AWS, we frequently look to expand our technical toolkit and the portfolio of services we offer our customers. Besides using AWS Lambda functions for small event driven tasks, such as file deliveries and ingestions, we recently migrated to Apache Airflow as our main ETL tool.

We typically use Airflow PythonOperators to connect to various systems and perform tasks in a designed ETL workflow called a DAG. These tasks could be fetching data, pulling that data into a pandas DataFrame, and then pushing that DataFrame to a target destination. For efficiency, we execute most of our heavy ETL logic (data cleansing, joins, aggregations, calculations, etc.) in SQL and rely on our databases to handle the manipulation of large datasets.

Recently, however, we encountered a project where we needed to use a fuzzy algorithm to identify patterns and correlations in a dataset of employee training records. This algorithm is too complex to perform using BI tools or SQL. We needed to execute our logic directly in Python using a specific library and the resources available in our Airflow environment. For simple logic and small to medium sized datasets, this would usually not be a problem. In this case, however, applying the algorithm to a large dataset was very time consuming. Even after provisioning additional resources to Airflow, our job took at least an hour.

This was not optimal for our customer. We needed a different solution for processing big data more efficiently.

AWS Glue vs. EMR vs. Coiled

After consulting our wise mentor ChatGPT, we determined we had a few options. In AWS some popular services for processing big data include AWS Glue and Elastic MapReduce (EMR) with Apache Spark. After some research, we realized most of these solutions involved a steep learning curve: managing new infrastructure, and/or essentially completely rewriting our existing Python code. After ruling out these options, another solution caught our eye: Coiled. ChatGPT summarized Coiled as “a managed Dask service that allows you to deploy Dask clusters in the cloud without the need to manage the underlying infrastructure.”

Not worrying about infrastructure sounded enticing. But would we have to rewrite our Python code to work with the Dask library? We pressed ChatGPT again. It assured us that since Dask DataFrame mimics pandas, we could easily scale to large datasets with minimal code changes.

This sounded too good to be true. We could use Coiled without having to worry about managing infrastructure in AWS and we wouldn’t need to completely rewrite our existing code. If this characterization of Coiled was accurate, we worried this could be expensive. Luckily, Coiled provides a free tier that was enough to cover our needs.

Getting started with Coiled

We decided to give Coiled a shot.

We began by connecting our AWS account to Coiled. The setup in our AWS environment was simple; Coiled offers a link to CloudFormation to create a preformatted stack of the resources. This stack now allowed Coiled to create Dask clusters on our behalf. We then generated a new API token for authorizing our Airflow environment.

Creating the Coiled Software Environment

When you start a Coiled cluster and do not specify any packages, the Coiled environment mirrors the library versions that exist in your local environment by default. It’s easier and faster than using Docker, and meant that we didn’t have to worry about Python package versions.

Authentication

To authenticate Coiled from Airflow, we set the Coiled token with dask.config.set({"coiled.token": <our-token>}). We also set the AWS_ACCESS_KEY_ID, and AWS_SECRET_ACCESS_KEY environment variables to allow Coiled to interact with AWS resources on our behalf.

Coiled Cluster

Starting the cluster is as easy as calling coiled.Cluster(), however, there are a number of customization options that allow us to tailor the cluster to our specific workload, optimizing both performance and cost. There are no standard rules here, so we tested different configurations to see how quickly our job ran using a varying number of workers and instance types in the cluster.

cluster = coiled.Cluster(
    n_workers=10,
    worker_vm_types=["c5.large"]    # VM for compute-intensive workloads
)
client = cluster.get_client()       # Connect to cluster

...                                 # Dask computations

cluster.shutdown()                  # Cleanup resources

Once the cluster is set up, we connect to it. This connection is necessary to submit tasks and manage computations across the cluster. The client acts as a gateway between our Airflow environment and the distributed Dask workers.

Once all tasks are completed, we shutdown the cluster (Coiled will also do this automatically after 20 minutes of inactivity).

Putting it all together

Once the cluster is up and running, we can execute our data processing logic on a Dask DataFrame, a parallelized version of pandas. After connecting to the cluster, we load our data into a Dask DataFrame and apply some group-by logic and the fuzzy algorithm to achieve our final dataset. The beauty of using Dask within this setup is that it automatically optimizes and distributes these operations across the cluster, making the processing highly efficient and easy to implement.

We used the Coiled UI to track a number of performance metrics and also keep track of uptimes, cloud costs, tasks processed, etc. During testing and set-up, we were regularly contacted by the Coiled team. They generously offered technical assistance and answered any questions we had. In one instance, the Dask library was updated in our Airflow environment, which caused a mismatch between Airflow and the Coiled software environment. A Coiled engineer promptly reached out after noticing our job had failed and pointed out the library discrepancies, which ultimately led to us using package sync to automatically handle the Python environment instead.

Conclusion

The results of utilizing Coiled to run our job on a Dask cluster were remarkable. The execution time of our code shrank from over 1 hour to around 10 minutes, a reduction of over 80%. Our customer was thrilled, and our group was able to confidently add big data processing to our portfolio of data engineering services. If you are like us, passionate programmers who have little interest in slogging through cloud infrastructure, consider our use case. Forget outdated solutions like Apache Spark or low-code services like AWS Glue. Call in the big guns with a few lines of code. Use Coiled.

Large Scale Geospatial Benchmarks: First Pass

2024-10-16T00:00:00+00:00

Summary

We implement several large-scale geo benchmarks. Most break. Fun!

This article describes those benchmarks, what they attempt, how they break, and the technical work necessary to make them run smoothly.

Context

Last month we published Large Scale Geoscience: Solicitation a community call for geo benchmarks that reflected common workloads in the 100 TiB range. We wanted to source benchmarks that were simple enough to understand, complex enough to be real, and generally very big.

Many people contributed to this lively GitHub discussion. Since then we’ve implemented several of those benchmarks and looked at how they run at scale. This post walks through the benchmarks implemented so far, how they perform at scale, and outlines directions for future improvement.

Now let’s go through what has been implemented so far, benchmark by benchmark.

Rechunking

Many applications’ performance depends strongly on data layout. Data laid out in one way (for example partitioned by time) may be inefficient for certain applications (like querying all data at a specific location) and vice versa. To resolve this we often rechunk datasets to have a more friendly structure for the computation at hand.

Rechunking from time-slice partitioned to spatially partitioned

This rechunking benchmark rechunks the ERA5 global climate reanalysis dataset from time-oriented chunks to spatial chunks. This rechunking pattern frequently occurs with Earth Observation (EO) data. Additionally, in this benchmark Xarray/Dask automatically chose to increase chunk sizes from 4 MiB to 100 MiB, which works better with cloud storage.

However, as this workflow scales up we encounter the following issues:

Large task graphs: As the data volume increases, the graph size increases, resulting in two issues:
- Initial Delay: Sending the large graphs from the client to the cluster scheduler can take minutes. This is bad both because workers idle during this time, and also because users are confused why there isn’t any progress. This delay is psychologically uncomfortable.
- Scheduler Memory: Each task takes up about 1kB in RAM. Normally this is fine, but as graph sizes climb to hundreds of millions this can require schedulers with more memory. This is fine in principle (it’s easy to get bigger schedulers) but it’s a common stumbling point for beginners.
Right-sizing disk size
- Historically, rechunking used to require a large amount of memory. This was a common pain point for novice users. Fortunately, this has been fixed and Dask array’s newer rechunking algorithm (made default just last month) now uses disk efficiently.
- Today we’re running out of disk because these benchmarks are so large. This was easy to fix (we can ask either for more workers, or for more storage per worker) but this is a stumbling point for beginners, or those for whom configuring distributed clusters is hard.

Cluster running out of disk space and crashing.

These challenges of large graph sizes and right-sizing clusters are common concerns to anyone familiar with these tools at this scale. We’re not surprised to see them in a large rechunking workload like this. There are workarounds here that experts use. We’re intentionally not invoking those workarounds. We’re looking forward to working to remove these concerns in a way that will impact less-sophisticated users.

Climatology

Calculating climatologies or seasonal averages is a standard operation when working with weather/climate data. Climatologies are both useful in their own right and also in related computations like anomalies (deviations from the average).

One slice of ERA5 to form a climatology

This climatology benchmark computes seasonal averages of the ERA5 global climate reanalysis dataset. There are a couple of ways to do this.

Ideally we would just write down normal Xarray code, but that fails hard. Instead it has become common-place for users to rechunk the data to align with the climatology one wants to compute, and then do full-chunk calculations. This approach scales better, but also fails. There’s some work to do here as we scale up.

Weighted rolling reductions scale poorly in Xarray (see this issue for more details).

Chunk sizes unexpectedly grew from MBs to GBs, which can cause out-of-memory errors.
Rechunking feels unnatural

This type of computation today is often reframed using Xarray’s map_blocks as a workaround for the poor rolling performance. However, this approach is tedious and can feel unnatural, especially for less sophisticated Xarray users.
Anomalies keep data in memory

This workload is often combined with calculating anomalies. When combining these two we find that Dask doesn’t know if it should drop intermediate data that is used twice (first for calculating the average, and then again for calculating the anomaly) or keep it. This results in a lot of disk engagement, which can be slow, especially on older versions of Dask.

NetCDF to Zarr

NetCDF data is ubiquitous, but often slow and expensive to access in the cloud. Because of this, people commonly convert netCDF data to Zarr, a cloud-optimized format. These workflows often also include a modest rechunking step to improve chunking for cloud access (cloud systems tend to prefer larger reads than POSIX filesystems).

This cloud-optimization benchmark processes the NASA Earth Exchange Global Daily Downscaled Projections (NEX-GDDP-CMIP6) dataset (which is stored as thousands of NetCDF files) by loading the dataset with Xarray, rechunking from time-oriented chunks to spatial chunks, and then writing the result to a Zarr dataset.

See https://cmip6.surge.sh/discoveries/crossingyear/global

As this workflow scales up, we encounter the following issues:

Xarray’s open_mfdataset doesn’t scale well when reading many files.
- Xarray today mostly calls open_dataset many times, and then concatenates those results. You’d think that would work well, but it doesn’t. The resulting task graph ends up being a bit unnatural. We can address this by writing a bit of custom code for open_mfdataset
- It can take non-trivial time to read the metadata of all the NetCDF files. Advanced users know to work around this with the parallel=True keyword argument, which parallelizes the metadata reading on the workers. Unfortunatley, novice users stumble on this.
As above, this workload inherits the issues around rechunking and right-sizing clusters.

Improving how xarray.open_mfdataset performs on several thousand files will have a broad impact beyond this benchmark on any large scale workload processing data formats like NetCDF.

Atmospheric circulation

This atmospheric circulation benchmark, proposed by Deepak Cherian, computes an Eulerian mean, depending on windspeed and temperature variables.

zonal_means = ds.mean("longitude")
anomaly = ds - zonal_means

anomaly["uv"] = anomaly.U * anomaly.V
anomaly["vt"] = anomaly.V * anomaly.T
anomaly["uw"] = anomaly.U * anomaly.W

temdiags = zonal_means.merge(anomaly[["uv", "vt", "uw"]].mean("longitude"))

This benchmark runs smoothly in small memory (mostly), but graph generation time is painfully slow, waiting almost ten minutes before anything happens.

Worker CPU and Memory usage. There’s a long delay before anything occurs.

During this period the scheduler is very busy.

All of this points to an inefficiently constructed task graph.

Regridding

Regridding, or remapping, is another common operation when working with geospatial data. It is the process of interpolating data when changing the underlying grid. This may involve a change in the grid resolution or transitioning to a different type of grid. To complicate matters, there is no single standard regridding library people seem to rely on. To learn more about the state-of-the-art, see Max Jones’ Pangeo Showcase. For our benchmark, we have chosen xESMF, which appears to be somewhat common.

Regridded dataset between fine and coarse resolutions

While this benchmark is embarrassingly parallel along the time dimension, it has problems when scaling the input or output grid due to their effect on the weight matrix used to interpolate the data:

Sequential Weight Matrix Calculation: The cost of computing the weight matrix is high, and this step precedes the submission of the computation to the cluster. When scaling to higher grid resolutions, this causes the benchmark to stall for minutes before ever submitting work to the cluster, leaving resources underutilized.
Large Weight Matrices: The size of the weight matrix grows significantly with larger grid sizes. While the matrix for regridding from a 1440×721 resolution to 240×121 takes up around 2 MiB, the matrix for regridding from 1440×721 to 3600×1801 takes up 400 MiB. This matrix is embedded in the task graph, both increasing the size of the graph and the time necessary for submitting it to the cluster.

Probably we should consider doing this work on the cluster instead.

Satellite Image Processing

Until now most of these benchmarks have been weather/climate/Xarray focused. However, we often encounter completely different use cases and problems in groups mostly focused on satellite imagery. Think of switching out Zarr for COG here.

People working on satellite imagery often don’t need big parallel Xarray jobs, instead they need an easy and efficient way to run a Python function many times on many individual files, each of which fits comfortably in memory. This is common with GeoTIFF data, or with NetCDF files storing L1 or L2 semi-processed satellite data. This user group cares about cloud computing and distributed computing, but the standard challenges around Xarray and Dask don’t apply as strongly.

This benchmark looks through Sentinel-2 data stored on the Microsoft Planetary Computer with the ODC-STAC package, computes a humidity index, and then computes a monthly average.

Additionally we have this historical workload, which lists tens of thousands of small netCDF files, and then does a simple data access computation on each one independently. For this we use coiled.function rather than Xarray (although any serverless function solution should work fine).

@coiled.function()
def process(filename):
    ...

filenames = s3.ls("...")
result = process.map(filenames)

As this workflow scales up, we encounter a different class of issues:

Limited Availability: We can do lots and lots of work here in parallel, but there are various limits stopping us from getting more workers. We encountered in order:
- Account quota: the account we were running on had a limited quota
- Spot availability: as mentioned below, this often becomes a cost optimization game, and even when we had quota we started running out of spot instances
- Instance type availability: if we’re willing to go beyond spot and pay full price, there are limits to any particular instance type
Aggressive spinup: it’s not immediately clear how many machines we can use efficiently in parallel.

By default we’re using Dask’s adaptive scaling heuristics, which run a few tasks and then based on that runtime information scale up accordingly. This is fine, but does mean that it takes a few minutes before the full army of workers arrives, which can be psychologically a little frurstrating. Maybe we can accelerate this a bit and be more aggressive.

This benchmark exposes more about the underlying cloud architecture we use, rather than the open source Python tooling. This makes it a little different than the benchmarks above, but still something that our team is interested in (and we encourage teams of similar deployment projects to take an interest as well). This benchmark is relatively small compared to the others (barely a terabyte) and we’d especially welcome more ideas and contributions in this direction.

GeoPhysical Inversions

Similarly distinct from the Xarray workloads we’ve run into users running thousands of geophysical inversions using SimPEG. This problem is common in mining and resource extraction, or wherever we have sparse measurements within a physically constrained system (is there oil under this rock? Is there water under this soil? Is there cancer under this muscle?). Increasingly as statistical methods become common we find that people want to do not one inversion problem, but thousands, which becomes computationally expensive.

We think that it might be possible to accelerate these workloads with GPUs, so we’re trying that in close collaboration with the SimPEG development team, and teams at NVIDIA. Sorry, no standard benchmark on this one just yet.

Performance Optimization

Once we’ve finished up an initial set of benchmarks (there are still more to do) we’ll start optimizing. This involves running benchmarks, identifying major bottlenecks, and then fixing those bottlenecks wherever we happen to find them in the various open-source projects (Xarray, Dask, s3fs, Zarr, …)

We’ve already identified several areas for improvement above. Let’s list some of them below:

Moving Large Graphs, which can be resolved both by making our graph representation more compact, or by moving symbolic representations of the graphs instead and expanding them only once they’re on the scheduler.
Reimplementations of core methods like xarray.open_mfdataset or rolling methods to be a bit more Dask-native where appropriate.
Better hinting of hardware failures like when rechunking runs out of disk so that users get more sensible error messages and can ask for different hardware
More aggressive autoscaling so that people feel less compelled to right-size clusters ahead of time and can trust adaptive scaling more (adaptive scaling seems like a nice idea in theory, but most people manually specify sizes still today).
Well-crafted sane defaults for keywords: Today Xarray experts often specify a mix of keywords when working at larger scales. We hope to be able to replace some of these with heuristics to help more novice users.
Less reliance on manual rechunking so that users can focus on expressing what they want to compute without worrying about the distributed memory layout or the footprint of individual steps. Instead, Xarray should become better at automatically dealing with chunking schemes that are not optimal.

Feedback

Feedback from the community is critical to ensuring these benchmarks reflect real-world use cases that users care about.

Is this the right set of benchmarks to use? Are key workloads missing?

If you have feedback on these benchmarks or would like to propose a new benchmark, please engage in this GitHub discussion.

Acknowledgments

Finally, thank you to everyone who has engaged in this effort so far. In particular, thank you to Deepak Cherian, Stephan Hoyer, Max Jones, Guillaume Eynard-Bontemps, and Kirill Kouzoubov for contributing benchmarks and discussions.

Scaling AI-Based Data Processing with Hugging Face + Dask

2024-10-09T00:00:00+00:00

Oct 9, 2024

4m read

Sarah Johnson, James Bourbeau, Quentin Lhoest, Daniel van Strien

This post was originally published on the Hugging Face blog.

The Hugging Face platform has many datasets and pre-trained models that make using and training state-of-the-art machine learning models increasingly accessible. However, it can be hard to scale AI tasks because AI datasets are often large (100s GBs to TBs) and using Hugging Face transformers for model inference can sometimes be computationally expensive.

Dask, a Python library for distributed computing, can handle out-of-core computing (processing data that doesn’t fit in memory) by breaking datasets into manageable chunks. This makes it easy to do things like:

Efficient data loading and preprocessing of TB-scale datasets with an easy to use API that mimics pandas
Parallel model inference (with the option of multi-node GPU inference)

In this post we show an example of data processing from the FineWeb dataset, using the FineWeb-Edu classifier to identify web pages with high educational value. We’ll show:

How to process 100 rows locally with pandas
Scaling to 211 million rows with Dask across multiple GPUs on the cloud

Processing 100 Rows with Pandas

The FineWeb dataset consists of 15 trillion tokens of English web data from Common Crawl, a non-profit that hosts a public web crawl dataset updated monthly. This dataset is often used for a variety of tasks such as large language model training, classification, content filtering, and information retrieval across a variety of sectors.

It can take > 1 minute to download and read in a single file with pandas on a laptop.

import pandas as pd

df = pd.read_parquet(
    "hf://datasets/HuggingFaceFW/fineweb/data/CC-MAIN-2024-10/000_00000.parquet"
)

Next, we’ll use the HF FineWeb-Edu classifier to judge the educational value of the web pages in our dataset. Web pages are ranked on a scale from 0 to 5, with 0 being not educational and 5 being highly educational. We can use pandas to do this on a smaller, 100-row subset of the data, which takes ~10 seconds on a M1 Mac with a GPU.

from transformers import pipeline

def compute_scores(texts):
    import torch

    # Select which hardware to use
    if torch.cuda.is_available():
        device = torch.device("cuda")
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
    else:
        device = torch.device("cpu")

    pipe = pipeline(
        "text-classification",
        model="HuggingFaceFW/fineweb-edu-classifier",
        device=device
    )
    results = pipe(
        texts.to_list(),
        batch_size=25,                    # Choose batch size based on data size and hardware
        padding="longest",
        truncation=True,
        function_to_apply="none"
    )
    
    return pd.Series([r["score"] for r in results])

df = df[:100]
min_edu_score = 3
df["edu-classifier-score"] = compute_scores(df.text)
df = df[df["edu-classifier-score"] >= min_edu_score]

Note that we also added a step to check the available hardware inside the compute_scores function, because it will be distributed when we scale with Dask in the next step. This makes it easy to go from testing locally on a single machine (either on a CPU or maybe you have a MacBook with an Apple silicon GPU) to distributing across multiple machines (like NVIDIA GPUs).

Scaling to 211 Million Rows with Dask

The entire 2024 February/March crawl is 432 GB on disk, or ~715 GB in memory, split up across 250 Parquet files. Even on a machine with enough memory for the whole dataset, this would be prohibitively slow to do serially.

To scale up, we can use Dask DataFrame, which helps you process large tabular data by parallelizing pandas. It closely resembles the pandas API, making it easy to go from testing on a single dataset to scaling out to the full dataset. Dask works well with Parquet, the default format on Hugging Face datasets, to enable rich data types, efficient columnar filtering, and compression.

import dask.dataframe as dd

df = dd.read_parquet(
    # Load the full dataset lazily with Dask
    "hf://datasets/HuggingFaceFW/fineweb/data/CC-MAIN-2024-10/*.parquet" 
)

We’ll apply the compute_scores function for text classification in parallel on our Dask DataFrame using map_partitions, which applies our function in parallel on each pandas DataFrame in the larger Dask DataFrame. The meta argument is specific to Dask, and indicates the data structure (column names and data types) of the output.

from transformers import pipeline

def compute_scores(texts):
    import torch

    # Select which hardware to use
    if torch.cuda.is_available():
        device = torch.device("cuda")
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
    else:
        device = torch.device("cpu")

    pipe = pipeline(
        "text-classification",
        model="HuggingFaceFW/fineweb-edu-classifier",
        device=device,
    )
    results = pipe(
        texts.to_list(),
        batch_size=768,
        padding="longest",
        truncation=True,
        function_to_apply="none",
    )

    return pd.Series([r["score"] for r in results])

min_edu_score = 3
df["edu-classifier-score"] = df.text.map_partitions(compute_scores, meta=pd.Series([0]))
df = df[df["edu-classifier-score"] >= min_edu_score]

Note that we’ve picked a batch_size that works well for this example, but you’ll likely want to customize this depending on the hardware, data, and model you’re using in your own workflows (see the HF docs on pipeline batching).

Now that we’ve identified the rows of the dataset we’re interested in, we can save the result for other downstream analyses. Dask DataFrame automatically supports distributed writing to Parquet. However, since Hugging Face uses commits to track dataset changes, we needed a custom function to allow writing a Dask DataFrame in parallel to Hugging Face storage.

from dask_hf import to_parquet

to_parquet(
    df,
    "hf://datasets/<your-hf-user>/<data-dir>"  # Update with your dataset location
)

In the future, we’ll create a direct integration with Dask and huggingface_hub, but for now, you can copy + paste this custom function for your own use. In this example, we’ve saved it to a file called dask_hf.py, which is the file referenced in the import clause above.

Multi-GPU Parallel Model Inference

There are a number of ways to deploy Dask on a variety of hardware. Here, we’ll use Coiled to deploy Dask on the cloud so we can spin up VMs as needed, and then clean them up when we’re done.

cluster = coiled.Cluster(
    region="us-east-1",                 # Same region as data
    n_workers=100,                      
    spot_policy="spot_with_fallback",   # Use spot instances, if available
    worker_vm_types="g5.xlarge",        # NVIDIA A10 Tensor Core GPU
    worker_options={"nthreads": 1},
)
client = cluster.get_client()

Under the hood Coiled handles:

Provisioning cloud VMs with GPU hardware. In this case, g5.xlarge instances on AWS.
Setting up the appropriate NVIDIA drivers, CUDA runtime, etc.
Automatically installing the same packages you have locally on the cloud VM with package sync. This includes Python files in your working directory, so we can import directly from dask_hf.py on the remote cluster.

The workflow took ~5 hours to complete and we had good GPU hardware utilization.

GPU utilization and memory usage are both near their maximum capacity, which means we’re utilizing the available hardware well.

Putting it all together, here is the complete workflow:

import dask.dataframe as dd
from transformers import pipeline
from dask_hf import to_parquet
import os
import coiled

cluster = coiled.Cluster(
    region="us-east-1",
    n_workers=100,
    spot_policy="spot_with_fallback",
    worker_vm_types="g5.xlarge",
    worker_options={"nthreads": 1},
)
client = cluster.get_client()
cluster.send_private_envs(
    {"HF_TOKEN": "<your-hf-token>"}             #  Send credentials over encrypted connection
)

df = dd.read_parquet(
    "hf://datasets/HuggingFaceFW/fineweb/data/CC-MAIN-2024-10/*.parquet"
)

def compute_scores(texts):
    import torch

    # Select which hardware to use
    if torch.cuda.is_available():
        device = torch.device("cuda")
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
    else:
        device = torch.device("cpu")

    pipe = pipeline(
        "text-classification",
        model="HuggingFaceFW/fineweb-edu-classifier",
        device=device
    )
    results = pipe(
        texts.to_list(),
        batch_size=768,
        padding="longest",
        truncation=True,
        function_to_apply="none"
    )

    return pd.Series([r["score"] for r in results])

min_edu_score = 3
df["edu-classifier-score"] = df.text.map_partitions(compute_scores, meta=pd.Series([0]))
df = df[df["edu-classifier-score"] >= min_edu_score]

to_parquet(
    df,
    "hf://datasets/<your-hf-user>/<data-dir>"               # Replace with your HF user and directory
)

Conclusion

Hugging Face + Dask is a powerful combination. In this example, we scaled up our classification task from 100 rows to 211 million rows by using Dask + Coiled to run the workflow in parallel across multiple GPUs on the cloud.

This same type of workflow can be used for other use cases like:

Filtering genomic data to select genes of interest
Extracting information from unstructured text and turning them into structured datasets
Cleaning text data scraped from the internet or Common Crawl
Running multimodal model inference to analyze large audio, image, or video datasets

Large Scale Geospatial Benchmarks

2024-09-09T00:00:00+00:00

Sep 9, 2024

3m read

James Bourbeau, Matt Rocklin

TL;DR: We need your help creating a geospatial benchmark suite. Please propose your workload on this GitHub discussion.

People love the Xarray/Dask/… software stack for geospatial workloads, but only up to about the terabyte scale. At the terabyte scale this stack can struggle, requiring expertise to work well and frustrating users and developers alike.

To address this, we want to build a large-scale geospatial benchmark suite of end-to-end examples to ensure that these tools operate smoothly up to the 100-TB scale.

We need your help.

Benchmarks

Dask developers have addressed similar issues of robust performance at scale for dataframes and tabular use cases by running existing industry standard benchmarks (see TPC-H benchmarks). Having these representative workloads was invaluable and enabled us to improve Dask DataFrame with much greater confidence.

GitHub comment with A/B test results for a Dask update on TPC-H benchmark performance.

This process has led to Dask DataFrame experiencing an average of 20x performance improvements at scale over the last year, which has been transformative to DataFrame users.

Performance improvements for Dask DataFrames

Over the last month we’ve shifted focus to arrays, particularly supporting Xarray geospatial users, in hopes of getting similar speedups and robustness at large scale. To aid in this endeavor, we’d like to develop a similar suite of large-scale, representative geospatial workflows, and we’d like your help.

The rest of this post will go through what we’re looking for.

What makes a good benchmark

There are many kinds of benchmarks ranging from microbenchmarks like the following:

%timeit x = np.random.random(...).mean()

to full, end-to-end benchmarks like the TPC-H benchmark suite:

import dask.dataframe as dd

date1 = datetime.strptime("2024-10-01", "%Y-%m-%d")
date2 = datetime.strptime("2024-07-01", "%Y-%m-%d")

line_item_ds = dd.read_parquet(dataset_path + "lineitem")
orders_ds = dd.read_parquet(dataset_path + "orders")

lsel = line_item_ds.l_commitdate < line_item_ds.l_receiptdate
osel = (orders_ds.o_orderdate < date1) & (orders_ds.o_orderdate >= date2)
flineitem = line_item_ds[lsel]
forders = orders_ds[osel]
jn = forders.merge(
    flineitem, how="leftsemi", left_on="o_orderkey", right_on="l_orderkey"
)
result = (
    jn.groupby("o_orderpriority")
    .size()
    .to_frame("order_count")
    .reset_index()
    .sort_values(["o_orderpriority"])
)

We’ve learned that microbenchmarks, or even normal-sized benchmarks, don’t provide the same high-quality direction that full end-to-end benchmarks provide. Historically we’ve done benchmarks that focused on specific parts of a workflow like rechunking, aggregations, slicing, and so on. However, while these simple benchmarks help us to accelerate parts of the system, they end up not having that big of an effect overall (performance at scale ends up being less about doing one thing well, and much more about doing nothing poorly).

What was transformative about the TPC-H DataFrame benchmarks for us was that they covered all aspects of a computation at once, from data loading, to a sequence of complex calculations, to collecting and plotting results, putting everything into perspective. These holistic benchmarks focused us in surprising directions, and uncovered flaws we wouldn’t have thought to look for. For example we spent far more time working on data ingestion than we thought, uncovered complex interactions between seemingly unrelated operations, and spent far less time working on in-memory performance than we expected.

As well as being full end-to-end, we think that benchmarks should be…

Relatively simple and easy to understand without sacrificing fidelity (we want other users to look at these benchmarks and see their own applications reflected within them)
Comprehensive and cover the space fairly well
Moderately independent (we don’t want many of the same kind of workflow)
Big (we’re mostly focused on TB or larger computations today)

Initial ideas

We have a couple of end-to-end examples today that we like:

National Water Model (blog, script)
Filtering satellite images (blog)
…

We like that they operate on real data, that they’re large (250TB and 1TB, respectively), and that we’ve seen lots of people do similar things (we share these blogposts frequently).

Each has also helped point us towards important rough edges at large scale that we weren’t previously aware of, improving a variety of tools in the PyData ecosystem including Xarray, Dask, and s3fs.

We want your help

We want your help to build a catalog of large scale, end-to-end, representative benchmarks. What does this help look like? We can use:

Ideas of what are the most common workflows in this space like: “People often need to take netCDF files that are generated hourly, rechunk them to be arranged spatially, and then update that dataset every day”.
Real datasets to work on. We want to work with real data, not fake data.
Real code that does it. We don’t know the space well enough to write like a user here.

This is a big ask, we know, but we hope that if a few people can contribute something meaningfully then we’ll be able to contribute code changes that accelerate those workflows (and others) considerably.

We’d welcome contributions as comments to this GitHub discussion.

DataFrames at Scale Comparison: TPC-H

2024-05-14T00:00:00+00:00

May 14, 2024

14m read

Hendrik Makait, Sarah Johnson, Matthew Rocklin

We run benchmarks derived from the TPC-H benchmark suite on a variety of scales, hardware architectures, and dataframe projects, notably Apache Spark, Dask, DuckDB, and Polars. No project wins.

This post analyzes results within each project and between projects. To start, here is the raw data of the results:

There is too much data here to understand all at once, so this post breaks down results by project and by comparing different projects.

Project results: How each project performs individually
- Spark
- Dask
- DuckDB
- Polars
Comparisons between Projects: Head-to-head comparisons
Experimental Details: How we ran these benchmarks

When to use DuckDB vs. Polars vs. Dask vs. Spark

Deciding which tool is right for you depends on many factors, some of which are objectively quantifiable (speed, scale), some of which are not (language preference, workload type, deployment pain).

In this section we’ll go project by project, explaining how each did.

Spark
Dask
DuckDB
Polars

First though, we’ll include an arbitrary table and some general subjective takes.

*DuckDB and Dask are the most robust choices today. Polars is developing rapidly, though, and we recommend their benchmarks for the latest performance details.*
	Spark	Dask	DuckDB	Polars
Fast Locally	❌	🤔	✅	✅
Fast on Cloud (1 TB)	✅	✅	✅	❌
Fast on Cloud (10 TB)	❌	✅	✅	❌
Scales Out	✅	✅	❌	❌
SQL	✅	🤔	✅	🤔
More than SQL	✅	✅	❌	✅
Sensible Defaults	❌	✅	✅	✅

Local + small? Choose DuckDB or Polars. For small data (10 GB or less) Polars and DuckDB (or pandas) are great. Conversely, Dask and Spark leave a lot of performance on the table in this configuration.

Local + big? Choose Dask or DuckDB. For larger data, Polars isn’t yet fully stable. Polars is improving quickly, although, and is the most dynamic project in this group. DuckDB is more reliable (so is Dask, but it’s slower). Spark fails here pretty easily. If you want to perform SQL queries on 100 GB of local data, DuckDB is the clear winner here.

Cloud at 1 TB or less? Choose Dask or DuckDB. For datasets that fit on a single cloud VM, Dask and DuckDB usually outperform Spark and Polars, respectively. Performance between the fast-local systems (Polars/DuckDB) and the distributed systems (Dask/Spark) converge a lot when we move to the cloud, largely due to S3 access bottlenecks (it doesn’t matter if you’re lightning fast if you’re waiting on S3 for data access). DuckDB does well across scales for the same hardware, but can’t scale beyond a single VM.

Cloud at large scale? Choose Dask or Spark. DuckDB can be as efficient as Dask or Spark on a single big VM, but if you want to process large volumes quickly, you eventually need a scalable system. Dask and Spark can outperform DuckDB and Polars on larger problems by adding more hardware. For large scale, Dask often outperforms Spark and scales well beyond DuckDB and Polars.

How does Spark do on TPC-H Benchmarks?

Spark is the de-facto standard technology in this space today, it is also the least performant in general. Spark’s days seem to be numbered.

Spark is able to complete queries at all scales, but is out-performed at large scale by Dask (by about 2x), and very out-performed at small scale by everyone (by about 10x+). Spark also performs poorly when given far less memory than the data it has to process, as was the case in the 10 TiB / 128 core case (which we included mostly so we could directly compare Spark against DuckDB, which is restricted to run on one machine).

Anecdotally, using Spark was also more painful than using the other systems. We spent more time configuring Spark than we spent on all other systems combined. Even then, we’re uncertain that we have it running optimally. Spark’s performance was highly sensitive to non-default settings (like explicitly setting an output number of shuffle partitions). Additionally we found that Spark was surprisingly inefficient. It had the highest CPU occupancy rates of any of the projects, while also some of the slowest performance.

Still, Spark’s history makes it a sane choice for any organization looking for a solid and well-known solution. We would recommend it mostly for inertia, and for its broad support among enterprise data software systems. We’re also curious about the various projects that rewrite Spark internals, like Photon from Databricks, and the open source Comet system, which were not evaluated here.

How does Dask do on TPC-H Benchmarks?

Disclaimer: We like Dask and know how to use it better than the other frameworks.

Locally, Dask is outperformed at small scale by DuckDB and Polars, but performs reliably across all scales. Dask performs universally better than Spark, and is more robust than Polars and DuckDB even if it’s slower at small scales with fast disk. Dask consistently outperforms other technologies at scale in the cloud, particularly as we increase the amount of hardware available.

Dask was the most robust system among the configurations tested (although again, there’s bias here because we’re better at Dask than at other systems).

Compared to Spark, Dask was easy to configure. Dask automatically inspected local hardware and configured itself appropriately. Additionally Dask’s out-of-core shuffling algorithms were stable even when data sizes were far larger than memory (compare against the 10 TiB 128 core case).

How does DuckDB do on TPC-H Benchmarks?

DuckDB performs commendably at all scales, both locally and in the cloud. It can be occasionally outperformed locally (by Polars, see comparisons below) and in the cloud (by Dask) but is a robust choice across all platforms.

However, two things stop us from recommending DuckDB universally (aside from self-interest):

Scale Limitations: If you need to do fast queries on 10+TB-scale data, DuckDB’s inability to use multiple machines is a limitation. Dask and Spark can easily get 5x speed increases over DuckDB at a 10x hardware cost difference.
SQL-only: Often people need more than just SQL (arguably this explain’s Python’s popularity today) and DuckDB is just SQL.

Using DuckDB was easy (it was pretty much plug-and-play). It is the only framework where we did not feel the urge to communicate with experts of the software, whereas we collaborated with Polars and Spark developers in producing these numbers (see the experiment details section below). It was however mildly inconvenient to add configuration to access data on S3, but documentation and LLMs sufficed to point us in the right direction. There are some small and straightforward things DuckDB can do to reduce friction for lazy users (and honestly, we’re all pretty lazy).

Kudos to the DuckDB team.

How does Polars do on TPC-H Benchmarks?

Updated June 17, 2025: The Polars streaming engine has undergone significant rewrites since these benchmarks were run. We recommend the Polars benchmarks for the latest performance details.

Polars does fantastically well at modest scale (10 GB) locally, where it can use fast local disk and fit tables in memory. The only contender at this scale is DuckDB. However, Polars struggles in a few situations:

Its cloud data readers are not yet very performant
Its unable to handle large multi-table joins well, which are unfortunately common in this benchmark suite

Additionally, when Polars fails, it tends to fail hard, often resulting in machines grinding to a halt. We unreservedly recommend Polars for the 10 GB scale on local machines in memory, and are excited to see how Polars develops in the future (the project is moving quickly).

Using Polars was an interesting case. It both broke frequently, but was also rapidly fixed by the Polars maintainers. This gave the impression that it’s not yet ready (again, this is with the experimental streaming backend, not in-memory), but moving very fast. This was especially true when deploying on the cloud, where Polars has not historically seen as much use. Note that we have not run benchmarks with Polars on the cloud for scales above 100 GB because initial attempts showed poor scaling behavior.

TPC-H Experimental Details

We ran the TPC-H benchmarks with Dask, Spark, DuckDB and Polars on a variety of scales locally and on the cloud. Implementation details and instructions on how to run the benchmarks yourself are in the public coiled/benchmarks GitHub repository.

Dask

import coiled

# Get Dask Cluster
cluster = coiled.Cluster(
    n_workers=32,
    worker_vm_types=["m6i.xlarge"]
)
client = cluster.get_client()

Learn more

PySpark

import coiled

# Get Spark Cluster
cluster = coiled.Cluster(
    n_workers=32,
    worker_vm_types=["m6i.xlarge"]
)
spark = cluster.get_spark()

Learn more

DuckDB

import coiled

# Run DuckDB query
@coiled.function(vm_types=["m6i.32xlarge"])
def query():
    connection = ...

Learn more

Polars

import coiled

# Run Polars Query
@coiled.function(vm_types=["m6i.32xlarge"])
def query():
    df = pl.read.parquet(...)
    ...

Learn more

While infrastructure setup was easy, no project “just worked” efficiently and everything required some degree of looking up documentation, tweaking configuration parameters to get decent performance at scale, passing credentials around, etc., but some more than others (gah! Spark.) Some projects degraded decently even when not optimally configured (DuckDB and Dask, nice job) while others failed ungracefully, grinding machines to a halt (Spark and Polars).

Bias

We’ve done our best to avoid intentionally tuning for our favorite project. We’ve chosen benchmarks chosen by other groups, run everything on standard hardware, and have asked the following experts in other systems to help ensure good performance for the different systems:

Polars: Ritchie Vink
Spark: Jeffrey Chou and Vinoo Ganesh of Sync Computing
DuckDB: no one, but this didn’t seem necessary (they do pretty well!)

Although of course everyone’s time is limited, and asking experts for input makes no guarantee that we’re still not missing important opportunities for performance in those projects.

Dataset Layout

The datasets used in our benchmarks are derived from the datasets generated for the TPC-H benchmarks at the corresponding scale factors stored in Parquet format. The scale reflects the total size of the entire dataset. To generate the data, we used the DuckDB TPC-H extension to generate the data and PyArrow to write partitions of ~100 MB each to disk. The resulting partitions are partially inhomogenous, which reflects the datasets we encounter in practice. All files are compressed with Snappy. See generate_data.py for details.

Measurements

We report three main metrics:

Runtime: This is the end-to-end runtime from generating the query to receiving the result (in seconds). This is our primary measurement.

Relative Difference in Runtime: This metric shows how large the performance difference between two systems is (in percent). We calculated the relative difference in runtime betwen two systems by dividing the absolute runtime difference by the runtime of the faster system:

\[\text{Relative Difference in Runtime} = abs(\frac{\text{Runtime}_X - \text{Runtime}_Y}{min(\text{Runtime}_X, \text{Runtime}_Y)})\]

Calculated this way, the measurement is symmetric, i.e., the value does not change if we compare X to Y or Y to X, and it is linear, i.e., each multiple of absolute difference will result in a multiple of the relative difference.

Query Completion: This boolean measurement shows whether the system completed the query by returning a result. If the system did not complete the query, it either encountered a hard failure (like processes being killed because they ran out of memory) or it timed out after several hours.

Local Configuration

The local benchmarks were run for scales 10 (10 GB) and 100 (100 GB). We executed them on a MacBook Pro with an M1 (8-core) Apple silicon CPU and 16 GB of RAM and the built-in SSD used for storage. We did our best to create reproducible benchmarks by closing noisy background processes where possible, but your mileage may vary here depending on what background processes you still have running on your machine. The software was installed with conda/mamba with builds from conda-forge. The datasets were stored on local SSD for fast access.

Cloud Configuration

The benchmarks were run on AWS on the cloud for scales 10 (10 GB), 100 (100 GB), 1,000 (1 TB), and 10,000 (10 TB) using EC2 instances from the M6i family (more on why we prefer non-burstable instances). DuckDB and Polars ran on a single, large machine with Coiled serverless functions while Dask and Spark ran on clusters. For scale 10,000 (10 TB), we performed two experiments and scaled the amount of vCPUs from 320 by a factor of 10 to 320. Since there is no instance type with 320 vCPUs available, this experiment was only run on the scale-out capable systems, Dask and PySpark.

*Cloud hardware used for running benchmarks at various scales (see the configuration for more details).*
Scale	vCPUs	Memory	Instances (single)	Instances (cluster)
100 GB	32	128 GiB	1 x `m6i.8xlarge`	16 x `m6i.large`
1 TB	128	512 GiB	1 x `m6i.32xlarge`	32 x `m6i.xlarge`
10 TB	128	512 GiB	1 x `m6i.32xlarge`	32 x `m6i.xlarge`
10 TB	1280	5 TiB	n/a	320 x `m6i.xlarge`

The datasets are stored on S3 in the public requester-pays bucket with subdirectories for each scale (e.g. s3://coiled-runtime-ci/tpc-h/snappy/scale-10/).

Library-Specific Configuration

Spark

For Spark, we manually configured memory limits in order to utilize existing memory while avoiding OOM errors. For cloud-based benchmarks, we configured spark.executor.memory to 80% of the available memory after tweaking for several iterations. For local benchmarks, we configured spark.driver.memory to 10g to avoid collect getting stuck. Locally, we also had to manually set spark.driver.host and spark.driver.bindAddress to 127.0.0.1 since the default does not work for M1 MacBooks. We used SparkSQL to run the benchmarks.

Dask

For Dask, we enabled Copy-on-Write (COW), which will become the default in pandas 3, which will be released soon. We also disabled task queueing. Since Dask is a dataframe-based system, we had to translate the reference SQL queries to idiomatic Dask queries. As we currently lack join order optimization (see dask-expr#1065), we had to manually order joins. For this, we diverged from the order in which tables were introduced in the SQL query to a manually optimized order that a careful user would typically choose when following the dimensions and facts defined in the TPC-H schema.

DuckDB

For DuckDB, we did not perform any configuration apart from passing security credentials.

Polars

Since Polars is a dataframe-based system, the reference SQL queries had to be translated to idiomatic Polars queries. For this, we reused the reference implementations published by Polars. To allow Polars to process larger-than-memory datasets, we executed all queries in streaming mode.

Learn More

This article covers quantitative results on the TPC-H Benchmark across common DataFrame libraries. For information on how to run these projects in the cloud yourself, consider looking at the Coiled documentation for each project:

We (Coiled) care about making data processing in the cloud easy, which involves running the right tool for the right job. That being said, we’re also biased towards Dask. If you’d like to read more about recent Dask performance, or how it compares to Spark today, we invite you to consider the following:

Dask DataFrame is Fast Now

2024-05-14T00:00:00+00:00

May 14, 2024

10m read

Patrick Hoefler

Performance Improvements for Dask DataFrames

Intro

Dask DataFrame scales out pandas DataFrames to operate at the 100GB-100TB scale.

Historically, Dask was pretty slow compared to other tools in this space (like Spark). Due to a number of improvements focused on performance, it’s now pretty fast (about 20x faster than before). The new implementation moved Dask from getting destroyed by Spark on every benchmark to regularly outperforming Spark on TPC-H queries by a significant margin.

Dask DataFrame workloads struggled with many things. Performance and memory usage were commonly seen pain points, shuffling was unstable for bigger datasets, making scaling out hard. Writing efficient code required understanding too much of the internals of Dask.

The new implementation changed all of this. Things that didn’t work were completely rewritten from scratch and existing implementations were improved upon. This puts Dask DataFrames on a solid foundation that allows faster iteration cycles in the future.

We’ll go through the three most prominent changes:

Apache Arrow support
Faster joins
Query optimization

We’ll cover how these changes impact performance and make it easier to use Dask efficiently, even for users that are new to distributed computing. We’ll also discuss plans for future improvements.

1. Apache Arrow Support: Efficient String Datatype

A Dask DataFrame consists of many pandas DataFrames. Historically, pandas used NumPy for numeric data, but Python objects for text data, which are inefficient and blow up memory usage. Operations on object data also hold the GIL, which doesn’t matter much for pandas, but is a catastrophy for performance with a parallel system like Dask.

The pandas 2.0 release introduced support for general-purpose Arrow datatypes, so Dask now uses PyArrow-backed strings by default. These are much better. PyArrow strings reduce memory usage by up to 80% and unlock multi-threading for string operations. Workloads that previously struggled with available memory now fit comfortably in much less space, and are a lot faster because they no longer constantly spill excess data to disk.

Memory Usage of the Legacy DataFrames Compared with Arrow Strings

2. Faster Joins with a New Shuffle Algorithm

Shuffling is an essential component of distributed systems to enable sorting, joins, and complex group by operations. It is an all-to-all, network-intensive operation that’s often the most expensive component in a workflow. Dask has a new shuffling system, which greatly impacts overall performance, especially on complex, data-intensive workloads.

A shuffle operation is intrinsically an all-to-all communication operation where every input partition has to provide a tiny slice of data to every output partition. Dask was already using it’s own task-based algorithm that managed to reduce the O(n * n) task complexity to O(log(n) * n) where n is the number of partitions. This was a drastic reduction in the number of tasks, but the non-linear scaling ultimately did not allow Dask to process arbitrarily large datasets.

Dask introduced a new P2P (peer-to-peer) shuffle method that reduced the task complexity to O(n) which scales linearly with the size of the dataset and the size of the cluster. It also incorporates an efficient disk integration which allows easily shuffling datasets which are much larger than memory. The new system is extremely stable and “just works” across any scale of data.

Memory Usage of the Legacy Shuffle Compared with P2P

3. Optimizer

Dask itself is lazy, which means that it registers your whole query before doing any actual work. This is a powerful concept that enables a lot of optimizations, but historically Dask wasn’t taking advantage of this knowledge in the past. Dask also did a bad job of hiding internal complexities and left users on their own while navigating the difficulties of distributed computing and running large scale queries. It made writing efficient code painful for non-experts.

The Dask release in March includes a complete re-implementation of the DataFrame API to support query optimization. This is a big deal. The new engine centers around a query optimizer that rewrites your code to make it more efficient and better tailored to Dask’s strengths. Let’s dive into some optimization strategies, how they make Dask run faster and scale better.

We will start with a couple general-purpose optimizations:

Column projection
Filter pushdown

And then dive into more specific techniques that are tailored to distributed systems generally and Dask more specifically:

Automatic partition resizing
Trivial merge and join operations

3.1 Column Projection

Most datasets have more columns than needed. Dropping them requires foresight (“What columns will I need for this query? 🤔”) so most people don’t think about this when loading data. This is bad for performance because carrying around lots of excess data slows everything down. Column Projection drops columns as soon as they aren’t needed anymore. It’s a straightforward optimization, but highly beneficial.

The legacy implementation always reads all columns from storage and only drops columns when explicitly specified by the user. Simply operating on less data is a big win for performance and memory usage.

The optimizer looks at the query and figures out which columns are needed for each operation. It looks at the final step of the query and then works backwards step by step to the data source, injecting drop operations to get rid of unnecessary columns.

Only require a subset of columns are needed. Replace doesn’t need access to all columns, so Dask drops unnecessary columns directly in the IO step.

3.2 Filter Pushdown

Filter pushdown is another general-purpose optimization with the same goal as column projection: operate on less data. The legacy implementation did not reorder filter operations. The new implementation executes filter operations as early as possible while maintaining the same results.

The optimizer identifies every filter in the query and looks at the previous operation to see if we can move the filter closer to the data source. It will repeat this until it finds an operation that can’t be switched with a filter. This is a bit harder than column projections, because Dask has to make sure that the operations don’t change the values of the DataFrame. For example, switching a filter and a merge operation is fine (values don’t change), but switching a filter and a replace operation is invalid, because the values might change and rows that would previously have been filtered out now won’t be, or vice versa.

Initially, the filter happens after the Dropna, but Dask can execute the filter before Dropna without changing the result. This allows Dask to push the filter into the IO step.

Additionally, if the filter is strong enough then Dask can potentially drop complete files in the IO step. This is a best-case scenario, where an earlier filter brings a huge performance improvement and even requires reading less data from remote storage.

3.3 Automatically Resizing Partitions

In addition to implementing the common optimization techniques described above, we’ve also improved a common pain point specific to distributed systems generally and Dask users specifically: optimal partition sizes.

Dask DataFrames consist of many small pandas DataFrames called partitions. Often, the number of partitions is decided for you and Dask users are advised to manually “repartition” after reducing or expanding their data (for example by dropping columns, filtering data, or expanding with joins) (see the Dask docs). Without this extra step, the (usually small) overhead from Dask can become a bottleneck if the pandas DataFrames become too small, making Dask workflows painfully slow.

Manually controlling the partition size is a difficult task that we, as Dask users, shouldn’t have to worry about. It is also slow because it requires network transfer of some partitions. Dask DataFrame now automatically does two things to help when the partitions get too small:

Keeps the size of each partition constant, based on the ratio of data you want to compute vs. the original file size. If, for example, you filter out 80% of the original dataset, Dask will automatically combine the resulting smaller partitions into fewer, larger partitions.
Combines too-small partitions into larger partitions, based on an absolute minimum (default is 75 MB). If, for example, your original dataset is split into many tiny files, Dask will automatically combine them.

Select two columns that take up 40 MB of memory out of the 200 MB from the whole parquet file.

The optimizer will look at the number of columns and the size of the data within those. It calculates a ratio that is used to combine multiple files into one partition.

The ratio of 40/200 results in combining five files into a single partition.

This step is currently limited to IO operations (like reading in a parquet dataset), but we plan to extend it to other operations that allow cheaply combining partitions.

3.4 Trivial Merge and Join Operations

Merge and join operations are typically cheap on a single machine with pandas but expensive in a distributed setting. Merging data in shared memory is cheap, while merging data across a network is quite slow, due to the shuffle operations explained earlier.

This is one of the most expensive operations in a distributed system. The legacy implementation triggered a network transfer of both input DataFrames for every merge operation. This is sometimes necessary, but very expensive.

Both joins are performed on the same column. The left DataFrame is already properly partitioned after the first join, so Dask can avoid shuffling again with the new implementation.

The optimizer will determine when shuffling is necessary versus when a trivial join is sufficient because the data is already aligned properly. This can make individual merges an order of magnitude faster. This also applies to other operations that normally require a shuffle like groupby().apply().

Dask merges used to be inefficient, which caused long runtimes. The optimizer fixes this for the trivial case where these operations happen after each other, but the technique isn’t very advanced yet. There is still a lot of potential for improvement.

The current implementation shuffles both branches that originate from the same table. Injecting a shuffle node further up avoids one of the expensive operations.

The optimizer will look at the expression and inject shuffle nodes where necessary to avoid unnecessary shuffles.

How do the improvements stack up compared to the legacy implementation?

Dask is now 20x faster than before. This improvement applies to the entire DataFrame API (not just isolated components), with no known performance regressions. Dask now runs workloads that were impossible to complete in an acceptable timeframe before. This performance boost is due to many improvements all layered on top of each other. It’s not about doing one thing especially well, but about doing nothing especially poorly.

Performance Improvements on Query 3 of the TPC-H Benchmarks from coiled/benchmarks

Performance, while the most enticing improvement, is not the only thing that got better. The optimizer hides a lot of complexity from the user and makes the transition from pandas to Dask a lot easier because it’s now much more difficult to write poorly performing code. The whole system is more robust.

The new architecture of the API is a lot easier to work with as well. The legacy implementation leaked a lot of internal complexities into high-level API implementations, making changes cumbersome. Improvements are almost trivial to add now.

What’s to come?

Dask DataFrame changed a lot over the last 18 months. The legacy API was often difficult to work with and struggled with scaling out. The new implementation dropped things that didn’t work and improved existing implementations. The heavy lifting is finished now, which allows for faster iteration cycles to improve upon the status quo. Incremental improvements are now trivial to add.

A few things that are on the immediate roadmap:

Auto repartitioning: this is partially implemented, but there is more potential to choose a more efficient partition size during optimization.
Faster Joins: there’s still lots of fine-tuning to be done here. For example, there is a PR in flight with a 30-40% improvement.
Join Reordering: Dask doesn’t do this yet, but it’s on the immediate roadmap

Learn more

This article focuses on a number of improvements to Dask DataFrame and how much faster and more reliable it is as a result. If you’re choosing between Dask and other popular DataFrame tools, you might also consider:

DataFrames at Scale Comparison: TPC-H which compares Dask, Spark, Polars, and DuckDB performance on datasets ranging from 10 GB to 10 TB both locally and on the cloud

Dask vs. Spark

2024-04-19T00:00:00+00:00

Apr 19, 2024

6m read

Sarah Johnson, Florian Jetter

People often ask “How does Dask compare to Spark?”

Apache Spark has been the industry standard for big data, although it’s not clear that anyone especially loves it, especially in the Python space. The promise of Spark is that it won’t topple over no matter how large your data gets, and for years it’s been the best big data tool available.

Today, though, there are a number of modern solutions for churning through TBs of data. Dask, in particular, is an attractive alternative to Spark because it’s:

Performant at large scale (and faster than Spark)
Easy to use (and easier to debug)
Highly flexible (does everything from ETL to machine learning)

We’re Dask fans and this post isn’t shy about showing that bias. For a less biased treatment of this same topic, see the OSS Dask documentation page where we’ve tried to keep a more balanced tone.

Dask is Faster than Spark

We compared Dask and Spark on the TPC-H benchmark suite and can confidently claim that Dask is not only easier to use, but often faster and more reliable than Spark. See DataFrames at Scale Comparison: TPC-H which compares Dask, Spark, Polars, and DuckDB performance to learn more.

We’ll show results from running the TPC-H benchmarks locally on a 10 GB dataset and on the cloud on a 10 TB dataset. These two cases are important because developers typically want to iterate locally on a small sample before scaling out to the full dataset on a production job on the cloud.

Dask consistently outperforms PySpark on a 10 GB dataset that fits on a typical laptop. Dask is up to 507% faster than Spark.

When running the TPC-H benchmarks locally on an M1 MacBook Pro, Dask is up to 507% faster than Spark. Since Dask is faster locally, this makes it easier for developers to iterate quickly.

There were some challenges with all queries successfully completing when scaling the benchmarks to a 10 TB dataset. Originally, we ran these queries on an underprovisioned cluster (32 x m6i.xlarge / 128 CPUs / 512 GB). Dask failed on three of the 22 queries, but handled everything else gracefully (albeit slowly). Spark, however, failed on all queries, either due to running out of memory, running out of disk, or running into an hour-long timeout.

Dask can finish many queries that PySpark cannot on a cluster with 5 TB of memory. PySpark was unable complete seven of the 22 queries (shown in grey).

We also ran the benchmarks on a cluster with 5 TB of memory and 1280 CPU cores. This should be sufficient, since systems should be able to deal with data that is modestly larger than RAM. While Dask completed 21 of 22 queries, Spark either failed or ran into an hour-long timeout on seven queries. This is problematic because it causes engineers to over-provision clusters with too large machines by habit, ultimately increasing cost.

Relative difference between Dask vs PySpark. Dask is up to 112% faster than Spark for queries they both completed.

For queries that Dask and PySpark completed, Dask is often faster. This is largely due to a number of engineering improvements in Dask, see Dask DataFrame is Fast Now to learn more.

Dask is easier to use than Spark

People are both slower and more expensive than computers, so we optimize for people. Users frequently praise how much easier Dask is to use than Spark.

A lot of the Spark internals are run on the JVM, so it’s actually not that easy to debug. You spend a lot of time on Stack Overflow trying to figure out what did I do wrong, why am I getting this traceback. It was painful and we wanted to use something that was not as difficult to actually debug.

- Sébastien Arnaud, CDO/CTO at Steppingblocks, on why they switched from Spark to Dask

Easy to learn

Dask is well-integrated in the larger Python ecosystem; it complements and enhances other popular libraries like pandas, NumPy, and scikit-learn. Many users are already familiar with Python for handling large datasets. Dask looks and feels a lot like pandas, so it’s easier to get started quickly.

import pandas as pd
import dask.dataframe as dd

# use pandas when your data fit in memory
df = pd.read_parquet("small-data.parquet")

# use Dask when you can't read in your data on a single machine
df = dd.read_parquet("large-data.parquet")

Easy to set up

Historically, Spark had better support for enterprise data sources like Delta Lake and Snowflake. Dask has caught up though, and today with projects like delta-rs (the Rust re-implementation of Delta Lake) and Snowflake’s bulk read/write capabilities, Dask gives a first-class user experience on-par with any other technology.

Delta Lake

You can create a Delta Lake table from a pandas dataframe and use dask_deltatable to read in a Delta Lake table into Dask DataFrame.

import deltalake
import pandas as pd

df = pd.read_csv("...")
deltalake.write_table("mytable", df, mode="append")

import dask_deltatable

df = dask_deltatable.read_deltalake("mytable", datetime="2020-01-01")

Snowflake

With dask_snowflake, you can read from a Snowflake table into a Dask DataFrame.

import dask_snowflake

example_query = """
    SELECT *
    FROM SNOWFLAKE_SAMPLE_DATA.TPCH_SF1.CUSTOMER;
"""
df = read_snowflake(
    query=example_query,
    connection_kwargs={
        "user": "...",
        "password": "...",
        "account": "...",
    },
)

Easy to Debug

Debugging Spark is painful, often accounting for most of the time spent with the technology. In contrast, Dask works hard to make this experience almost as easy as it is in a single-machine context. Dask users especially love the Dask dashboard, which gives them a lot of understanding about how their computation is going.

Easy to Deploy

Deploying Spark, or even configuring it locally, is a pain. Spark users often mention it requires digging through docs and finding obscure configuration settings just to get anything running. We really don’t see it in the wild, except under some SaaS deployment solution like Databricks, EMR, Synapse, Dataproc, etc.

Dask on the other hand has SaaS solutions like Coiled, but is also often deployed “raw” using systems like Kubernetes or HPC job managers (see the Dask Deployment docs). We’re a bit biased, but the main SaaS solution, Coiled, is way easier to use than any of the Spark SaaS solutions.

Dask does more than Spark

The first questions you ask yourself are:

“Can this tool actually solve my problem?”

“How painful will it be?”

For Spark, if your problem is SQL queries, DataFrames, or streaming analytics then the answer is “yes, this tool definitely solves my problem”. For Dask these all apply as well, and also more applications like:

Training ML models with XGBoost or PyTorch
HyperParameter optimization with Optuna
Workflow management with tools like Prefect or Dagster
Highly custom parallelism and ad-hoc workflows
Wrangling large multi-dimensional data, such as in imaging or processing simulation output

Internally, Spark is fundamentally a MapReduce engine. It applies the same function across many inputs uniformly, shuffles those results around, and reduces the results. This pattern is both simple and powerful, and lets Spark build up abstractions like DataFrames and SQL.

Map-Shuffle-Reduce: Apache Spark always operates on lots of data in bulk using the same operations.

However, many applications (like those listed above) can’t be written down as a MapReduce style problem. These problems can not be efficiently solved with Spark, at least not without a lot of pain, and so people either build their own solution, or they build on tools like Dask.

Instead of using MapReduce, Dask follows a task scheduling approach. This provides more flexibility, which has made Dask a powerful backend for a modern generation of tools, ranging from ML training systems like XGBoost, to workflow managers like Prefect and Dagster, all of which can back themselves with Dask clusters for scalable computing. Additionally, disciplines with a high degree of complexity, like finance and geoscience, prefer Dask because they can build arbitrarily complex computations on a fine-grained task-by-task basis.

Task graph of a credit model. Dask uses dynamic task scheduling, which can handle weirder situations (which come up a lot in Python workloads).

Conclusion

In general, both Spark and Dask can probably do what you want (unless you’re doing something very strange). Speaking with our Dask bias, we’re pretty proud of the decisions and tradeoffs Dask has made and are excited to see Dask continue to make impressive performance gains.

It’s also pretty easy to try out for yourself, either locally or on the cloud:

Local

Dask can set itself up easily in your Python session with a LocalCluster.

from dask.distributed import LocalCluster

cluster = LocalCluster()
client = cluster.get_client()

Cloud

It’s easy to sign up for Coiled to run Dask on the cloud.

from coiled import Cluster

cluster = Cluster(n_workers=100)
client = cluster.get_client()

Once you have a Dask cluster, you can run Python code on that cluster. Here are some examples:

Easy Scalable Production ETL

2024-04-08T00:00:00+00:00

We show a lightweight scalable data pipeline that runs large Python jobs on a schedule on the cloud.

Data	Cost
575 million records / day	$10/day

It’s common to run regular large-scale Python jobs on the cloud as part of production data pipelines. Modern workflow orchestration systems like Prefect, Dagster, Airflow, Argo, etc. all work well for running jobs on a regular cadence, but we often see groups struggle with complexity around cloud infrastructure and lack of scalability.

Here we show a scalable data pipeline that runs regular jobs on the cloud with Coiled and Prefect. The pipeline consists of a few Python files in this GitHub repository which you can run yourself either locally or on the cloud. This approach is:

Easy to deploy on the cloud
Scalable across many cloud machines
Cheap to run on ephemeral VMs and a small always-on VM

Use Python

Defining tasks and cloud hardware in Python makes it easy for Python data developers to write and deploy their processing code quickly.

import coiled
from prefect import task, flow
import pandas as pd

@task
@coiled.function(region="us-east-2", memory="64 GiB")  # <--- Define hardware in script
def json_to_parquet(filename):
    df = pd.read_json(filename)                        # <--- Use simple pandas functions
    df.to_parquet(OUTDIR / filename.with_suffix(".parquet").name)

@flow
def convert_all_files():
    files = list_files()
    json_to_parquet.map(files)                         # <--- Map across many files at once

Run jobs regularly

Prefect runs Python jobs on a regular cadence. Coiled deploys Prefect on a cheap, always-on t3.medium instance (~$1/day) on AWS.

from datetime import timedelta
from prefect import serve

preprocess = json_to_parquet.to_deployment(
    name="preprocess",
    interval=timedelta(minutes=15),             # <--- Run job on a regular schedule
)
reduce = query_reduce.to_deployment(
    name="reduce",
    interval=timedelta(hours=24),               # <--- Different jobs at different schedules
)

...

serve(data, preprocess, reduce, ...)

Scale with Dask

Coiled creates on-demand Dask clusters for processing larger-than-memory datasets. This lets you scale out when you need to and only pay for hardware while you’re using it.

import coiled
from prefect import task, flow
import dask.dataframe as dd

@task
def unshipped_orders_by_revenue(bucket):
    with coiled.Cluster(
        n_workers=200,                  # <--- Ask for hundreds of workers
        region="us-west-2",             # <--- Run anywhere on any hardware
    ) as cluster:
        with cluster.get_client() as client:
            df = dd.read_parquet(bucket)
            result = ...   # Complex query
            result.to_parquet(...)

Data pipelines

How do I process large-scale data continuously in production?

While details will vary across use cases, data pipelines typically look something like:

Step 1: Acquire raw data from the real world
- Example: A customer buys a loaf of bread at your grocery store
- Example: A satellite measurements sea temperature in Antarctica
Step 2: Regularly process that raw data
- Example: Apply data quality cuts followed by a feature engineering step
- Example: Convert to a different storage format
- Extra requirement: Retry on failures, run close to data in the cloud, alert me if processing fails
Step 3: Scalably analyze that data periodically
- Example: Large-scale query of your data that addresses a business question
- Example: Train an ML model for weather forecasting
- Extra requirement: Data is too large to fit on a single machine
Step 4: Publish results
- Example: Serve an interactive dashboard with a real-time view of customer orders
- Example: Host an ML model

Running these steps on a regular cadence, task retries, and alerts are all handled well by modern workflow management systems like Prefect, Dagster, Airflow, Argo, etc. Managing cloud infrastructure and scalability is where most people tend to struggle.

Pain points of common data pipelines

Pipelines can be overly complex or lack scalability.

Workflow management systems tend to struggle with:

Complexity around managing cloud infrastructure:
- Provisioning / deprovisioning cloud machines
- Software environment management
- Handling cloud data access
- Cost monitoring and spending limits
Scale limitations:
- Scaling existing Python code across many cloud machines in parallel
- Computing on larger-than-memory datasets (e.g. 100 TB Parquet dataset in S3)

Because of these issues, it’s common for systems to be overly complex or expensive. Below we show an example that connects Prefect and Coiled to run a lightweight, scalable data pipeline on the cloud.

Easy deployment for data pipelines

Coiled and Prefect make it easy to run regularly scheduled jobs on the cloud at scale.

Coiled is a lightweight tool that makes deploying and scaling data pipelines on the cloud straightforward. This example data pipeline combines Coiled and Prefect to run regularly scheduled jobs on the cloud. It looks like this:

Step 1: Data generation – New TPC-H JSON files with customer orders and supplier information appear in an S3 bucket (every 15 minutes)
Step 2: Regular processing
- JSON gets transformed into Parquet / Delta (every 15 minutes)
- Data compaction of small Parquet files into larger ones for efficient downstream usage (every 6 hours)
Step 3: Scalable analysis – Run large-scale multi-table analysis query to monitor unshipped, high-revenue orders (every 24 hours)
Step 4: Serve dashboard – Results from latest business query are served on a dashboard (always on)

This pipeline processes ~575 million records daily and costs ~$10/day to run on the cloud. Let’s look at how Coiled and Prefect work together to this work.

Define tasks and hardware in Python

Data-intensive Prefect tasks are common for per-file preprocessing steps, like those in pipeline/preprocess.py. We combine them with Coiled Functions for remote execution on cloud VMs. Defining tasks and cloud hardware in Python makes it easy for Python data developers to write and deploy their code quickly.

# pipeline/preprocess.py

import coiled
from prefect import task, flow
import pandas as pd

@task
@coiled.function(
    region="us-east-2",       # <-- Define hardware in script
    memory="64 GiB",
)
def json_to_parquet(filename):
    df = pd.read_json(filename)     # <-- Use simple pandas functions
    df.to_parquet(OUTDIR / filename.with_suffix(".parquet").name)

@flow
def convert_all_files():
    files = list_files()
    json_to_parquet.map(files)     # <-- Run across many files at once

Run jobs regularly

Related tasks are connected in a Prefect flow which defines their dependency to each other. This flow can then be scheduled to run regularly at different cadences, see for example in workflow.py. We deploy this workflow on a cheap, always-on t3.medium instance (~$1/day) using Coiled’s coiled prefect serve CLI.

# workflow.py

from datetime import timedelta
from prefect import serve

preprocess = json_to_parquet.to_deployment(
    name="preprocess",
    interval=timedelta(minutes=15),      # <--- Run job on a regular schedule
)
reduce = query_reduce.to_deployment(
    name="reduce",
    interval=timedelta(hours=24),         # <--- Different jobs at different schedules
)

...

serve(data, preprocess, reduce, ...)

Scale with Dask clusters

The large-scale, multi-table analysis query in pipeline/reduce.py uses Coiled to create an on-demand Dask cluster to handle processing our larger-than-memory data. Once the query is done all the VMs in the cluster are automatically shut down.

# pipeline/reduce.py

import coiled
from prefect import task, flow
import dask.dataframe as dd

@task
def unshipped_orders_by_revenue(bucket):
    with coiled.Cluster(
        n_workers=200,                    # <--- Ask for hundreds of workers
        region="us-west-2",               # <--- Run anywhere on any hardware
    ) as cluster, cluster.get_client() as client:
        df = dd.read_parquet(bucket)
        result = ...   # Complex query
        result.to_parquet(...)

Summary

Data pipelines with regular large-scale Python jobs are common, but can be difficult to deploy on the cloud or don’t scale well. We showed how to:

Use Python for defining tasks and cloud hardware
Run jobs regularly on the cloud using Prefect and coiled prefect serve
Scale with Dask when processing large volumes of data

in a scalable example data pipeline that processes ~575 million records daily and costs ~$10/day to run on the cloud.

Next steps

You can run this example data pipeline yourself, either locally or on the cloud.

To run things locally, make sure you have a Prefect cloud account and have authenticated your local machine. Then clone the pipeline GitHub repo, install dependencies, and run the pipeline:

# Run data pipeline locally
git clone https://github.com/coiled/etl-tpch
cd etl-tpch
mamba env create -f environment.yml
mamba activate etl-tpch
python workflow.py

To run the same pipeline on the cloud, follow these instructions. Running the pipeline for several days fits comfortably within the Coiled free tier. We hope this example is easy to copy and modify for your needs.

You might also be interested in these additional use cases:

One Trillion Row Challenge

2024-02-05T00:00:00+00:00

Feb 5, 2024

4m read

Sarah Johnson

Last month Gunnar Morling launched the One Billion Row Challenge with the task of writing a Java program for retrieving temperature measurement values from a text file and calculating the min, mean, and max temperature per weather station. This took off greater than anyone would expect, gathering dozens of submissions from different tools.

We joined in the fun and made our own unofficial Python submission for Dask, which ran in 32.8 seconds on a MacBook. Fellow Dask maintainer Jacob Tomlinson took things a step further, and ran the challenge on GPUs in 4.5 seconds.

This got us thinking, what about a trillion row challenge? The 1BRC was fun, and it’d be interesting to compare performance across different sets of big data tools at a much larger scale.

How can you participate?

Your task is to use any tool(s) you’d like to calculate the min, mean, and max temperature per weather station, sorted alphabetically. The data is stored in Parquet on S3: s3://coiled-datasets-rp/1trc. Each file is 10 million rows and there are 100,000 files. For an extra challenge, you could also generate the data yourself.

Open an issue in the 1TRC repo with your submission and enough details for someone else to be able to run your implementation. This includes things like:

Hardware
Runtime
Reproducible code snippet

There is no prize and everyone is a winner. Really, the idea is to solicit ideas and generate discussion.

The next sections dive into the dataset generation process in more detail and an example of running the challenge using Dask.

Data Generation

Our goal is to make the dataset easy for you to generate yourself, but also easy to access if you don’t have many TBs of data storage available. Using the 1BRC dataset as a template, we scaled up to one trillion rows, and wrote the data to Parquet on S3 (thanks to Jacob Tomlinson for the data generation script).

You’ll be working with 100,000 Parquet files, each with 10 million rows. The dataset is accessible on S3: s3://coiled-datasets-rp/1trc.

Here’s what you’ll be working with:

100,000 Parquet files. Each files has 10 million rows and is 24 MiB on disk, for a total of 2.5 TiB.
Stored on S3 s3://coiled-datasets-rp/1trc. This is a requester-pays bucket in AWS region US East (N. Virginia). Those using AWS can avoid incurring charges by reading in the data from a cloud VM that’s also running in us-east-1.

Choosing good chunk sizes

The first decision was how to split up a trillion rows of data. There’s not really one right answer here. You need small enough chunks that many could fit on a single worker’s available memory at one time. If the chunks are too small, this could mean an unnecessarily large task graph. If the chunks are too big, then you would be limited to specific high-memory machine types or risk running out of memory. For Dask, this usually means starting with somewhere around 100–500 MiB memory usage per chunk (note, this is when loaded into memory, not the size on disk).

If you read a single file into memory, it takes up ~228 MB, which is a good place to start.

import dask.dataframe as dd

df = dd.read_parquet(
    "s3://coiled-datasets-rp/1trc/measurements-0.parquet"
)
df.info(memory_usage=True)

<class 'dask.dataframe.core.DataFrame'>
Columns: 2 entries, station to measure
dtypes: float64(1), string(1)
memory usage: 228.5 MB

CSV vs. Parquet

We first tried to use text files, but found Dask to be prohibitively slow at this scale. We generated 100,000 CSVs, each 131.6 MiB on disk for a total of 12.5 TiB. Creating the task graph for 1,000 files took ~30 seconds. There is certainly more we could dig into here, but we decided using Parquet is more realistic for working with this scale of data.

CSVs are everywhere and fine solution if you’re not working with much data. Parquet is more performant, though, especially if you have large datasets. This is in part because Parquet is more compact to store and faster to read. Since we’re processing a lot more data than the 1BRC, it makes sense to use Parquet.

Storing data on S3

The dataset is publicly available on AWS S3 in s3://coiled-datasets-rp/1trc in AWS region US East (N. Virginia). The dataset is publicly available in a requester-pays bucket, which means you don’t have to worry about storing the data, but you are responsible for any data transfer or download costs (more details in the AWS documentation).

If you’re using AWS, you can avoid data transfer costs by running your computation in us-east-1, the same region as the data.

1TRC with Dask

Here’s an example of how we ran this challenge using Dask and Coiled. It took 8.5 minutes and cost $3.26 in AWS cloud costs (10 minutes if you include cluster start time).

import dask_expr as dd
import coiled

cluster = coiled.Cluster( # Start Coiled cluster
    n_workers=100,        # 100 m6i.xlarge instances
    region='us-east-1',   # Faster data access and avoid data transfer costs
)
client = cluster.get_client()

df = dd.read_parquet(
    "s3://coiled-datasets-rp/1trc/",
    dtype_backend="pyarrow",
    storage_options={"requester_pays": True}
)
df = df.groupby("station").agg(["min", "max", "mean"])
df = df.sort_values("station").compute()

Similar to the 1BRC submission, most of the time is spent in I/O, in this case on read_parquet. Looking at the task stream it seems the cluster resources are well utilized, since there’s little white space.

Dask dashboard during the 1TRC, most of the time is spent on read_parquet tasks (yellow).

Optimizing costs

Processing 2.5 TB cost us $3.26. This is cheap. We can bring the cost down ~3x to $1.10 by:

Using Spot instances, which are excess capacity offered at a discount.
Using ARM-based instances, which are less expensive than Intel-based equivalents with the similar, or better, performance.

cluster = coiled.Cluster(            # Start Coiled cluster
    n_workers=100,                   # 100 m7g.xlarge instances
    region='us-east-1',              # Faster data access and avoid data transfer costs
    arm=True,                        # Use ARM-based instances
    spot_policy="spot_with_fallback" # Use spot, when available
)

In addition to being less expensive, using more performant ARM-based instances lowers computation runtime to 5.8 minutes (8 minutes if you include cluster start time).

Conclusion

Looking forward to seeing all your challenge submissions in the 1TRC repo. Thanks to Jacob Tomlinson for helping kick things off (looking forward to seeing a GPU implementation soon).

Real-world Grocery Demand Forecasting

2024-01-31T00:00:00+00:00

Jan 31, 2024

5m read

Jack Solomon

Every day, grocers are tasked with predicting the future: they have to predict how many items will sell in the future to order the correct amount of inventory today.

Since grocery retailers sell fresh food, which expires within a few days of being on the shelf, grocers need to be hyper-accurate in predicting item-level demand for each day.

If grocers over-predict, they have tons of unsold food that gets thrown out; today, if food waste was a country it would be the third largest emitter behind the US and China. On the other hand, if grocers under-predict, they’re left with empty shelves, which leads to immediate revenue loss but also long-term a drop in customer loyalty.

Parallel Time Series Forecasting in Python

Fortunately, we can use machine learning algorithms with large, real-world datasets to make SKU-level predictions with greater accuracy and efficiency. In this article, we’ll discuss how we’re using models like XGBoost with Dask and Coiled for time series forecasting with large datasets in Python.

XGBoost (Extreme Gradient Boosting) is a powerful machine learning algorithm that uses an ensemble learning method combining multiple decision trees to make predictions. Dask is a general purpose library for parallel computing. Dask is tightly integrated with XGBoost, which makes it easy to train models at scale and run predictions on large datasets. With Coiled, we can easily run Dask in the cloud.

Challenges of Traditional Forecasting Methods

The standard solution to demand forecasting is to use historical sales data to predict future demand. These traditional forecasting methods based exclusively on historical sales data have several limitations.

Inaccurate forecasts lead to both excess inventory and shortages. Image from Nicolas Vandeput’s blog post.

Inability to Adapt to Real World Events

Standard time series forecasting relies on using historical data to make predictions. But history doesn’t always repeat itself.

On one hand, trends change over time as macro factors like economic conditions change. On the other hand, specific events spring up that forecasting models have no knowledge of - things like major sports games, one-off concerts or conferences, or even just school term dates.

So, to forecast accurately, models need knowledge of what’s happening in the real world.

Is it going to be unseasonably warm tomorrow? Did school start back up and kids need snacks? Did a promotion on substitute products just end? Are the 49ers playing down the road from the local store? There are hundreds of real world factors behind a shopper’s every single purchase.

Time-Consuming and Prone to Human Error

Forecasting models that are unable to account for real-world events (like holidays or promotions) require manual, ad-hoc adjustments. These adjustments not only require a lot of effort, but are also prone to human mistakes.

That’s especially true given how much interplay there is between demand drivers. Maybe this year’s 4th of July celebrations will be small because of the bad weather and local road closures.

Limited by Data Size and Complexity

Traditional forecasting methods like regression models are not designed to handle large and complex datasets and multivariate problems. As a result, they may not be able to capture all the relevant factors that influence demand, leading to less accurate predictions.

Real-World Example: Grocery Demand Forecasting

The problem of bad forecasting is especially acute in grocery. The spikiness that we see in the charts above has a large effect. It’s the difference between having empty shelves in the guacamole section on the SuperBowl, or stocking and wasting too much guacamole during the rest of the year.

Fortunately, the world is full of data, and we can use that data to get really well-informed models:

Sporting events
- Good: calendar of events, joined with proximity to local stadiums
- Better: sports betting odds and activities on each game so we know when the stadium will be empty or full
Concerts:
- Good: Find out when concerts are in town
- Better: And cross-reference Spotify data to see how popular the artists are to predict attendance
Ethnic holidays
- Good: a calendar of which days are public holidays
- Better: a calendar of holidays for every ethnic group, combined with demographic information about the surrounding neighborhood

At Guac we specialize in customizing algorithms with hyperlocal data for our different grocery customers.

Take the following example of a real 2023 sales graph from one of our customers. Running a standard historics-based approach to forecasting, the model totally misses the increase in demand driven by a big 49ers vs. Steelers game on the weekend of 10 September.

Actual (blue) and predicted (red) grocery sales before including data on sporting events.

But once we enriched the data with information about the game and its expected popularity, we can almost perfectly capture the true sales on the day.

Actual (blue) and predicted (green) grocery sales after augmenting the dataset with hyperlocal data. Accounting for the popularity of the 49ers vs. Steelers game provided much more accurate predictions.

Distributed XGBoost for Large Tabular Datasets

Unfortunately, adding all of those extra features greatly increases the size of our data. For one of our customers, this meant a 15x increase in the size of the dataset from 5 GB to 75 GB.

Setting up this infrastructure was a bit daunting. We considered using Spark and Databricks, but that would have required rewriting a huge amount of our pandas-heavy codebase. We instead latched onto Dask and Coiled.

Dask was great because …

It was almost a drop-in replacement for our existing pandas code
It had great integration with a wide range of ML libraries, such as XGBoost

Coiled was great because …

It handled all the messy dev-ops work for us
Setup was super fast and we had our Dask clusters up and running in under a day
Support has been really helpful, working with us on migrating our codebase and infrastructure

Conclusion

If this kind of rich real-world demand forecasting is something that you could use, then talk to us at Guac. We invite you to try it out with a trial on your past sales data, showcasing its ability to cut food waste.

If you’re looking for super-easy distributed XGBoost, look no further than Coiled. It’s easy to get started and free for modest use.

Schedule Python Jobs with Prefect and Coiled

2024-01-23T00:00:00+00:00

Jan 23, 2024

3m read

James Bourbeau

We show how to deploy Prefect workflows on the cloud with Coiled.

Prefect is a popular workflow management system. Its intuitive API makes writing workflows straightforward, especially for Python developers, who can usually start running on their laptop in just a few minutes.

It’s common to want to run a workflow, or specific task within a workflow, on a VM in the cloud instead of locally for different reasons:

You want an always-on machine for long-running, or regularly scheduled, jobs
You want to run close to your cloud-hosted data
You want specific hardware (e.g. GPU, lots of memory)
You want many machines for running tasks in parallel

This post shows how to deploy a Prefect workflow on the cloud using Coiled. The workflow runs a regularly scheduled job on a dataset that’s updated daily. While the real-world example we show in this post highlights a geospatial use case, the overall pattern of running regularly scheduled jobs is common in many fields (e.g. machine learning, finance, remote sensing, sports betting, etc.) and looks something like the following:

import coiled
from prefect import flow, task


# Step 1: Find new data files to be processed.
@task
def get_new_data_files():
    files = ...
    return files


# Step 2: Process input data file.
#         This often requires specific computing resources.
@task
@coiled.function(...)
def process(file):
    results = ...
    return results


# Step 3: Combine tasks into a full workflow that
#         processes all new data files.
@flow
def workflow():
    files = get_new_data_files()
    futures = process.map(files)
    for future in futures:
        print(f"Saved file {future.result()}")


# Step 4: Run the workflow on some regular cadence. 
if __name__ == "__main__":
    workflow.serve(interval=...)

Workflow

Each day NASA satellites collect sea surface temperature measurements that are then processed and stored in a new file added to the Global Sea Surface Temperature dataset on AWS S3. Our workflow consists of a daily job that finds the newly added file, does additional processing on that data, and saves the processed result to an S3 bucket.

Following the pattern above we have:

# workflow.py

import datetime
import os
import tempfile

import coiled
import earthaccess
import numpy as np
import s3fs
import xarray as xr
from prefect import flow, task

IN_BUCKET = "s3://podaac-ops-cumulus-protected/MUR-JPL-L4-GLOB-v4.1"
OUT_BUCKET = "s3://openscapes-scratch/MUR-SST"


# Step 1: Find new data files to be processed.
@task
def get_new_data_files() -> set[str]:
    # Get all files from data source bucket and output destination bucket
    fs_in = earthaccess.get_s3fs_session(provider="POCLOUD")
    infiles = set(map(os.path.basename, fs_in.glob(IN_BUCKET + "/*.nc")))
    fs_out = s3fs.S3FileSystem()
    outfiles = set(map(os.path.basename, fs_out.glob(OUT_BUCKET + "/*.nc")))
    # Find new files that haven't been processed yet
    return infiles - outfiles


# Step 2: Process input data file.
#         We run this task on a VM in the same region as the data
#         that's large enough to reliably process files.
@task
@coiled.function(
    region="us-west-2",     # Same region as data
    vm_type="m6i.large",    # Big enough VM to process data
)
def process(file: str) -> str:
    # Open file with `xarray`
    fs_in = earthaccess.get_s3fs_session(provider="POCLOUD")
    ds = xr.open_dataset(fs_in.open(IN_BUCKET + f"/{file}"))
    # Select specific region of interest + data validation
    ds = ds.sel(lon=slice(-93, -76), lat=slice(41, 49))
    cond = (ds.sea_ice_fraction < 0.15) | np.isnan(ds.sea_ice_fraction)
    result = ds.analysed_sst.where(cond)
    # Save processed result to S3
    with tempfile.TemporaryDirectory() as tmpdir:
        local_file = os.path.join(tmpdir, file)
        result.to_netcdf(local_file)
        fs_out = s3fs.S3FileSystem()
        outfile = OUT_BUCKET + f"/{file}"
        fs_out.put(local_file, outfile)

    return outfile


# Step 3: Combine tasks into a full workflow that
#         processes all new data files.
@flow
def workflow():
    files = get_new_data_files()
    futures = process.map(files)
    for future in futures:
        print(f"Saved file {future.result()}")


# Step 4: Run the workflow once a day.
if __name__ == "__main__":
    workflow.serve(name="NASA-SST", interval=datetime.timedelta(days=1))

There are a couple of key points to highlight:

We lightly annotate our process function with the @coiled.function decorator to have it run on a separate cloud VM with enough resources to process our data files. We also make sure the functions runs in the same region where our data is to avoid data transfer costs:
```
@task
@coiled.function(
    region="us-west-2",     # Same region as data
    vm_type="m6i.large",    # Big enough VM to process data
)
def process(file: str) -> str:
    ...
```
See Coiled Functions to learn more about using serverless functions.
Since our dataset is updated once a day, we set the workflow’s interval= to automatically run daily.
```
workflow.serve(name="NASA-SST", interval=datetime.timedelta(days=1))
```

Deploy on the cloud

With our workflow ready to go, we now deploy it to an always-on cloud VM using the coiled prefect serve CLI:

coiled prefect serve \                   # Run on the cloud
    --region us-west-2 \                 # Same region as data
    --vm-type t3.medium  \               # Cheap always-on VM for serving the workflow
    -e AWS_ACCESS_KEY_ID="..." \         # AWS credentials to write to output bucket
    -e AWS_SECRET_ACCESS_KEY="..." \
    -e EARTHDATA_USERNAME="..." \        # Earthdata credentials to read NASA data bucket
    -e EARTHDATA_PASSWORD="..." \
    workflow.py

Because our resource-heavy process function is dispatched to a separate VM using Coiled Functions, the rest of the workflow doesn’t need significant computing resources and can be run on a small, inexpensive t3.medium instance (costs ~$0.04/hr or less than a dollar a day). We also include two sets of secrets, one for reading data from the NASA bucket and one for writing to our processed data bucket, using the -e option for securely setting environment variables secrets.

Now once a day an m6i.large instance is spun up to run our process function on a new data file and then automatically spun down. We can see our flow running in both Coiled and Prefect:

Workflow running in Coiled UI.

Workflow running in Prefect UI.

Summary

We ran a common example of a regularly scheduled job on the cloud using Prefect and Coiled. This was:

Easy: Minimal code changes were needed to run Prefect code on the cloud.
Cheap: Costs ~$1/day to run.

We hope this post shows that with just a handful of lines of code you can run your Prefect workflows on the cloud. We also hope it provides a template you can copy and adapt for your own use case.

Want to run this example yourself?

Get started with Coiled for free at coiled.io/start. This example runs comfortably within the free tier.
Run the code snippet above.

One Billion Row Challenge (1BRC) in Python with Dask

2024-01-16T00:00:00+00:00

Jan 16, 2024

3m read

Sarah Johnson

Last week Gunnar Morling launched the One Billion Row Challenge and it’s been fun to follow along. Though the official challenge is limited to Java implementations, we were inspired by an unofficial Python submission and have our own unofficial submission for Dask.

Dask Results

We ran the challenge on an M1 Mac (8-core CPU, 16 GB memory). Dask came in at 32.8 seconds (± 2.43 seconds, averaged across 7 trials). Here’s our implementation:

from dask.distributed import Client
import dask.dataframe as dd

client = Client()

df = dd.read_csv(
    "measurements.txt",
    sep=";",
    header=None,
    names=["station", "measure"],
    engine='pyarrow',
    dtype_backend='pyarrow'
)
df = df.groupby("station").agg(["min", "max", "mean"])
df.columns = df.columns.droplevel()
df = df.sort_values("station").compute()

Dask is a general purpose Python library for parallel and distributed computing. It’s especially nice for handling large datasets that don’t fit into memory on your laptop. For tabular data, the syntax is designed to feel very similar to pandas, so it’s nearly a drop-in replacement.

Where Dask does well, though, is when handling hundreds of GB of data distributed across many machines. We were actually pleasantly surprised to see this level of performance on a single machine for only a 13 GB dataset. This is largely due to a number of recent improvements in Dask like Arrow strings, new shuffling algorithms, and query optimization (see Dask DataFrame is Fast Now for more details on these engineering efforts).

How does Dask compare?

Amongst the other 1BRC Python submissions, Dask is pretty squarely in the middle. It’s faster than using Python’s multiprocessing (except for the PyPy3 implementation) and slower than DuckDB and Polars. This is not too surprising given Polars and DuckDB tend to be faster than Dask on a smaller scale, especially on a single machine.

1BRC challenge performance for Python implementations of Polars, DuckDB, Dask, and Spark. Polars and DuckDB values are copied from this repo.

We also ran the challenge for PySpark, which took ~2 minutes. We are by no means experts in this area though, and would welcome ideas for improvement. Here’s our implementation:

from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, FloatType, StringType

spark = SparkSession.builder.master("local[*]").getOrCreate()
schema = StructType(
    [
        StructField("station", StringType(), True),
        StructField("measure", FloatType(), True),
    ]
)

df = spark.read.options(delimiter=";", header=False).csv(
    "repos/1brc/measurements.txt", schema=schema
)
df.createOrReplaceTempView("measurements")

query = """
SELECT station, MAX(measure) as max, MIN(measure) as min, AVG(measure) as mean
FROM measurements
GROUP BY station ORDER BY station;
"""

result = spark.sql(query).show()

TPC-H Benchmarks at Scale

Digging a little deeper into Dask performance, we see most computation time is spent on I/O. Yes, this challenge is for fun! But a comparison predominantly based on which tool is faster at reading a CSV into memory does have its limitations. We thought we’d mention some other more meaningful performance comparisons.

Dask dashboard during the 1BRC, most of the time is spent on read_csv tasks (purple).

We’ve run a more robust series of TPC-H Benchmarks comparing Dask, Polars, Spark, and DuckDB ranging in scale from 10GB to 10TB. We ran these benchmarks on the cloud using Coiled, which makes it easy to run your Python code in the cloud. Coiled creates ephemeral VMs in cloud account that match your local environment and then cleans everything up when you’re done.

No project wins, and it’s interesting analyzing the results.

Run the challenge yourself

First, you’ll need to generate the dataset. For those who are also less familiar with Java, the setup instructions in Robin Moffat’s post are quite helpful.

Next, you’ll need a Python environment with latest version of dask installed:

pip install dask-expr "dask[complete]"

Then you can run the challenge for Dask:

from dask.distributed import Client
import dask.dataframe as dd

client = Client()

df = dd.read_csv(
    "measurements.txt",
    sep=";",
    header=None,
    names=["station", "measure"],
    engine='pyarrow',
    dtype_backend='pyarrow'
)
df = df.groupby("station").agg(["min", "max", "mean"])
df.columns = df.columns.droplevel()
df = df.sort_values("station").compute()

You might also consider running the one trillion row challenge which scales up this dataset to one trillion rows, or ~ 2.5 TB.

Acknowledgments

Thanks to both of the earlier Python submissions from Matthew Aylward and Italo Nesi. And, of course, to Gunnar Morling for holding the challenge in the first place.

Xarray at Large Scale: A Beginner’s Guide

2023-12-21T00:00:00+00:00

Dec 21, 2023

9m read

Coiled Team

Parallel Xarray Can be Easy

You don’t need to be a distributed systems expert to process terabytes of geospatial data. The tools today are good enough to make this accessible to everyone, but most people still stay on a single computer, even though it’s often more expensive.

In fact the video below condenses 250TiB of data, and cost only $25 to produce over 20 minutes, while it would have taken thousands of dollars and weeks of download time on a single machine. Let’s walk through the common toolset used today by Python researchers to process large volumes of geospatial raster data.

What is Xarray?

Xarray is an open-source project and Python library that makes working with labeled multi-dimensional arrays simple, efficient, and fun. It’s well suited to work with raster geospatial data which are often used in climate and meteorological research, as well as other forms of environmental, oceanographic, and atmospheric data. Xarray brings the power of pandas to n-dimensional arrays, which makes it a great tool for handling complex datasets.

Key features of Xarray include:

N-Dimensional Data and Labels: Xarray manipulates multi-dimensional arrays like NumPy, but with dimension labels, making data manipulation and analysis more intuitive.
Pandas-like Operations: Xarray supports many of the same data manipulation operations as pandas, such as group-by, merge, and reshaping.
Data Readers for Common Formats like NetCDF, GeoTIFF, Zarr, HDF, and others.
Integration with Other Libraries: Xarray integrates well with other data analysis libraries like pandas, NumPy, and Matplotlib, making it a versatile tool for data science projects.
Parallel Computing: Xarray can leverage Dask to support parallel computing, allowing efficient analysis of very large datasets.

For a practical example, let’s go through reading a netCDF file and performing some simple analysis using Xarray:

import xarray as xr
ds = xr.open_dataset('example_data.nc')
mean_temperature = ds['temperature'].mean(dim='time')
mean_temperature.plot()

See the Xarray documentation for more introductory-level examples.

What is Dask?

Dask works with Xarray by providing a backend for parallel and out-of-core computation. This integration allows Xarray to handle datasets that are larger than the available memory, which is essential for processing the large, multi-dimensional arrays commonly used in scientific fields like climate science, oceanography, and geospatial analysis.

Integration Mechanism:

Chunking: When an Xarray dataset is loaded, it can be chunked into smaller parts using Dask. These chunks are essentially small, manageable blocks of the larger dataset.
Lazy Evaluation: Xarray operations on Dask-backed arrays become lazy. This means that computations are not executed immediately but are instead queued up as tasks in a Dask graph.
Task Graph: Each operation on the dataset creates a task in the Dask graph. This graph represents the sequence of operations needed to obtain the final result.
Parallel Computation: When a result is requested (e.g., for plotting, saving, or explicitly computing), Dask executes the task graph. The computations are carried out in parallel, with each chunk being processed independently. This parallel execution is key to handling large datasets efficiently.
Memory Efficiency: By working on small chunks, Dask ensures that the memory usage is kept in check, avoiding out-of-memory errors that would occur if the entire dataset were loaded into memory.

Practical Usage:

In practical terms, this means that you can work with Xarray in a way that feels familiar (similar to using in-memory NumPy arrays) but on much larger datasets. The Dask integration is transparent, so you don’t need to manage the parallelism or the chunking directly; Xarray and Dask handle these aspects behind the scenes. This makes it possible to write code that scales from small, in-memory datasets on a single machine to large datasets that are distributed across a cluster, without changing your code.

Overall, the synergy between Xarray and Dask empowers users to efficiently process and analyze large datasets that would otherwise be impractical or impossible to handle on standard hardware, making it a powerful combination for data-intensive scientific applications.

Example: How to load many NetCDF files with Xarray and Dask

To use Dask with Xarray, especially in a scenario where you want to open multiple netCDF files as a single dataset, you would use the xarray.open_mfdataset function with a chunks argument. This chunks argument specifies how the dataset should be divided into Dask chunks.

Here’s an example of how you might use xarray.open_mfdataset with Dask:

import xarray as xr

# Specify the path to your netCDF files (can use a wildcard for multiple files)
file_path = 'path/to/your/netcdf/files/*.nc'

# Open the dataset with Dask chunking
# The 'chunks' argument specifies the chunk size for each dimension
# For example, {'time': 10} would chunk the data along the time dimension in blocks of 10
ds = xr.open_mfdataset(file_path, chunks={'time': 10})

# At this point, computations on 'ds' will use Dask for parallel, out-of-core computation

This approach enables handling larger-than-memory datasets on your machine or scaling up to a cluster without changing your code significantly.

Deploy Dask and Xarray onto many machines

You can use Xarray with Dask to process large volumes of data on a single machine. This lets you process as much data as you can download and fit onto your hard drive. However, this has two limitations:

You can only fit so much data on your hard drive (around 100 GiB probably)
Downloading data can be slow and expensive

Often it is better instead to run your computations close to where your data lives. Often large datasets are stored either on HPC Network File Systems (NFS) or on cloud object stores, like AWS’s S3 or Google’s GCS. It’s best to run your computations on data where the data is stored for reasons of scale, cost, and speed.

Deploy Xarray and Dask on an HPC Cluster

It is easy to run Xarray and Dask on HPC machines using common job schedulers like SLURM, PBS, LSF, SGE, or others. Typically we recommend using the library dask-jobqueue for this.

Dask-jobqueue enables you to create and manage Dask clusters in HPC environments using common job schedulers like SLURM, PBS, SGE, or LSF. It allows dynamic scaling of Dask workers as HPC jobs, adhering to the HPC’s resource allocation and scheduling policies.

Steps to Deploy:

Install dask-jobqueue

pip install dask-jobqueue

Configure Dask Cluster

In your Python script, configure a Dask cluster that submits jobs via your HPC’s scheduler. For example, using SLURM:

from dask_jobqueue import SLURMCluster

cluster = SLURMCluster(
    queue='regular', cores=24, memory='120GB', walltime='02:00:00'
)
client = cluster.get_client()

Scale the Cluster

Adjust the number of worker jobs based on your needs.

cluster.scale(jobs=10)

Run Xarray Computations

Perform your data analysis with Xarray, leveraging the Dask cluster for parallel processing.

import xarray as xr

ds = xr.open_dataset('data.nc', chunks={'time': 100})

result = ds.mean(dim='time').compute()

This setup allows you to efficiently utilize HPC resources for processing large datasets with Xarray and Dask, scaling from a single machine to multiple nodes in an HPC environment.

Deploy Xarray and Dask on AWS

Agencies like NASA, NOAA, USGS, and others are pumping massive datasets to Amazon and making them publicly available. Researchers are using this data today to ask questions like the following:

What does air pollution look like over the 9000 most populated cities on earth?
Where should we build factories to take advantage of cheap renewable energy while bypassing local utilities?
What is the capacity of China’s fleet of solar farms, and how has this been changing over time?
How much oil has been flowing out of Venezuela on tanker ships?
How is drought likely to affect grain production in Northern Africa?

This data is plentiful and freely available, provided you can run computations close to it.

Fortunately, there are several ways to run Xarray and Dask on AWS.

Dask-Cloudprovider

Summary: Automates the setup of Dask clusters on AWS, including options like AWS Fargate for serverless deployment. See the documentation.

Code Sample:

from dask_cloudprovider.aws import FargateCluster
from dask.distributed import Client

cluster = FargateCluster(n_workers=4)
client = Client(cluster)

Dask-Kubernetes

Summary: Ideal for deploying Dask within a Kubernetes environment on AWS (Amazon EKS), allowing integration with Kubernetes resources. See the documentation.

Code Sample:

from dask_kubernetes import KubeCluster
from dask.distributed import Client

cluster = KubeCluster.from_yaml('worker-spec.yml')
client = Client(cluster)

Coiled

Summary: Provides a managed service for Dask on AWS, handling cluster provisioning and management, suitable for users preferring a simplified approach. See the documentation.

Code Sample:

import coiled
from dask.distributed import Client

cluster = coiled.Cluster(n_workers=10)
client = Client(cluster)

Each method offers a different balance of control and convenience, allowing you to choose the one that best fits your project’s requirements and your level of expertise with AWS and Dask.

We of course like Coiled the best. It’s the only option that’s secure out of the box, avoids common costly mistakes like idle clusters or lingering network resources, and it’s also the easiest to set up. It’s nice knowing though that you can always go to OSS options if you prefer.

Should I use HPC or Cloud?

Often people want to know where they should run their computations. If you’re a user the answer is simple:

Run wherever your IT admins are willing to give you resources 🙂
If you have choices, then run wherever your data is. Data gravity is a big deal.

So if you’re processing data that just came out of a simulation on a big HPC machine, and the data is sitting right there on the NFS, then use HPC. If you’re processing data that is hosted on AWS by NASA or NOAA or ESA then please run on the cloud.

If you’re a data provider and you’re planning where to host data then the answer comes down to reuse:

Is this dataset going to be generated once and then read once? If so, then host the data wherever it’s generated, often this is HPC.
Is the dataset going to be generated once but then used many times, possibly by many different groups? If so, then the cloud might lower barriers to easier sharing and collaboration. There’s no data web portal you’ll be able to build which will be nearly as interoperable or accessible as just dumping raw files into an S3 bucket with requester pays public access.

How much does the cloud cost?

This varies by region and tier, but we’ll include some rule of thumb numbers below to store 1 TB of data:

Storage: $10-20
Egress: $100
Cloud Compute: $0.10

Storing large volumes of data in the cloud can be moderately expensive, but generally these costs don’t compare to, say, the costs of an FTE. Egress costs can be very substantial if the primary way that people access your data is to download it outside of the cloud, or to different regions than where the data is stored.

Once data is stored in the cloud though, it is incredibly cheap to process, at around the cost of $0.10 per 1 TB, if you do everything right. The cost to rent an EC2 instance in the same region, where accessing the data is fast and cheap.

If done correctly, processing cloud data is incredibly cheap.

Example: Process 250 TB for $25

As an example, we’ll include below some code that we ran to produce this video:

# Start a Coiled cluster
cluster = coiled.Cluster(
    name="nwm-1979-2020",
    region="us-east-1",  # close to data
    n_workers=200
)
client = cluster.get_client()

# Access 250 TB NOAA National Water Model in the cloud with Xarray
ds = xr.open_zarr(
    fsspec.get_mapper('s3://noaa-nwm-retrospective-2-1-zarr-pds/rtout.zarr', anon=True),
    consolidated=True,
    chunks={'time': 896, 'x': 350, 'y': 350},
).zwattablrt.sel(time=slice('1979-02-01', '2020-12-31')) # water table depth

# Merge on US counties
fs = fsspec.filesystem('s3', requester_pays=True)

counties = rioxarray.open_rasterio(
    fs.open('s3://nwm-250m-us-counties/Counties_on_250m_grid.tif'), chunks='auto'
).squeeze()

_, counties_aligned = xr.align(ds, counties, join='override')

# Calculate average US county water table depth
county_id = np.unique(counties_aligned.data).compute()

county_mean = flox.xarray.xarray_reduce(
    ds,
    counties_aligned.rename('county'),
    func='mean',
    expected_groups=(county_id,),
)

county_mean.load()

This example processes 250 TiB for $25 in 20 minutes on 200 cloud machines. If you have an AWS account you should be able to run it yourself (you’ll have to set up Coiled first, which takes around 5 minutes). You can read more on technical details in our blog post.

What is Pangeo?

When Xarray and Dask started to take off in the geoscience community Xarray and Dask maintainers got together and started Pangeo, a community organization for large scale geoscience development.

Since then the Pangeo community has grown to researchers and developers at dozens of organizations working on …

OSS software: with developers collaborating on libraries like Xarray, Dask, Zarr, and other projects.
community support including frequent showcase meetings describing technical or scientific research, as well as online forums to help new researchers
Distributed infrastructure with publicly hosted cloud infrastructure

One of the strongest aspects of Pangeo today is the community. We’ll include a few links below for further study:

Summary

Easy-to-use tools like Xarray and Dask represent a significant leap in data processing, making the analysis of extensive geospatial datasets feasible and affordable. By moving computations closer to where the data is stored, whether in HPC environments or cloud platforms like AWS, researchers can now handle vast datasets efficiently. This democratizes access to high-performance computing, allowing a broader range of researchers to engage in complex data analysis tasks without the burden of prohibitive costs or the need for specialized knowledge in distributed systems. As technology continues to evolve, it’s clear that tools like Xarray and Dask will play a crucial role in unlocking new possibilities in data-driven research.

Polars queries in the cloud with Coiled

2023-11-17T00:00:00+00:00

Data	Local	Coiled	Speed Up
150 GB Parquet	30+ mins	7 mins	5x

Local machines can struggle to process large datasets due to memory and network limitations. Coiled Functions provide a cloud-based solution that allows for efficient and cost-effective handling of such extensive datasets, overcoming the constraints of local hardware for complex data processing tasks. Incorporating libraries like Polars can further enhance this approach, leveraging optimized computation capabilities to process data more quickly and efficiently.

In this post we’ll use Coiled Functions to process the 150 GB Uber-Lyft dataset on a single cloud machine with Polars.

Query S3 data locally with Polars

The process starts with creating queries as though we would for local execution against the dataset.

def load_data():
    # Define the S3 path and storage options and load the data into memory
    s3_path = 's3://coiled-datasets/uber-lyft-tlc/*'
    storage_options = {"aws_region": "us-east-2"}
    return pl.scan_parquet(s3_path, storage_options=storage_options)

def compute_percentage_of_tipped_rides(lazy_df):
    # Compute the percentage of rides with a tip for each license number
    result = (lazy_df
              .with_columns((pl.col("tips") > 0).alias("tipped"))
              .group_by("hvfhs_license_num")
              .agg([
                  (pl.sum("tipped") / pl.count("tipped")).alias("percentage_tipped")
              ]))
    return result.collect()

def query_results():
    # Load and process the dataset to compute the results
    lazy_df = load_data()
    return compute_percentage_of_tipped_rides(lazy_df)

This query is pretty basic; it is meant as a simple example of how we can process these large files. If we execute this query locally, we would pull all of the data onto the machine, process it, and return a small fraction of the data. Unfortunately for me, at least, this would hit the physical memory limits of my machine. It would also take a long time, as I would have to download the whole dataset and be limited by my network bandwidth. In other situations this would also be an expensive operation as I would have to pay egress costs on the data. Let’s see how Coiled can help solve these issues.

Query S3 data in the cloud with Coiled serverless functions

Coiled Functions allow us to use cloud machines that are large enough to process this dataset. Coiled connects to AWS or GCP, so we can access any virtual machine (VM) that is available there. We can also co-locate the VM with the data, speeding up processing and avoiding egress costs.

All we need to do is add a @coiled.function decorator to our main function. The decorator tells Coiled that it should spin up a large VM on AWS, run the query there, and then return the result locally.

@coiled.function(
    vm_type="m6i.4xlarge",  # VM with 64GB RAM
    region="us-east-2",     # AWS region to match the data location
    keepalive="5 minutes",  # VM keepalive time for potential subsequent queries
)
def query_results():
    # Load and process the dataset to compute the desired metric
    lazy_df = load_data()
    return compute_percentage_of_tipped_rides(lazy_df)

coiled.function() takes a variety of arguments. Let’s take a look at a few of the important ones:

vm_type: This specifies the type of VM that we want the computation to run on.
region: This specifies the AWS or GCP region that we want to start the VM in.
keepalive: This keeps the VM alive so that we can run multiple queries against the data in memory.

Let’s execute our queries and pull the results back to our local machine:

result = query_results()
print(result)

[
  ('HV0002', 0.08440046492889111),
  ('HV0003', 0.1498555901186066),
  ('HV0004', 0.09294857737045926),
  ('HV0005', 0.1912300216459857),
]

That’s all you need to do to run this function in the cloud. There is no need to adjust the other functions or manually spin up a VM.

What just happened

coiled.function() first scanned my local environment for the installed packages and replicated those on the VM that I requested. We didn’t need to specify a Docker container or a specific Python environment (though we could). Coiled then started a VM in AWS with the specified EC2 instance. The VM took about two minutes to start up, and then downloading the dataset and running the computations in the cloud took another five minutes. Not too bad for processing 150 GB! Finally, Coiled returned the results back to our local machine.

If we look at the Coiled Dashboard, we can see that all the data was pulled into the VM in the cloud, processed, and returned.

VM Memory Utilization Increasing as the Computation Runs

Coiled would normally shut down the VM immediately after the Python interpreter shuts down in order to reduce costs. We specified keepalive="5 minutes" to keep the VM running in the cloud so that we could use the same VM to do other queries. This avoids the normal two minute boot time; we call this “warm start.”

Conclusion

coiled functions and Polars provide a solution for processing large datasets in the cloud. This approach allows data practitioners to overcome local machine limitations and reduce operational costs. This example demonstrates how simple it is to execute large-scale queries using Coiled. By deploying a @coiled.function decorator, we can swiftly transition from local to cloud processing, gaining the benefits of powerful cloud computing resources. The outcome is a faster, more scalable data processing workflow that enhances productivity and focuses on data analysis rather than infrastructure challenges.

Want to run this example yourself?

Get started with Coiled for free at coiled.io/start. This example runs comfortably within the free tier.
Copy and paste the first code snippet and add the coiled.function() decorator.

You can check out the docs or take a look at how to utilize Coiled Functions to run the same query using DuckDB.

Processing Terabyte-Scale NASA Cloud Datasets with Coiled

2023-11-01T00:00:00+00:00

We show how to run existing NASA data workflows on the cloud, in parallel, with minimal code changes using Coiled. We also discuss cost optimization.

Data	Local	Coiled	Speed Up	Cost
500 GB NetCDF	383 mins	9 mins	42x	$0.36

People often run the same function over many files in cloud storage. This looks something like the following:

# Step 1: Get a list of all of our files (typically thousands of NetCDF files).
files = ...

# Step 2: Create a function to process each file.
#         Then either return the processed result
#         or save it to some external data store like S3.
def process(file):
    result = ...
    return result

# Step 3: Process all of the files.
results = []
for file in files:
    result = process(file)
    results.append(result)

# Step 4: Analyze results locally.
plot(results)

This straightforward but important pattern occurs frequently in groups working with NASA Earthdata and NASA Distributed Active Archive Centers (DAACs) to move scientific workloads to the cloud. Here are some examples:

Extract greenhouse gasses over urban areas from TROPOMI L2 data: TROPOMI measures trace gasses and aerosols relevant for air quality and climate. It generates a large volume of L2 data stored as millions of files on S3. In our particular project we’re interested in greenhouses over highly populated urban areas, and how they affect quality of life. We need to select those files over many cities and then construct a timeseries over them to see how air quality changes.
Grid sea ice freeboard data from NASA’s ICESat-2 mission: This workflow from NSIDC researchers loads L3 sea ice freeboard HDF5 data, selects a lat-lon bounding box of interest, and regrids the dataset to the EASE-Grid v2 6.25 km projected grid using drop-in-the-bucket resampling. The resulting processed dataset is then written out to a local Parquet file to be used in further analysis. This is particularly valuable as it allows researchers to easily create their own specialized datasets with spatial and temporal resolutions relevant for their analysis.

In this post we’ll show how to run this same “function on many files” analysis pattern on the cloud with Coiled, which you can copy and modify for your own use case. We’ll also highlight cost optimization strategies for processing terabyte-scale data for fractions of a dollar.

Comparing cost and duration between running the same workflow locally on a laptop (data egress costs), running on AWS, and running with cost optimizations on AWS.

Processing Locally

In this example we process many NetCDF files stored on S3 from the MUR Global Sea Surface Temperature NASA dataset to look at surface temperature variation over the US Great Lakes region:

import os
import tempfile
import earthaccess
import numpy as np
import xarray as xr

# Step 1: Get a list of all files.
#         Use earthaccess to authenticate and find data files (total of 500 GB).
granules = earthaccess.search_data(
    short_name="MUR-JPL-L4-GLOB-v4.1",
    temporal=("2020-01-01", "2021-12-31"),
)


# Step 2: Create a function to process each file.
#         Load and subset each data granule / file.
def process(granule):
    results = []
    with tempfile.TemporaryDirectory() as tmpdir:
        files = earthaccess.download(granule, tmpdir)
        for file in files:
            ds = xr.open_dataset(os.path.join(tmpdir, file))
            ds = ds.sel(lon=slice(-93, -76), lat=slice(41, 49))
            cond = (ds.sea_ice_fraction < 0.15) | np.isnan(ds.sea_ice_fraction)
            result = ds.analysed_sst.where(cond)
            results.append(result)
    return xr.concat(results, dim="time")


# Step 3: Run function on each file in a loop
results = []
for granule in granules:
    result = process(granule)
    results.append(result)


# Step 4: Combine and plot results
ds = xr.concat(results, dim="time")
ds.std("time").plot(figsize=(14, 6), x="lon", y="lat")

This specific example uses earthaccess to authenticate with NASA Earthdata and download our dataset (~500 GB of NetCDF files) and Xarray to load and select the region of the data we’re interested in, but we’ve seen many different approaches here, including file formats like HDF, Zarr, and GeoTIFF, as well as libraries like Xarray, rasterio, and pandas. In each case though the underlying pattern is the same.

We then combine the results from processing each data file and plot the results.

Standard deviation of sea surface temperature over the Great Lakes region in the US.

Running this locally on my laptop takes ~6.4 hours and costs NASA ~$25 in data egress costs (based on a $0.05/GB egress charge rate – normal is $0.10/GB, but presumably NASA gets a good deal).

Processing on the Cloud with Coiled

Let’s run the exact same analysis on the cloud. Running on the cloud helps in a couple of ways:

Data Proximate Computing: Running computations closer to where the data is stored increases performance and avoids data transfer costs.
Scale: Distributing the processing over many machines in parallel lets us tackle larger volumes.

Using Coiled to run the same workflow on the cloud involves lightly annotating our process function with the @coiled.function decorator:

import coiled

@coiled.function(
    region="us-west-2",                  # Run in the same region as data
    environ=earthaccess.auth_environ(),  # Forward Earthdata auth to cloud VMs
)
def process(granule):
    # Keep everything inside the function the same
    ...

This tells Coiled to run our processing function on a VM in AWS and then return the result back to our local machine.

Additionally, we switch to using the Coiled Function .map() method (similar to Python’s builtin map function) to run our process function across all the input files in parallel:

# Step 3: Run function on each file in a loop
results = []
for granule in granules:
    result = process(granule)
    results.append(result)
results = process.map(granules)     # This runs on the cloud in parallel

With these two small code changes, the same data processing runs in parallel across many VMs on AWS. The runtime drops to ~9 minutes (a factor of ~42x speedup) and costs ~$1.52 (a factor ~16x less).

Optimizing Costs

Processing 500 GB of data costs us $1.52. This is cheap.

We bring this cost down further to $0.36 with the following techniques:

Spot: Use spot instances, which are excess capacity offered at a discount.
ARM: Use ARM-based instances, which are less expensive than their Intel-based equivalents, often with similar, or better, performance.
Instance Type Selection: Use single-core instances (only available when using ARM), because this particular computation doesn’t require large computing resources.

With this additional configuration made to the Coiled Functions decorator, our full cost-optimized cloud workflow looks like the following:

import os
import tempfile
import coiled
import earthaccess
import numpy as np
import xarray as xr

# Step 1: Get a list of all files.
#         Use earthaccess to authenticate and find data files (total of 500 GB).
granules = earthaccess.search_data(
    short_name="MUR-JPL-L4-GLOB-v4.1",
    temporal=("2020-01-01", "2021-12-31"),
)


# Step 2: Create a function to process each file.
#         Load and subset each data granule / file.
@coiled.function(
    region="us-west-2",                  # Run in the same region as data
    environ=earthaccess.auth_environ(),  # Forward Earthdata auth to cloud VMs
    spot_policy="spot_with_fallback",    # Use spot instances when available
    arm=True,                            # Use ARM-based instances
    cpu=1,                               # Use Single-core instances
)
def process(granule):
    results = []
    with tempfile.TemporaryDirectory() as tmpdir:
        files = earthaccess.download(granule, tmpdir)
        for file in files:
            ds = xr.open_dataset(os.path.join(tmpdir, file))
            ds = ds.sel(lon=slice(-93, -76), lat=slice(41, 49))
            cond = (ds.sea_ice_fraction < 0.15) | np.isnan(ds.sea_ice_fraction)
            result = ds.analysed_sst.where(cond)
            results.append(result)
    return xr.concat(results, dim="time")


# Step 3: Run function on each file in parallel
results = process.map(granules)     # This runs on the cloud in parallel


# Step 4: Combine and plot results
ds = xr.concat(results, dim="time")
ds.std("time").plot(figsize=(14, 6), x="lon", y="lat")

and, while code runtime stays the same at ~9 minutes, cost drops to ~$0.36 (a factor of ~70x less than running locally).

Summary

We ran a common data subsetting-style workflow on NASA data in the cloud with Coiled. This demonstrated the following:

Easy: Minimal code changes to migrate existing code to run on the cloud.
Fast: Acceleration of the workflow runtime by ~42x.
Cheap: Cost reduction of ~70x.

We hope this post shows that with just a handful of lines of code you can migrate your data workflow to the cloud in a straightforward manner. We also hope it provides a template you can copy and adapt for your own analysis use case.

Comparing cost and duration between running the same workflow locally on a laptop (data egress costs), running on AWS, and running with cost optimizations on AWS.

Want to run this example yourself?

Get started with Coiled for free at coiled.io/start. This example runs comfortably within the free tier.
Copy and paste the code snippet above.

Acknowledgements

This work was done in collaboration with NASA Openscapes, a community supporting research teams’ cloud migration with different Cloud environments and technologies, including Coiled.

We’d also like to thank Amy Steiker, Luis Lopez, and Andy Barrett from NSIDC and Aronne Merrelli from the University of Michigan for sharing their scientific use cases which motivated this post.

Run Jupyter Notebooks on a GPU on the Cloud

2023-10-10T00:00:00+00:00

Oct 10, 2023

3m read

Sarah Johnson

GPUs accelerate tasks like ML training, computer vision, and analytics, but can be challenging to access. Using the cloud can open up access to GPUs, but comes with its own challenges including:

Complex Cloud Infrastructure: AWS SageMaker, for example, requires navigating AWS IAM roles, VPC settings, and other cloud configuration steps before even starting a GPU instance.
Limited Customization: Tools like SageMaker and Google Colab provide managed environments, but they can feel restrictive, making it hard to install custom dependencies or optimize GPU usage.
Poor Developer Experience: Compared to tools like Jupyter Notebooks, the UI for SageMaker or Colab can feel clunky and unintuitive, adding friction to the development process.

Coiled can launch a GPU-enabled Jupyter Notebook for you, from your own AWS, GCP, or Azure account, in a couple minutes.

$ coiled notebook start --gpu --sync

This has a few benefits:

GPU access: You’ll get temporary access to a GPU machine for as long as the notebook is running. When you’re done, it’ll turn off to minimize costs.
Software synchronization: You’ll get all of your current Python packages, including the same version of libraries like PyTorch, but those libraries will be properly configured (with CUDA drivers and everything) to operate on the underlying GPU. You don’t need to mess around with Docker.
File synchronization: You’ll also get all of your local files, including your notebooks, live-synced up to the cloud machine running Jupyter. Additionally, any changes you make in the notebook on the cloud machine will be synchronized back down to your computer, saving all of your changes on your local hard drive.

This creates a seamless experience for interactive work.

Getting Started

If you haven’t already signed up for Coiled, it’s easy to get started. You’ll get 500 free CPU hours each month, though you’ll still have to pay your cloud provider.

Next, install Coiled alongside additional notebook dependencies like jupyterlab and your machine learning library of choice.

conda

$ conda create -n env python=3.11 pytorch torchvision coiled jupyterlab jupyter-server-proxy
$ conda activate env

pip

$ pip install "coiled[notebook]" torch torchvision

Then start a JupyterLab instance on a GPU-enabled VM on the cloud:

coiled notebook start --gpu

Under the hood, Coiled handles:

Provisioning a cloud VM with GPU hardware. In this case, a g4dn.xlarge (Tesla T4) instance on AWS.
Setting up the appropriate NVIDIA drivers, CUDA runtime, etc.
Automatically installing the same packages you have locally on the cloud VM.
Making sure to install the CUDA-compiled version of PyTorch on the cloud VM.

You can also include additional arguments like:

--sync to live-sync your local files, including your Jupyter Notebooks, to the cloud machine. Any changes you make in the notebook running remotely will be synchronized back down to your computer. You’ll need Mutagen installed.
--vm-type to request a specific VM type.
--container to specify a Docker container.
--region to specify a region. We find GPUs are often easier to get in AWS region us-west-2.

See our documentation for more details.

Example: Define a PyTorch Model

Now that we have a notebook running, we can define a PyTorch neural network and use it for model inference. We’ve adapted this example from this introductory PyTorch tutorial.

import os
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# Select what hardware to use
if torch.cuda.is_available():
    device = torch.device("cuda")        # NVIDIA GPU
elif torch.backends.mps.is_available():
    device = torch.device("mps")         # Apple silicon GPU
else:
    device = torch.device("cpu")         # CPU
print(f"{device = }")

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

We can then use the model to generate predictions based on the input data.

model = NeuralNetwork().to(device)

X = torch.rand(1, 28, 28, device=device)
logits = model(X)
pred_probab = nn.Softmax(dim=1)(logits)
y_pred = pred_probab.argmax(1)

Next steps

This example was primarily illustrative, to show how you can get up and running with Jupyter Notebooks on GPUs. For more extensive machine learning examples you might consider:

Ten Cents Per Terabyte

2023-10-06T00:00:00+00:00

Oct 6, 2023

2m read

Matthew Rocklin

The optimal cost of cloud computing

I often get asked the question:

Q: What should the following workflow cost in the cloud?

If the computation is some sophisticated ML workflow then the answer is who knows 🤷? We’ll have to run it and see. But most workloads are simple and so, when tuned, are primarily bound by S3 bandwidth. In this case the answer is pretty easy.

A: If you do everything right, the cost is ten cents per terabyte.

Let’s do a back-of-the-envelope calculation:

~60 MB/s S3 bandwidth for a single core VM
~~$0.05 per hour~~
~~$0.04 per hour if you use ARM~~
$0.02 per hour if you use ARM + Spot (m7g.medium)

So 1 TB / (60 MB/s) / 3600 s/hr * $0.02 / hr = $0.10

This assumes linear scaling, which for these workloads is usually a good assumption.

But only if you do everything right

We recently wrote about scaling an Xarray computation out to 250 TB. After optimization it cost $25, so right on target.

However, before optimization that workload didn’t run at all. After some work it ran, but was slow and expensive. It was only after more tuning that the workload ran at the magic ten-cents-per-terabyte number. Tuning is hard.

We see this often. People come with wildly unoptimized cloud workflows running on EC2 or AWS Lambda or Databricks or something, asking if we can make them cheaper. We’re often able to get these computations down to a few cents (most problems are small).

Common problems include:

Downloading data locally rather than operating near the data (Egress costs can be $0.10 per GB rather than TB)
Really inefficient pure-Python code
Many small reads to S3, especially with non-cloud-optimized file formats or bad partitioning
Provisioning big machines for small tasks
Provisioning small machines for big tasks and swapping to disk
Idling resources (the harder a system is to turn on, the more people avoid turning it off)
Not using ARM (recompiling software is hard)
Not using Spot (failing machines are annoying)

But of course, the biggest cost here is human time. People spend days (weeks?) wrangling unfamiliar cloud services (EC2, Batch, Lambda, SageMaker) unsuccessfully, at the personnel cost of thousands of dollars (people are expensive).

Optimize for Ease, Visibility, and Flexibility

I like compute platforms that optimize for …

Ease of use and rapid iteration
Visibility of performance metrics
Flexibility of hardware

This combination enables fast iteration cycles that enable humans to explore the full space of solutions (code + hardware) effectively.

Putting on my sales hat, that’s roughly what we’ve built at Coiled for Python code at scale. It’s been a joy to use.

TPC-H Benchmarks for Query Optimization with Dask Expressions

2023-10-05T00:00:00+00:00

Dask-expr is an ongoing effort to add a logical query optimization layer to Dask DataFrames. We now have the first benchmark results to share that were run against the current DataFrame implementation.

Dask-expr is up to 3 times faster and more memory efficient in its current state than the status quo. This is a very promising result for our future development efforts.

Results

We are comparing Dask 2023.09.02 with the main branch of Dask-expr. Both implementations will use the P2P shuffling algorithm. The results were produced on 100GB of data, e.g. scale 100 for the TPC-H benchmarks. The cluster had 20 workers with 16GB each.

(Runtime on a per-query basis of Dask DataFrame and Dask-expr)

We can see that Dask-expr performs better on every single query, up to a 3-times improvement compared to the status quo.

Dask-expr reorders the query that was given by the user to only execute things that are strictly necessary. It will, for example, drop unnecessary columns and thus reduce the amount of data that is required. We can see that the biggest improvements are for the queries that contain merge operations. These queries write a lot of data to disk and send it over the network. Reducing the size of the data that is transferred has a big impact on the performance.

The reordering results in a reduced memory footprint as well. This is especially helpful for users that struggle with memory pressure on their clusters.

(Memory usage on a per query basis of Dask DataFrame and Dask-expr)

We did another run on 1TB of data, which showed similar results but with a few notable differences:

Getting Dask-expr to compute the results successfully was very easy. 20 workers very sufficient to compute the results, which is the same number of workers as we used for the 100GB benchmarks.
The original version needed a bigger cluster with 50 workers and with 32GB memory each.
Some queries didn’t complete the original version at all, because there were memory spikes on single clusters. Dask-expr was able to complete all queries successfully.

Summarizing, getting these queries to complete was significantly easier with Dask-expr.

These results show us that we are on the right course and motivates us to improve the performance of Dask-expr further. There is still a lot of untapped potential.

About the benchmarks

The TPC-H benchmarks are a set of queries that are commonly used to compare performance of different databases. Thus, they are very merge and groupby heavy. Historically, the current implementation didn’t perform very well on these types of queries. The new P2P shuffling algorithm that was introduced earlier this year and now Dask-expr improves the performance of Dask a lot.

These queries only capture a small part of the Dask DataFrame API but are helpful to compare performance.

The full implementation of all queries is available here. A single example query is copied below, so you can get a sense of what they look like:

var1 = datetime.strptime("1995-03-15", "%Y-%m-%d")
var2 = "BUILDING"

lineitem_ds = read_data("lineitem")
orders_ds = read_data("orders")
cutomer_ds = read_data("customer")

flineitem = lineitem_ds[lineitem_ds.l_shipdate > var1]
forders = orders_ds[orders_ds.o_orderdate < var1]
fcustomer = cutomer_ds[cutomer_ds.c_mktsegment == var2]

jn1 = fcustomer.merge(forders, left_on="c_custkey", right_on="o_custkey")
jn2 = jn1.merge(flineitem, left_on="o_orderkey", right_on="l_orderkey")
jn2["revenue"] = jn2.l_extendedprice * (1 - jn2.l_discount)
total = jn2.groupby(["l_orderkey", "o_orderdate", "o_shippriority"])[
    "revenue"
].sum()
total.reset_index().compute().sort_values(["revenue"], ascending=False).head(10)[
    ["l_orderkey", "revenue", "o_orderdate", "o_shippriority"]
]

What’s next

We have some additional steps planned over the next couple of weeks. These include:

Benchmarking Dask expressions against Spark
Scaling the benchmarks to 1TB, 10TB and 100TB.

We are hoping that we can get additional insights into our implementation and are hoping to identify bottlenecks.

Can I use Dask-expr?

Though Dask-expr is still under active development, the project is in a state where users can try it out. The API coverage is pretty good already. We are still adding new optimizations, so performance will improve in the future.

You can try it out with:

pip install dask-expr

We don’t expect that there exists any hard-to-find-bugs. We did not reimplement any of the algorithms from Dask DataFrame, so we don’t expect hard to identify issues or an unexpected loss of data.

Conclusion

Dask Expressions provides a faster and more memory efficient implementation of Dask DataFrames. The project is still under active development, but it already outperforms the current status quo by a significant margin on the TPC-H benchmarks.

Coiled observability wins: Chunksize

2023-09-19T00:00:00+00:00

Distributed computing is hard, distributed debugging is even harder. Dask tries to simplify this process as much as possible. Coiled adds additional observability features for your Dask clusters and processes them to help users understand their workflows better.

This series of posts will illustrate how we are building metrics that help us detect clusters that are underutilized or run into errors. We are using this information to help users with optimizing their workflows.

This post will illustrate how small chunks of data can slow your workflow down and lead to underutilizing your resources.

Coiled collect all metrics that we describe in these posts. We only use them to monitor the health of Dask clusters and to help users improve their Dask workflows. Coiled users can turn the metric collection off at any time without compromising their workflows.

Using bigger chunks to avoid bottlenecks on the scheduler

Dask processes your data in chunks. Every chunk creates one task per operation on your data (we will ignore fusion of tasks for simplicity). The processing time per task grows with the size of your chunks. Finding the appropriate size for these chunks is not always trivial.

A common guideline for tasks within a Dask workflow is that tasks should run for at least 100ms. The scheduler has some overhead while figuring out how to schedule your workflow efficiently. Having tasks that run significantly faster would strain your scheduler.

Dask tries to optimize the default chunksize, but this might conflict with other parameters. One common example is many small input files. Dask generally parallelizes over your files to speed up the I/O step. This will result in small chunks in this case.

Let’s look at some of our observations from clusters that had tasks with a processing time significantly below that. We will create an example with synthetic data here.

import dask.array as da

arr = da.random.randint(1, 1_000_000, (3500, 3500, 3500), chunks=(50, 50, 50))
arr.mean().compute()

This example creates a random Dask array with tiny chunks. Processing one of these chunks takes only 2ms, which is very short for a Dask task. Let’s look at how our task stream looks while processing this query:

This is a snapshot from the Dask Dashboard. We can see that the task stream has many small dots and a lot of white in between. The white spots represent periods where your workers are idle because they don’t have anything to process. This is generally a sign that your resources aren’t used efficiently. This workflow is clearly constrained by the scheduler.

Automatically detecting these patterns with Coiled

Coiled allows us to monitor clusters of our customers and to detect some of these patterns. This is a tool we use to help customers to run their code faster and use their resources more efficiently.

We have two metrics that allow us to spot these types of clusters automatically. The first collects the number of tasks on the scheduler that are queued. This metric notifies us when the number of tasks reaches a certain threshold.

That is not a foolproof system, large task graphs can have a lot of tasks while using an efficient chunksize. We have another metric that checks CPU utilization on the scheduler periodically and flags clusters that have long periods of time when the scheduler uses close to 100% CPU. This is a sign that the scheduler constrains the workload on your cluster.

Again, this is not foolproof. There could be a number of reasons why scheduler utilization is so high. Using both metrics together enables us to flag these clusters with a high degree of certainty.

We will then reach out to our users if we detect this pattern on one of their clusters. One of our engineers will work with them to make their workload more efficient and thus run their workflow more efficiently.

In this scenario, we would recommend increasing the chunksize of the individual chunks.

import dask.array as da

arr = da.random.randint(1, 1_000_000, (3500, 3500, 3500), chunks=(400, 400, 400))
arr.mean().compute()

This now runs in a couple seconds, rather than minutes.

The scheduler is also pretty bored while processing our query. That’s exactly what we were aiming for.

Conclusion

Coiled offers us a lot of insights into Dask clusters and their metrics. We can use this to identify clusters that aren’t used as efficiently as possible. Having relatively small chunksizes can lead to underutilizing your resources and thus spend more time than necessary on a give computation.

Monitoring the number of queued tasks and scheduler CPU allows us to detect such clusters automatically and help users adjust their code to make better use of their hardware.

Parallel Serverless Functions at Scale

2023-09-07T00:00:00+00:00

Data	Local	Coiled	Speed Up
500 GB NetCDF	305 mins	4.7 mins	65x

The cloud offers amazing scale, but it can be difficult for Python data developers to use. This post walks through how to use Coiled Functions to run your existing code in parallel on the cloud with minimal code changes.

Example: Processing Many Parquet Files on S3

In the code snippet below we load NYC Taxi Data from AWS (data lives in region us-east-1), process that data, and then upload the processed data back to AWS. This type of data processing / cleaning is fairly common in data engineering / science workloads.

import coiled
import pandas as pd
import s3fs
from dask.distributed import print

# Collect files to process
s3 = s3fs.S3FileSystem()
filenames = s3.ls("s3://nyc-tlc/trip data/")  # Data is in us-east-1
filenames = [
    f"s3://{fn}" for fn in filenames
    if "yellow_tripdata_201" in fn   # Data files from the 2010s
    or "yellow_tripdata_202" in fn   # Data files from the 2020s
]

# Define our processing function
def process(filename):
    print(f"Processing {filename}")
    df = pd.read_parquet(filename)
    df = df[df.tip_amount != 0]

    outfile = "s3://oss-shared-scratch/jrbourbeau/" + filename.split("/")[-1]
    df.to_parquet(outfile)
    return outfile

# Process files
print(f"Processing {len(filenames)} files")
results = map(process, filenames)
for outfile in results:
    print(f"Finished {outfile}")

print("Done!")

Running this code on my laptop (M1 MacBook Pro) takes a little over five hours.

This is slow for a couple of reasons:

We spend time transferring data from AWS to my laptop and back.
Data files process sequentially, one at a time.

Let’s use Coiled Functions to address both these issues.

Run Functions on the Cloud

To run the process function on the cloud, we lightly annotate it with the @coiled.function decorator. Now process automatically runs on a VM in AWS and returns the result locally.

@coiled.function(
    region="us-east-1",    # Same region as data
    vm_type="m6i.xlarge",  # Enough memory to process files
)
def process(filename):
    # Everything inside the function stays the same
    ...

Note that we set region="us-east-1" and vm_type="m6i.xlarge" to make sure the VM the process function runs on is in the same region as our data and has enough memory to process the data files (m6i.xlarge instances have 16 GiB of memory).

Adding the @coiled.function decorator reduces the overall runtime to 35.9 minutes (8.5x faster than the original version that ran locally). This is a nice performance improvement, but we can do even better by scaling out and running in parallel.

Run in Parallel

By default @coiled.function decorated functions will run in the cloud, serially, and return the results back to your local machine. When we want to apply our function many times across different inputs, we can use the Coiled Funtion .map method to run our function in parallel across inputs on multiple VMs.

In this example, we swap out Python’s builtin map function for, process.map (everything else stays the same):

# results = map(process, filenames)  # Replace this
results = process.map(filenames)     # With this  

Now the number of VMs used to run process calls will adaptively scale up and down depending on the workload (see Adaptive Deployments in the Dask docs for more information), with each VM running a separate process call across inputs.

Running the process function in parallel, we briefly scale up to 33 VMs, and the overall runtime dropped to 4.7 minutes (65x faster than the original version that ran locally).

VMs adaptively scaling up to process files quickly in parallel and then scaling down after all files are processed

Summary

In this post we demonstrated how to take existing Python code that runs locally and run it in parallel on the cloud with Coiled Functions. With minimal code changes, we were able to improve performance by a factor of 65x, without needing to maintain any additional cloud infrastructure. This makes developers’ lives a little easier and lowers the bar for transitioning code to run on the cloud.

Comparing code runtime between a laptop, single cloud VM, and multiple cloud VMs in parallel

Want to run this example yourself?

Get started with Coiled for free at coiled.io/start. This example runs comfortably within the free tier.
Copy and paste the code snippet above.

Processing a 250 TB dataset with Coiled, Dask, and Xarray

2023-09-05T00:00:00+00:00

We processed 250TB of geospatial cloud data in twenty minutes on the cloud with Xarray, Dask, and Coiled. We do this to demonstrate scale and to think about costs.

Data	Local	Coiled	Speed Up	Cost
252 TB Zarr	500h+[1]	6 mins	5,000x	$25

Many people use Dask with Xarray to process large datasets for geospatial analyses, however, there is pain when operating at a large scale. This comes up in a few ways:

Large task graphs due to too-small chunks causing the scheduler to become a bottleneck
Running out of memory due to one of:
- Dask being too eager to load in data
- Rechunking
- Unmanaged memory
- Too-large chunks

Many of these problems are now fixed (e.g. too-eager data loading), but not all. We wanted to run a big-but-simple Xarray problem to experience this pain ourselves and demonstrate feasibility. We’re also interested in optimizing costs, so we decided to play around with that too.

In this example we’ll process 250 TB of data from the National Water Model, calculating the ground saturation level in every US county. The calculation will take ~20 minutes and cost ~$25.

National Water Model Dataset

The National Water Model (NWM) is a complex hydrological modeling framework that simulates observed and forecasted streamflow across the US at a fine spatiotemporal scale. Many organizations use the NWM to make decisions around water management and to predict droughts and flooding. For more background, read America Is Draining Its Groundwater Like There’s No Tomorrow from The New York Times, published just last week.

In this example, we use a dataset from a 42-year retrospective simulation available on the AWS Open Data Registry. These historical data models provide context for real-time estimations and as a historical baseline for training predictive models.

We calculate the mean depth to soil saturation for each U.S. county:

Years: 1979-2020
Temporal resolution: every 3 hours
Spatial resolution: 250m grid
Data size: 252 terabytes

Processing one year of data

We first run this computation for one year (drawn from this example by Deepak Cherian). We process 6 TB of data, which takes about 8 minutes and costs $0.77 on a Coiled cluster with 40 workers, each with 16 GiB of memory. You can run this example yourself in our documentation. This analysis produces the plot below:

Since this is only a subset of the full dataset, we thought it’d be fun to run this for the full 42 years rather than just 1.

Scaling up to the full dataset

We now run this computation on all 42 years of the 250 TB dataset. This takes 24 minutes and costs $25 using 200 workers each with 64 GiB of memory. You can see the example here.

This presents a few challenges at first:

Large task graph (~344 MB)
Running out of memory (and spilling to disk)
Unnecessary rechunking

Median CPU usage is low with a large range. Memory usage is high and consistently spilling to disk, which significantly slows runtime.

So we make a few adjustments.

Larger chunk sizes

By default each chunk size is 200 MiB, resulting in a Dask graph of 1.3 million chunks. By increasing the chunksize to 840 MiB, we reduce the Dask graph to 320k chunks:

ds = xr.open_zarr(
    fsspec.get_mapper("s3://noaa-nwm-retrospective-2-1-zarr-pds/rtout.zarr", anon=True),
    consolidated=True,
    chunks={"time": 896, "x": 350, "y": 350}
)

Memory-optimized VMs

We further improve memory usage with memory-optimized instance types with a higher memory to CPU ratio. Since peak RAM was ~20 GB, we use r7g.2xlarge:

cluster = coiled.Cluster(
    ...,
    scheduler_vm_types="r7g.xlarge",
    worker_vm_types="r7g.2xlarge",
)

Launch from a remote VM

Lastly, we use Coiled Run to submit this workflow from a VM in the cloud to speed up uploading the task graph to the scheduler.

coiled run --region us-east-1 --vm-type m6g.xlarge python xarray-water-model.py

Dask uploads the task graph from the client to the scheduler, and if this were running locally we’d be relying on home internet speeds (which aren’t bad for me in Seattle, but this step helps speed things up a bit).

Results

This improves runtime and CPU- and memory-utilization quite a bit.

Median CPU usage is much higher at ~55%. Memory usage is consistently between 5 and 20 GB and we’re no longer spilling to disk.

Optimizing costs

In addition to trying to improve performance, we also want to optimize costs. There are a few ways we do this:

Use spot instances, which are excess capacity offered at a discount.
Use ARM-based instances, which are less expensive than their Intel-based equivalents, often with similar, or better, performance.
Use a region close to our data to avoid egress charges.

cluster = coiled.Cluster(
    region="us-east-1", # choose region close to data
    worker_vm_types="r7g.2xlarge", # ARM
    compute_purchase_option="spot_with_fallback" # spot instances, replaced with on-demand
)

We save $25 by using r6g ARM-based instance types rather than the x86-based equivalent. This is in part because runtime is nearly 10 minutes faster for the ARM-based instance types (others have similar findings).

Comparing r6i x86-based vs. r6g and r7g ARM-based EC2 instance types.

We also use more spot instances, since r6g instances have a much lower frequency of spot interruption than r6i (<5% vs. >20%, respectively). Switching to the newer r7g ARM Graviton3 instances saves us an additional $9, mostly due to a faster runtime of ~6 minutes.

Conclusion

We ran a large Xarray workflow with three goals in mind:

Demonstrate feasibility
Expose rough edges at scale
Explore cost dimension of these problems.

Mission accomplished. We processed a quarter-petabyte for $25.

We also found opportunities for optimization like large graph sizes, chunksize tuning, and hardware choices. We hope this article helps point others in this situation in good directions. It’s also helping to inform the internal Dask team’s work prioritization for the coming couple of months.

In the future we’d like to run similar calculations, but with a more complex workflow (like rechunking). We hope that that exposes even more opportunities for optimization.

Want to try this experiment yourself?

Download the notebook to run one year or the script to run all 40 years.
Follow the Coiled setup process (this runs comfortably within the Coiled free tier).

Reduce training time for CPU intensive models with scikit-learn and Coiled Functions

2023-09-01T00:00:00+00:00

Sep 1, 2023

3m read

Patrick Hoefler

You can use Coiled Run and Coiled Functions for easily running scripts and functions on a VM in the cloud.

In this post we’ll use Coiled Functions to train a RandomForestClassifer on a cloud-hosted machine that has enough cores to speed up our training process. The model parallelizes very well, which means that training time on my local machine is only bound by the number of cores available. Offloading this step onto a machine with 128 cores will save a lot of time during the training process.

Getting started

We will use random generated data for the purpose of this blog post. We will first train the model locally, before we use Coiled Funtions to offload the calculation to AWS.

from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import make_classification

def train():
    X, y = make_classification(n_samples=2_000_000, n_features=30, random_state=0, shuffle=False)
    clf = RandomForestClassifier(random_state=0, n_jobs=-1)
    clf.fit(X, y)
    return clf

We use 2 million samples and 30 features. This model takes around 12 minutes on my MacBook Air with 8 cores. The runtime can grow increasingly large if our dataset increases in size. We can speed up that process through using a more powerful VM.

Using `coiled functions` to train the model on a large VM

Coiled Functions come into the equation since we need access to a machine with a lot of cores. Coiled can connect to AWS or GCP and thus, use all resources that are available there. Let’s train the same model as above, but this time on an EC2 instance with 128 cores.

Adding a @coiled.function decorator to the function that executes our training step is the only modification we have to do. The decorator will tell Coiled that it should spin up a large VM on AWS and train the model there, and then return the trained model to our local machine.

import coiled
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import make_classification

@coiled.function(
    vm_type="c6i.32xlarge", # 128 cores, compute optimized
    keepalive="5 minutes",  # keep alive to train more models if necessary
)
def train():
    X, y = make_classification(n_samples=2_000_000, n_features=30, random_state=0, shuffle=False)
    clf = RandomForestClassifier(random_state=0, n_jobs=-1)
    clf.fit(X, y)
    return clf

Let’s execute the training step and return the model back to our local machine:

result = train()

RandomForestClassifier(n_jobs=-1, random_state=0)

Let’s take a look at the CPU utilization during the training step:

There is no need to adjust the other functions. Coiled will run our function on a VM in the cloud with enough resources.

Let’s take a brief look at the arguments to coiled.function():

vm_type: This specifies the type of EC2 instance. We are looking for an instance that as many cores as possible to speed up our training step.
keepalive: Keeps the VM alive so that we can run multiple queries against the data in memory.

coiled.function() will now start a VM in AWS with the specified EC2 instance. The VM is normally up and running in 1-2 minutes. Coiled will scan our local environment and replicate the same dependencies on this machine. We don’t have to specify an explicit Python environment. Inputs of your function are serialized and sent to the VM as well. Coiled will return our results back to our local machine.

Coiled would normally shut down the VM immediately after the Python interpreter finishes. This is mostly to reduce costs. We specified keepalive="5 minutes" to keep the VM alive for a few minutes after our Python interpreter finished. This ensures that new local runs can connect to the same VM avoiding the boot time of up to 2 minutes; we call this a warm start.

Conclusion

coiled functions enables you to run queries on a machine with as many cores as you need. This grants you access to computational resources that can be hard to replicate locally.

For more resources:

Check out the documentation
Use Coiled functions to train a neural net on a powerful GPU
Process hundreds of GBs of data with DuckDB using Coiled functions

Fine Performance Metrics and Spans

2023-08-23T00:00:00+00:00

While it’s trivial to measure the end-to-end runtime of a Dask workload, the next logical step - breaking down this time to understand if it could be faster - has historically been a much more arduous task that required a lot of intuition and legwork, for novice and expert users alike. We wanted to change that.

There are a few questions that someone who wants to optimize the performance of their Dask workflow typically asks:

How much time am I losing to GIL contention? Is it a problem?
How much time am I wasting in network transfer? Should I spend money increasing network throughput, or focus on optimizing data locality? What’s the potential benefit to end-to-end runtime?
I see a gray bar in the memory plot on the dashboard! This means that Dask is spilling data to disk as there isn’t enough RAM available. But how much of the end-to-end runtime is actually wasted on it? Should I spend money increasing available RAM, or developer time trying to reduce the memory footprint, or is it just not enough to bother?
The worker logs are saying that the event loop is getting blocked for many seconds at a time! Should I worry? Again, how much time, exactly, am I losing to this?
I just changed my algorithm, the version of Dask, or the version of other dependencies and now the end-to-end runtime is slower! Where am I spending more time than before?
I just doubled the size of my cluster but my workload is nowhere near twice as fast! Where did I accrue new inefficiency?

We recently introduced two complementary tools, Fine Performance Metrics and Spans, to answer these questions. They are available to all Dask users.

Fine Performance Metrics

Since version 2023.7.0, Dask automatically collects metrics which break down the end-to-end runtime of a computation by task and, within each task, by a series of activities taken to complete it.

In order to observe these metrics, you can simply

run your workload end-to-end
open the Dask dashboard (default for a LocalCluster: http://localhost:8787)
select More... -> Fine Performance Metrics

Alternatively, if you’re using JupyterLab and dask-labextension, you can drag the Fine Performance Metrics widget onto your Jupyter dashboard.

The grand total of the time shown should roughly add up to the end-to-end runtime of your workload, multiplied by the number of CPU cores on the cluster.

The central panel (Task execution, by activity) shows what activities the cluster spent its time on, cumulatively for all currently visible functions. The most important activities are:

thread-cpu

CPU time spent by tasks while running on workers. This is typically “good” time; in other words it’s the same time you would have spent if you ran the workload serially on a single CPU - but parallelized over however how many CPUs are available on your cluster.

thread-noncpu

Difference between wall clock time and CPU time spent by tasks while running on workers. This is typically I/O time, GPU time, CPU contention, or GIL contention. If you observe large amounts of this in your overall workload, you probably want to break it down by task and isolate those that are known to perform I/O or GPU activity.

If you observe large amounts of thread-noncpu time in what should be pure CPU tasks, the first thing you should investigate is GIL contention.

idle

Time during which a worker had a free thread, but nothing to run on it. This is typically caused by the workload not being able to fully utilize all threads on the cluster, network latency between scheduler and workers, or excessive CPU load on the scheduler. This measure does not include time spent while the whole cluster was completely idle.

disk-read, disk-write, compress, decompress

Time spent spilling data to disk due to not having enough memory available.

executor, offload, other

This is overhead from the Dask code and should be typically negligible. However, it can be inflated by GIL contention and by spill/unspill activity.

The left panel (Task execution, by function) shows the same information as the central one, but broken down by function.

The right panel (Send data, by activity) shows network transfer time. Note that most of it should be pipelined with task execution, so it may not have an impact. You should worry about this only if you have a very large idle time.

There is a filter that allows you to show only selected tasks. In this sample workflow, you can observe that most of the thread-noncpu time is concentrated - as expected - in tasks that are known to be I/O heavy. Here they are singled out:

And here’s all other functions that take a non-trivial amount of time:

In this sample workflow, this plot tells us an important piece of information: why is astype, which is a pure CPU function, spending so much time occupying workers’ threads but without accruing any CPU time? The answer, almost certainly, is that it doesn’t properly release the GIL.

Fine Performance Metrics collect more than wall clock timings. We can change the unit to bytes:

The above gives us insights on spill-to-disk activity. In this workflow, in 99% of the cases there was enough RAM to accommodate all the data and it was unnecessary to retrieve it from disk; in other words we have cache hits for 99% of the data, which means that we would not get any benefit if we increased RAM, but we would likely start seeing slowdowns if we reduced it.

This is corroborated by the time plots shown above, which show negligible time spent in disk-read or disk-write.

Spans

It can be non-trivial to correlate a task to an instruction on the client. The same function (e.g. a certain reduction) could also be applied in different sections of the code, to very different data. This problem becomes more severe as the complexity of the algorithm running on Dask increases.

Alternatively - or in addition - to the above, you may be in a multitenancy setup, where part of the load on the cluster was not caused by your client code, and you want to concentrate only on your own problem.

Spans let you break down fine performance metrics into macro-blocks (e.g. data load, preprocessing, etc.) by attaching meaningful tags to them.

To do so, you can use the span context manager inside the client code.

For example:

import dask.config
import dask.dataframe as dd
from distributed import Client, span

dask.config.set({"optimization.fuse.active": False})  # Don't lose task annotations
client = Client()

with span("Alice's workflow"):
    with span("data load"):
        df = dd.read_parquet(...)

    with span("ML preprocessing"):
        df = preprocess(df)
    
    with span("Model training"):        
        model = train(df)

model = model.compute()

The example will create the following spans on the scheduler:

Alice's workflow
Alice's workflow / data load
Alice's workflow / ML preprocessing
Alice's workflow / Model training

Each of the spans will be mapped to the tasks matching the segment of the Dask graph that was defined inside its context manager. The parent span will be mapped to all tasks of its children.

Tags are arbitrary and nothing stops you from parameterizing them; for example

>>> with span(f"{user}'s workflow"):
...     ...

Which may give you Alice's workflow / data load, Bob's workflow / data load, etc. This is useful for example if you want to observe either all the workload submitted by Alice, while hiding Bob’s activity, or alternatively all the data loading activity, regardless of who submitted it. You could use tags to store information on what dataset you’re processing, etc.

Viewing the spans

You can filter by span tags in the Fine Performance Metrics dashboard widget to filter your workload:

The Coiled dashboard

Coiled users enjoy full integration of both of these new features into the Coiled Dashboard. In addition to Fine Performance Metrics, the Coiled Dashboard also uses Spans to highlight portions of the workload plot and to tag code snippets.

The functionality works out of the box; all you need to do is make sure you’re running dask 2023.7.0 or later:

You can get started with Coiled for free at coiled.io/start.

Conclusions

These two new features should drastically improve users’ understanding of Dask performance. Are they useful to you? Are they easy to use? Is there a killer feature that you feel it’s missing? Let us know!

Data-proximate Computing with Coiled Functions

2023-08-10T00:00:00+00:00

Coiled Functions make it easy to improve performance and reduce costs by moving your computations next to your cloud data.

It’s common practice for data scientists and researchers to analyze data on their local work computer (often a laptop). This works great for data that’s stored locally on their machine. However, increasingly data is moving to cloud storage system like AWS S3 and Google Cloud Storage. When analyzing cloud data on a laptop, users end up needing to transfer their data from cloud storage to their laptop, which is both slow and expensive.

This post walks through how to use Coiled Functions to move your data-intensive computations to the cloud, alongside your data, to reduce costs and improve performance.

Example: Analyzing NYC taxi data

In the code snippet below we load NYC Taxi Data from AWS (data lives in us-east-1), process that data, and then upload the processed version of the data back to AWS. This type of data processing / cleaning is fairly common in data science workloads.

import time
import coiled
import pandas as pd
import s3fs
from dask.distributed import print

t_start = time.time()  # Get starting time

# Collect files to process
s3 = s3fs.S3FileSystem()
filenames = s3.ls("s3://nyc-tlc/trip data/")  # Data is in us-east-1
filenames = [
    f"s3://{fn}" for fn in filenames
    if "yellow_tripdata_2022" in fn
]

# Define our processing function
def process(filename):
    df = pd.read_parquet(filename)
    df = df[df.tip_amount != 0]

    outfile = "s3://oss-shared-scratch/jrbourbeau/" + filename.split("/")[-1]
    df.to_parquet(outfile)
    print(f"Finished {outfile}")

# Process files
print(f"Processing {len(filenames)} files")
for filename in filenames:
    print(f"Processing {filename}")
    process(filename)

print("Done!")

# Print total runtime
t_end = time.time()
print(f"Took {t_end - t_start / 60:0.2f} minutes")

Running this code on my laptop (M1 MacBook Pro) in Madison, WI takes 11.25 minutes.

In this case, most of the time is spent transferring data from AWS to my laptop and then uploading the processed data on my laptop back to AWS.

Run functions on the cloud

Given our data is already on AWS, let’s update our process function to also run on the cloud in AWS. We can use Coiled Functions to decorate our existing process function and have it automatically run on a VM in AWS.

@coiled.function(
    region="us-east-1",   # Same region as data
    vm_type="m6i.large",  # Enough memory to process files
)
def process(filename):
    # Everything inside the function stays the same
    ...

Note that we set region="us-east-1" and vm_type="m6i.large" to make sure the VM this process function runs on is in the same region as our data and has enough memory to process the data files (m6i.large instances have 8 GiB of memory).

When I add the @coiled.function decorator and rerun the script, it takes 2.82 minutes to run (note that includes the startup time for the cloud VM), which is about a 4x overall speedup. We also never had to transfer data outside of AWS, so we didn’t need to worry about data transfer costs.

Summary

Moving computations that process cloud-based data to be co-located alongside the data can both improve code runtime performance and reduce data transfer costs.

Coiled Functions is a simple API for moving existing Python functions to run on the cloud with minimal code changes. This makes developers’ lives a little easier and lowers the bar for transitioning code to run on the cloud.

Learn more

Check out the Coiled Functions documentation.
See the “How to Train a Neural Network on a GPU in the Cloud with Coiled Functions” blogpost from Patrick Hoefler for another handy use case for @coiled.function.
We highly value user feedback. If you try Coiled Functions, feel free to let us know about your experience in the coiled/feedback GitHub repo or by emailing support@coiled.io.
Get started with Coiled for free at coiled.io/start.

Dask, Dagster, and Coiled for Production Analysis at OnlineApp

2023-08-09T00:00:00+00:00

tl;dr

We show a simple integration between Dagster and Dask+Coiled. We discuss how this made a common problem, processing a large set of files every month, really easy.

The User and the Problem

Olá 👋 I’m Lucas, the data science tech lead at OnlineApp, a B2B company serving the Brazilian market.

Every month, the Brazilian government publishes a large set of CSV files with information on Brazilian companies that we use to better understand our market. We have lots of services that want to read this data, but before that’s possible we need to pre-process the data a bit:

Filter out some rows and columns that aren’t of interest to us;
Clean some values in several columns;
Join against other internal datasets we have;
Convert to Parquet and store in our own S3 storage;

We have many other services in our company that then rely on this data.

Architecture

For regular jobs like this that run every month, we typically use Dagster. We like Dagster because it’s easy to use and feels natural. This runs on a VM in the cloud, so we have fast access to other cloud data (like these CSV files).

Normally, we would use pandas for this kind of work, but this set of files is too large to fit into memory (around ≅150 million rows with several columns), so we switched to Dask. Unfortunately, the dataset is also too large for our VM running Dagster, so we use Coiled to ask for a cluster of larger machines, each with a fair amount of memory.

In principle, we want to do something like this:

import coiled
import dask.dataframe as dd
from dagster import op


@op
def run_operation(context):
    '''Processes all CSV files and saves the result as Parquet.'''

    cluster = coiled.Cluster(
        region="sa-east-1",
        memory="64 GiB",
        n_workers=100,
        ... # any other configuration options you need
    )

    client = cluster.get_client()

    df = dd.read_csv("...")

    ... # custom code specific to our problem

    df_result.to_parquet("...")

This would work, but we wanted to make this more natural by integrating it with Dagster. We’re excited to share what we did. It wasn’t a lot of code and resulted in a smooth experience.

Dagster + Coiled Integration

Dagster has the concept of a Resource. Resources typically model external components that Assets and Ops (short for “Operations”, like in the example above) interact with. For example, a resource might be a connection to a data warehouse like Snowflake or a service like a REST API.

To make Dask+Coiled’s resources more reusable and also take advantage of all the extra functionality that Dagster provides when working with resources (more on that later), we transformed our Coiled cluster into a Dagster resource. This is straightforward to do:

from contextlib import contextmanager
from dagster import resource
import coiled

@resource(
    config_schema={
        "n_workers": int,
        "worker_memory": str
    }
)
@contextmanager
def dask_coiled_cluster_resource(context):
    '''Yields a Dask cluster client.'''

    with coiled.Cluster(
        name=f"dagster-dask-cluster-{context.run_id}",
        n_workers=context.resource_config["n_workers"],
        worker_memory=context.resource_config["worker_memory"],
        region="sa-east-1",  # use your preferred region
        compute_purchase_option="spot_with_fallback", # save money with spot
    ) as cluster:
        with cluster.get_client() as client:

            context.log.info(
                f"Cluster created, dashboard link: {client.dashboard_link}"
            )

            yield client

To use this resource we add it to our job definition:

from dagster import job

@job(
    resource_defs={
        "dask_cluster": dask_coiled_cluster_resource,
    },
    config={
        "resources": {
            "dask_cluster": {
                # set default values, but you can change these
                # in the Dagster UI or later in the code
                "config": { 
                    "n_workers": 5,
                    "worker_memory": "64GiB"
                },
            },
        },
    }
)
def my_job():
    '''My job DAGs definition.'''
    ...
    result = my_op()
    ...

And finally, in the operations that need to have a Dask cluster to run computations, we inform these dependencies as follows:

from dagster import op


@op(
    # specify the resource key that you want to use in this op
    required_resource_keys={"dask_cluster"}
)
def my_op(context):
    '''Run my computation with a Dask cluster.'''

    # every computation inside this op will be executed with the Dask cluster
    ...

    # we can also access and use our cluster client directly
    context.resources.dask_cluster("...")

And that’s it, we now have full integration between Dagster, Dask and Coiled. Very easy to set up and maintain.

As a bonus, this is how we define a local Dask cluster for development and testing purposes:

from dagster import resource
from contextlib import contextmanager
from dask.distributed import Client, LocalCluster


@resource
@contextmanager
def dask_local_cluster_resource(context):
    '''Yields a local Dask cluster client.'''

    with LocalCluster() as cluster:
        with cluster.get_client() as client:
            context.log.info(
                f"Cluster created, dashboard link: {client.dashboard_link}"
            )
            yield client

Results

We now have nodes in our Dagster graph that are annotated as Dask+Coiled nodes. When Dagster runs these, they fire up a temporary Dask cluster, do larger scale processing than our Dagster VM can handle, and then clean up.

This image shows how the operation resource is listed in the Dagster UI:

(Click to enlarge)

And this is the Dagster Launchpad, here we can change several configurations before running our job, and one of those, thanks to defining our Coiled cluster as a resource, is choosing our cluster size, in a simple and parameterized way:

We’re now able to use our existing infrastructure to handle much larger datasets than we were able to use previously. Everything was very easy and cheap.

In Conclusion

In this post we walked through an example of how to use Dagster and Dask+Coiled to orchestrate a recurring monthly data processing workflow. With Dask and Coiled, we were able to process the data in parallel using a cluster of VMs in the cloud. Dagster is able to take advantage of these cloud resources, enabling us to process much larger datasets than we otherwise could have.

We hope others looking to scale their Dagster workflows to the cloud will find this guide useful. For more resources on getting started with Coiled, check out our getting started guide.

Process Hundreds of GB of Data with DuckDB in the Cloud

2023-08-07T00:00:00+00:00

Aug 7, 2023

4m read

Patrick Hoefler

DuckDB is great tool for running efficient queries on large datasets. When you want cloud data proximity or need more RAM, Coiled makes it easy to run your Python function in the cloud. In this post we’ll use Coiled Functions API to process the 150 GB Uber-Lyft dataset on a single machine with DuckDB.

Query Parquet files with DuckDB

We start with creating the SQL queries that we want to run against the data locally. We’ll use the DuckDB Python API.

def load_data(conn):
  # Load data into memory so that subsequent queries are fast
  conn.execute(
    '''
    CREATE TABLE test AS 
    SELECT * FROM read_parquet("s3://coiled-datasets/uber-lyft-tlc/*")
    '''
  )


def compute_percentage_of_tipped_rides(conn):
    # Run the actual query
   return conn.execute(
       '''
       SELECT hvfhs_license_num, sum(tipped) / count(tipped)
        FROM (select 
           *,  
           CASE WHEN tips > 0.0 then 1 ELSE 0 end as tipped
           from test) GROUP BY hvfhs_license_num
       '''
   ).fetchall()


def create_conn():
    import duckdb

    return duckdb.connect()


def query_results():
    conn = create_conn()
    load_data(conn)
    return compute_percentage_of_tipped_rides(conn)

These queries aren’t particularly fancy, they are meant to illustrate how we can process these files. If we execute these queries as is, it would pull all the data onto our machine. The whole dataset won’t fit in memory on most workstations, so let’s look at how Coiled can make this work. The loading would take a long time, even if we had enough memory.

How do I query data stored in S3 using DuckDB?

Coiled Functions come into the equation since we need access to machines that have enough resources and are also close to our data. Coiled can connect to AWS and thus, use all resources that are available there. These steps also work for data in Google Cloud Storage or Azure Blob Storage. We will go through the necessary steps execute these queries on a VM in the same region as our data with enough memory available.

The dataset is stored in the S3 bucket s3://coiled-datasets/uber-lyft-tlc in AWS region us-east-2. We’ll have to adapt our create_conn function to use load_aws_credentials and set the region. This DuckDB AWS extension builds on the httpfs extension to add authentication. Coiled forwards your credentials as short-lived (refreshable) STS tokens, which is more secure than putting long-lived AWS access keys (which consist of access key ID and secret access key).

def create_conn():
    import duckdb
    
    conn = duckdb.connect()
    conn.execute("CALL load_aws_credentials()")
    return conn

The next step is adding the @coiled.function decorator to the function that executes our queries. This syntax will tell Coiled that it should spin up a large VM on AWS and run the query there, and then return the result locally.

@coiled.function(
    vm_type="m6i.16xlarge", # 256 GB of RAM
    region="us-east-2",     # region of our data
    keepalive="5 minutes",  # keep alive to run multiple queries if necessary
)
def query_results():
    conn = create_conn()
    load_data(conn)
    return compute_percentage_of_tipped_rides(conn)

Let’s execute our queries and pull the results back to our local machine:

result = query_results()

print(result)
[
  ('HV0005', 0.1912300216459857), 
  ('HV0003', 0.1498555901186066), 
  ('HV0004', 0.09294857737045926), 
  ('HV0002', 0.08440046492889111),
]

The data are now all in memory on our VM in the cloud:

Memory usage goes up to 150 GB as we load the dataset on a big VM in the cloud.

There is no need to adjust the other functions. Coiled will run our query on a VM in the cloud with enough resources and close to our data.

Let’s take a brief look at the arguments to coiled.function():

vm_type: This specifies the type of AWS EC2 instance. We are looking for an instance that has enough memory to hold our data. This instance has 256GB, so this should be sufficient.
region: The region specifies the AWS region that our VM is started in. Our data are also in "us-east-2".
keepalive: Keeps the VM alive so that we can run multiple queries against the data in memory.

Coiled would normally shut down the VM immediately after the Python interpreter finishes. This is mostly to reduce costs. We specified keepalive="5 minutes" to keep the VM alive for a few of minutes after our Python interpreter finished. This ensures that new local runs can connect to the same VM avoiding the boot time of up to 2 minutes; we call this a warm start.

Conclusion

You can use Coiled serverless functions to run queries on a machine with as much memory as you want. This grants you access to computational resources that can be very close to your data. Doing data processing in the cloud becomes very easy with this functionality. Check out the docs for more details.

Want to run this example yourself? Get started with Coiled for free at coiled.io/start. This example run comfortably within the free tier.

Additional use cases:

Query data stored in S3 with Polars
Process 5,000 files in parallel on the cloud. Run the example yourself in the coiled/examples GitHub repo.
Train a RandomForest classifier with scikit-learn on multiple cores in the cloud

High Level Query Optimization in Dask

2023-08-04T00:00:00+00:00

Introduction

Dask DataFrame doesn’t currently optimize your code for you (like Spark or a SQL database would). This means that users waste a lot of computation. Let’s look at a common example which looks ok at first glance, but is actually pretty inefficient.

import dask.dataframe as dd

df = dd.read_parquet(
    "s3://coiled-datasets/uber-lyft-tlc/",  # unnecessarily reads all rows and columns
)
result = (
    df[df.hvfhs_license_num == "HV0003"]    # could push the filter into the read parquet call
    .sum(numeric_only=True)
    ["tips"]                                # should read only necessary columns
)

We can make this run much faster with a few simple steps:

df = dd.read_parquet(
    "s3://coiled-datasets/uber-lyft-tlc/",
    filters=[("hvfhs_license_num", "==", "HV0003")],
    columns=["tips"],
)
result = df.tips.sum()

Currently, Dask DataFrame wouldn’t optimize this for you, but a new effort that is built around logical query planning in Dask DataFrame will do this for you. This article introduces some of those changes that are developed in dask-expr.

You can install and try dask-expr with:

pip install dask-expr

We are using the NYC Taxi dataset in this post.

Dask Expressions

Dask expressions provides a logical query planning layer on top of Dask DataFrames. Let’s look at our initial example and investigate how we can improve the efficiency through a query optimization layer. As noted initially, there are a couple of things that aren’t ideal:

We are reading all rows into memory instead of filtering while reading the parquet files.
We are reading all columns into memory instead of only the columns that are necessary.
We are applying the filter and the aggregation onto all columns instead of only "tips".

The query optimization layer from dask-expr can help us with that. It will look at this expression and determine that not all rows are needed. An intermediate layer will transpile the filter into a valid filter-expression for read_parquet:

df = dd.read_parquet(
    "s3://coiled-datasets/uber-lyft-tlc/",
    filters=[("hvfhs_license_num", "==", "HV0003")],
)
result = df.sum(numeric_only=True)["tips"]

This still reads every column into memory and will compute the sum of every numeric column. The next optimization step is to push the column selection into the read_parquet call as well.

df = dd.read_parquet(
    "s3://coiled-datasets/uber-lyft-tlc/",
    columns=["tips"],
    filters=[("hvfhs_license_num", "==", "HV0003")],
)
result = df.sum(numeric_only=True)

This is a basic example that you could rewrite by hand. Use cases that are closer to real workflows might potentially have hundreds of columns, which makes rewriting them very strenuous if you need a non-trivial subset of them.

Let’s take a look at how we can achieve this. dask-expr records the expression as given by the user in an expression tree:

result.pprint()

Projection: columns='tips'
  Sum: numeric_only=True
    Filter:
      ReadParquet: path='s3://coiled-datasets/uber-lyft-tlc/'
      EQ: right='HV0003'
        Projection: columns='hvfhs_license_num'
          ReadParquet: path='s3://coiled-datasets/uber-lyft-tlc/'

This tree represents the expression as is. We can observe that we would read the whole dataset into memory before we apply the projections and filters. One observation of note: It seems like we are reading the dataset twice, but Dask is able to fuse tasks that are doing the same to avoid computing these things twice. Let’s reorder the expression to make it more efficient:

result.simplify().pprint()

Sum: numeric_only=True
  ReadParquet: path='s3://coiled-datasets/uber-lyft-tlc/' 
               columns=['tips'] 
               filters=[('hvfhs_license_num', '==', 'HV0003')] 

This looks quite a bit simpler. dask-expr reordered the query and pushed the filter and the column projection into the read_parquet call. We were able to remove quite a few steps from our expression tree and make the remaining expressions more efficient as well. This represents the steps that we did manually in the beginning. dask-expr performs these steps for arbitrary many columns without increasing the burden on the developers.

These are only the two most common and easy to illustrate optimization techniques from dask-expr. Some other useful optimizations are already available:

len(...) will only use the Index to compute the length; additionally we can ignore many operations that won’t change the shape of a DataFrame, like a replace call.
set_index and sort_values won’t eagerly trigger computations.
Better informed selection of merge algorithms.
…

We are still adding more optimization techniques to make Dask DataFrame queries more efficient.

Try it out

The project is in a state where interested users should try it out. We published a couple of releases. The API covers a big chunk of the Dask DataFrame API, and we keep adding more. We have already observed very impressive performance improvements for workflows that would benefit from query optimization. Memory usage is down for these workflows as well.

We are very much looking for feedback and potential avenues to improve the library. Please give it a shot and share your experience with us.

dask-expr is not integrated into the main Dask DataFrame implementation yet. You can install it with:

pip install dask-expr

The API is very similar to what Dask DataFrame provides. It exposes mostly the same methods as Dask DataFrame does. You can use the same methods in most cases.

import dask_expr as dd

You can find a list of supported operations in the Readme. This project is still very much in progress. The API might change without warning. We are aiming for weekly releases to push new features out as fast as possible.

Why are we adding this now?

Historically, Dask focused on flexibility and smart scheduling instead of query optimization. The distributed scheduler built into Dask uses sophisticated algorithms to ensure ideal scheduling of individual tasks. It tries to ensure that your resources are utilized as efficient as possible. The graph construction process enables Dask users to build very flexible and complicated graphs that reach beyond SQL operations. The flexibility that is provided by the Dask futures API requires very intelligent algorithms, but it enables users to build highly sophisticated graphs. The following picture shows the graph for a credit risk model:

The nature of the powerful scheduler and the physical optimizations enables us to build very complicated programs that will then run efficiently. Unfortunately, the nature of these optimizations does not enable us to avoid scheduling work that is not necessary. This is where the current effort to build high level query optimization into Dask comes in.

Conclusion

Dask comes with a very smart distributed scheduler but without much logical query planning. This is something we are rectifying now through building a high level query optimizer into Dask DataFrame. We expect to improve performance and reduce memory usage for an average Dask workflow.

This API is read for interested users to play around with. It covers a good chunk of the DataFrame API. The library is under active development, we expect to add many more interesting things over the coming weeks and months.

Easy Heavyweight Serverless Functions

2023-08-01T00:00:00+00:00

tl;dr

What is the easiest way to run Python code in the cloud, especially for compute jobs?

We briefly compare common options (Lambda, EC2, Fargate, Modal) and then pitch a new contender: Coiled Run.

Motivation to Run Python Code in the Cloud

We want to run jobs in the cloud for a few common reasons:

Data proximity: Our data lives in the cloud, and we want to process it where it lives to avoid costly egress charges
Speed and Scale: We have many things to process and don’t want to wait
Hardware access: We want GPUs or big-memory machines
Always-on Coordination: We want to respond to actions in the real world, like in a web server or in response to data arriving

This post focuses on the first three use cases which require more heavyweight computing. As an example, maybe we want to load, manipulate, and store many files living in cloud storage:

def process(filename):                         # Expensive function
    df = pd.read_csv(filename)
    df = df[df.value > 0]
    df.to_parquet(filename[:-4] + ".parquet")

filenames = s3.glob("s3://my-bucket/*.csv")    # Lots of files
for filename in filenames:                     # Run many times
    process(filename)

Let’s see what kinds of options we have.

Cloud Infrastructure to Run Single Jobs

In general we have two classes of options:

Big VMs (EC2, ECS) which give you a big VM
- Pros:
  - Hardware flexibility: you can use any hardware available in the cloud
  - Cheap: typically $0.05 per CPU-hour
- Cons:
  - Annoying to set up / inaccessible
  - Takes at least a minute to turn on (but more likely an hour of human time)
Serverless technologies (AWS Lambda) which run a function for you, hiding the VM
- Pros:
  - Much easier to use
  - Turns on quickly (seconds)
- Cons:
  - Easier, but not yet easy. These are still kinda annoying to set up.
  - Hardware limitations
  - Expensive: typically $0.20 per CPU-hour, can’t use Spot

Most people I see using cloud for ad-hoc computing choose VMs. They spin up a big instance, use it for a while, and then spin it down (hopefully). People who start with serverless tend to move over to VMs, usually due to some flexibility limitation. AWS Lambda is often great, except if any of the following are true:

You need a big machine
You want accelerated hardware like GPUs
You have a large software environment
Your functions last tens of minutes
Your functions last more than a minute, and you’re cost conscious

Third party orchestrators

Products like Argo, Prefect, Airflow, Coiled+Dask (us!), Kubeflow, Anyscale+Ray, Dagster and more will also happily run jobs for you, but this is as part of a broader abstraction. If you just want to run a single job, they’re all overkill.

The best product experience I’ve seen in this space is Modal Labs. It’s pretty slick. Like the others though, it still has non-trivial ceremony. It also has some other limitations like you have to run in their cloud accounts (so data privacy is a concern), it’s only on one region (so data egress costs can be non-trivial) and they cost $0.20 per CPU-hour (which is a 4x markup). Generally a great product experience though, and we hear only positive things. I encourage people to check them out.

Coiled Run

So now we get to our new solution. There are two APIs, they look like this:

CLI:
```
$ coiled run echo "Hello, world"
```

Python:

import coiled

@coiled.function(memory="512 GiB", region="us-west-2")
def f():
   return "Hello, world!"

f()

Coiled was originally built to deploy distributed Dask clusters in the cloud, a harder problem than serverless execution. However, the infrastructure we built for Dask transfers over surprisingly well. We started this effort because we saw Coiled users spinning up single-node Dask clusters (rather than hundred-node clusters) to run a single task (rather than millions). We took the hint and built these new APIs on top of the existing infrastructure. Here you go:

Run a script

The easiest thing is to run a standalone script in the cloud. Here we run a script in Frankfurt on a machine with 128 cores.

$ coiled run python myscript.py --region eu-central-1 --vm-type m6i.32xlarge

Optionally specify useful keywords:

$ coiled run python myscript.py \
    --region us-east-1 \
    --vm-type m6i.32xlarge \

$ coiled run python pytorch-train.py \
    --region us-west-1 \
    –-vm-type g5.2xlarge \  # GPU instance type
    --container ...

Doesn’t even need to be a Python script

$ coiled run echo "Hello, world"

Coiled creates a VM, synchronizes all of your local software packages, cloud credentials, files, etc.., runs your script, and then shuts down the VM.

Here are some more fully worked examples:

Run a Python function

Alternatively you can run your Python code locally, but decorate individual functions to run remotely, as with my_function below.

import coiled

@coiled.function(
    memory="256 GiB",
    region="us-east-1",
)
def my_function(...):
    return ...

result = my_function(*args, **kwargs)

Here are some more fully worked examples.

Pros and Cons

These APIs were accidental, but turned out to be easy and efficient cloud functions.

Pros:

Easiest API we can find
Any hardware / architecture (uses EC2 instances)
Runs securely in your account
Runs in any region, close to your data
Base cloud cost (uses EC2 instances)

Cons:

Minute-long start times for your first function call (but subsequent calls have only millisecond delays)
Third-party service (requires separate sign-up/accounting)

This essentially offers the flexibility and cost efficiency of EC2 instances, but without the setup pain.

We think that the UX is smooth, mostly inheriting from other features we built for Dask clusters like the following:

Automatic package synchronization: we transmit the same package versions you have locally to your cloud workers
Credential forwarding: we securely forward your AWS permissions to remote workers (safely using AWS STS)
Cost/user management controls: we help you constrain costs on a per-user basis
General progress / UX polish: things just look nice.

Opportunity

We’re excited about Coiled functions because they are …

Easy to use for Python developers unfamiliar with cloud computing
Flexible with their ability to use any hardware in any region
Cheap being backed by foundational technologies like EC2 Fleets.

Dask is super-powerful (far more powerful than what’s described in this post) but there are lots of folks for whom Dask is overkill. We hope that these simpler interfaces are a good onramp to data-proximate and parallel cloud computing.

Interested in playing? The following should work if you have cloud credentials:

pip install coiled
coiled setup  # this will connect Coiled to your cloud account
coiled run echo "Hello, world"

For more on how to get started, see

How to Train a Neural Network on a GPU in the Cloud with coiled functions

2023-07-24T00:00:00+00:00

Jul 24, 2023

4m read

Patrick Hoefler

We recently pushed out two new and experimental features coiled run and coiled functions which is a deviation of coiled run. We are excited about both of them because they:

Allow users to scale up any given program on any hardware in the cloud.
Make GPUs easily accessible without going through the pains of setting up environments in the cloud.

This post will provide an example how to utilize coiled functions to seamlessly train a neural network on a GPU that is hosted in the cloud.

Getting started

We have to start with creating a model on our local machine before we can start worrying about training it. This blog post is not dedicated to figuring out a fancy model, we will utilize the Net model that is given in the PyTorch tutorials.

We can simply add the model definition to our python script. There is no need to do anything different. Similarly, we will use the transformer that is given there as well.

The next step is creating a function that we can use to train the model:

def train(transform):
    device = torch.device("cpu")
    net = Net()
    net = net.to(device)

    trainset = torchvision.datasets.CIFAR10(
        root='./data', train=True, download=True, transform=transform,
    )
    trainloader = torch.utils.data.DataLoader(
        trainset, batch_size=4, shuffle=True, num_workers=2,
    )
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

    for i, data in enumerate(trainloader, 0):
        inputs, labels = data[0].to(device), data[1].to(device)

        optimizer.zero_grad()

        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

    return net

We can now train our model:

if __name__ == "__main__":
    train(transform)

This will train our model on the CPU of our local machine. The training is reasonably quick for such a small model, but training time will grow exponentially as our model gets larger or if we are using a significantly bigger dataset. Training the model on the CPU won’t be sufficient anymore. Additionally, there are a lot of machines out there that don’t have GPUs built into them. For example, I’m using a MacBook Pro with an M2 CPU, which means my machine doesn’t support cuda. Consequently, we need a different solution to make these steps accessible for folks who don’t have access to a local GPU.

Using `coiled functions` to train the model on a cloud-hosted GPU

Coiled functions come into the equation if you need access to resources that aren’t available locally. Coiled can connect to AWS or GCP and thus, use all resources that are available there. We will go through the necessary steps to train our model on a GPU that is hosted on AWS instead of our local CPU.

The first step includes defining a Python environment to run our computations. We simply include PyTorch, CUDA, and Coiled, that’s it. Generally, you should use the same Python version that is installed locally.

import coiled

coiled.create_software_environment(
    name="pytorch",
    conda={
        "channels": ["pytorch", "nvidia", "conda-forge", "defaults"],
        "dependencies": [
            "python=3.11",
            "coiled",
            "pytorch",
            "torchvision",
            "torchaudio",
            "cudatoolkit",
            "pynvml",
        ],
    },
    gpu_enabled=True,
)

Coiled will create a Python environment for you. This step is only necessary when running your script for the first time. The resulting environment is cached which makes further runs more efficient.

The next step is adding the @coiled.function decorator to our training function that tells our program we want to execute said function on a machine in the cloud.

@coiled.function(
    vm_type="g5.xlarge", # GPU instance type
    region="us-west-2",
    software="pytorch",
)

Additionally, we have to tell PyTorch that we want to train the model on the GPU.

def train():
    import torch
    # tell PyTorch to use the GPU
    device = torch.device("cuda:0")
    ...
    return net.to(torch.device("cpu"))

Putting this all together:

import coiled
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))])

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

@coiled.function(
    vm_type="g5.xlarge",
    region="us-west-2",
    software="pytorch",
)
def train(transform):
    import torch
    device = torch.device("cuda:0")

    net = Net()
    net = net.to(device)

    trainset = torchvision.datasets.CIFAR10(
        root='./data', train=True, download=True, transform=transform,
    )
    trainloader = torch.utils.data.DataLoader(
        trainset, batch_size=4, shuffle=True
    )
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

    for i, data in enumerate(trainloader, 0):
        inputs, labels = data[0].to(device), data[1].to(device)

        optimizer.zero_grad()

        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

    return net.to(torch.device("cpu"))


if __name__ == "__main__":
    train(transform)

Let’s take a brief look at the arguments to coiled.function():

vm_type: This specifies the type of EC2 instance. We are looking for an instance that has a GPU attached to it. The G5 family has Nvidia GPUs attached to it. The smallest version is sufficient for our example, but you can choose instances with up to 8 GPUs.
region: The region specifies the AWS region that our VM is started in. We observed that "us-west-2" is a region where GPUs are easier to get.
software: This specifies the coiled software environment that is installed. This corresponds to the environment that we previously created.

coiled.function() will now start a VM in AWS with the specified EC2 instance. The VM is normally up and running in 1-2 minutes. The previously specified Python environment is installed automatically. Coiled executes the function on said VM. Inputs of your function are serialized and sent to the VM as well. It makes sense to download the training data on the VM to reduce time that is spent sending data to AWS. The function returns our model back to our local machine so that we can use it locally without depending on AWS.

Coiled will shut down the VM immediately after the Python interpreter finishes. This is mostly to reduce costs. You can specify a certain amount of time that the VM is kept alive through keepalive="5 minutes". This ensures that new local runs can connect to the same VM avoiding the boot time of up to 2 minutes; we call this a warm start.

Conclusion

coiled functions enables you to seamlessly port the training process for a neural network from your local machine to AWS or GCP. This grants everyone access to multiple GPUs or huge machines independent of the local machine that is actually used. Training a neural network on a GPU becomes as easy as adding a decorator to the training function.

Dask performance benchmarking put to the test: Fixing a pandas bottleneck

2023-06-23T00:00:00+00:00

Jun 23, 2023

4m read

Patrick Hoefler, Hendrik Makait

Getting notified of a significant performance regression the day before release sucks, but quickly identifying and resolving it feels great!

We were getting set up at our booth at JupyterCon 2023 when we received a notification: An engineer on our team had spotted a significant performance regression in Dask. With an impact of 40% increased runtime, it blocked the release planned for the next day!

Luckily, the other attendees still focused on coffee and breakfast, so we commandeered an abandoned table next to our booth and got to work.

Performance testing at Coiled

The performance problem had been flagged by the automated performance testing for Dask that we developed at Coiled. If you have not read Guido Imperiale’s blog post on our approach to performance testing, here is a summary: With coiled/benchmarks, we created a benchmark suite that contains a variety of common workloads and operations with Dask, including standardized ones like the h2oai/db-benchmark.

It also contains tooling that allows us to do two things:

Automatically detect performance regressions in Dask and raise them as issues.
Run A/B tests to assess the performance impact of different versions of Dask, upstream packages, or cluster configurations.

While the former started this journey, the latter will also come in handy soon.

Identifying the problem

Our automated regression testing had alerted us that test_h2o.py::test_q8 had experienced a significant increase in runtime across all data sizes and file formats. From the historical report of our benchmarking suite, we could see that dask/dask and dask/distributed were unlikely to be the culprit: Nothing had changed on dask/dask when the performance started to degrade, and there was only one unrelated change on dask/distributed. That left us with the Coiled platform and upstream packages as possible candidates.

After digging deeper into the cluster data, we noticed that pandas had been upgraded from 1.5.3 to 2.0.1. A major upgrade to pandas at the same time a dataframe-based workload shows degrading performance? That’s suspicious!

To confirm this suspicion, we ran an A/B test based on the current Dask release (2023.4.1 at the time), testing the impact of the pandas upgrade. The results were clear: The runtime increased significantly with pandas=2.0.1 (sample).

Having shown that pandas caused for the performance degradation and that we could reproduce it with the current Dask release, our release process for 2023.5.0 was cleared.

To further analyze the problem, we also derived a minimal local reproducer from the original workload:

from dask.distributed import Client

client = Client()
uri = "s3://coiled-datasets/h2o-benchmark/N_1e7_K_1e2_parquet/*.parquet"
ddf = dd.read_parquet(uri, engine="pyarrow", storage_options={"anon": True}).persist()
wait(ddf)

ddf = ddf[["id6", "v1", "v2", "v3"]]
(
    ddf[~ddf["v3"].isna()][["id6", "v3"]]
    .groupby("id6", dropna=False, observed=True)
    .apply(
        lambda x: x.nlargest(2, columns="v3"),
        meta={"id6": "Int64", "v3": "float64"},
    )[["v3"]]
).compute()

Investigating the pandas performance degradation

The only user-visible thing that changed between pandas 1.5.3 and pandas 2.0.1 was the default value of group_keys in GroupBy. Switching to group_keys=False with version 2.0.1 got us back to the initial runtime. Now that we knew that pandas was to blame for the performance degradation, we had to create a reproducer in plain pandas to help fix the issue.

df = pd.DataFrame(
    {
        "foo": np.random.randint(1, 50_000, (100_000, )), 
        "bar": np.random.randint(1, 100_000, (100_000, )),
    },
)

df.groupby(
    "foo", group_keys=False
).apply(lambda x: x.nlargest(2, columns="bar"))

group_keys=False: approx. 11 seconds
group_keys=True: approx. 15 seconds

Experimenting a bit showed us that the bottleneck got even worse while increasing the number of groups during the groupby calculation. We settled on this version which is 30% slower with group_keys=True, enough to be able to troubleshoot the problem. There was no obvious reason why the changed value should bring a significant slowdown.

%prun showed us that the time was almost exclusively spent in a post-processing step that combines all groups via concat.

Addressing the performance degradation

Let’s look at how both cases differ. The new version passes the grouping levels to concat, which are used to construct the resulting Index levels. This shouldn’t be that slow though. Investigations showed that this runs through a code-path that is very inefficient!

Looking closer at the results of %prun pointed us to one specific loop that took up most of the runtime. This loop calculates the codes for the resulting index based on the provided levels. It’s slow, really slow! Every single element provided as keys, which represent the number of groups, is checked against the whole level, which explains our previous observation that the runtime got worse with an increasing number of groups. You can check out the pandas user guide if you are interested in situations where this is useful. Fortunately, we have a convenient advantage in case of groupby. We know beforehand that every key equals the specific level. We added a fast-path that exploits this knowledge getting the runtime of this step more or less to zero.

This change resulted in a small PR that cut the runtime of group_keys=True to approximately 11 seconds as well.

Conclusion

Now that we made our pandas reproducer run 30% faster, we have to check whether we accomplished our initial objective. Re-running the local Dask reproducer should give us an idea about the impact on Dask. We got the performance down to 22 seconds as well! Promising news that saved our plans for the evening!

Unfortunately, we had to wait until pandas 2.0.2 was released to run a proper benchmark.

This looks great! Our small pandas change translated to our Dask query and got performance back to the previous level!

Coiled notebooks

2023-06-14T00:00:00+00:00

We recently pushed out a new, experimental notebooks feature for easily launching Jupyter servers in the cloud from your local machine. We’re excited about Coiled notebooks because they:

Allow users to easily scale up their computations on a single machine in the cloud.
Make GPUs more accessible.

How to launch a notebook

You can launch a Coiled notebook with:

pip install "coiled>=0.6.8"

coiled notebook start

This will provision an instance on your cloud provider, set up a JupyterLab session, and connect you to that session in a web browser.

In addition to the advantages you get from Coiled Dask clusters like:

Automatic software environment syncing (docs)
The same cloud permissions you have locally
Run in any region on AWS/GCP

Coiled notebooks also support automatic file syncing between your local machine and the remote Jupyter server in the cloud. This lets you bring your local files with you, edit them in the cloud, and have those changes brought back locally.

To run a notebook with automatic file syncing, install Mutagen

brew install mutagen-io/mutagen/mutagen

and use the --sync flag:

coiled notebook start --sync  # Automatically syncs your files

Example: Scaling up

Sometimes users just want to spin up a big machine in the cloud and use a local Dask cluster on that machine for their computations. This is a great fit when working with ~10-100 GB datasets where a single, large machine can be effective.

With Coiled notebooks, you can use the --vm-type option to specify what instance type you want to use:

coiled notebook start --vm-type m6i.32xlarge  # 128 Cores and 512 GB of RAM

Then on the Jupyter server, Dask’s LocalCluster will utilize the full resources of the instance (e.g. 128 core and ~500 GB of memory in this case)

Example: Easy access to GPUs

Several libraries in the PyData space like PyTorch, XGBoost, and the entire RAPIDS ecosystem (e.g. cuDF, cuML, etc) support GPU acceleration. This can lead to large performance improvements in a variety of applications. However, today many data scientists and developers don’t readily have access to GPUs.

Coiled notebooks make it easy to run on GPU-enabled instances, lowering the barrier to entry for GPU accelerated computations. For example, you can use the --gpu flag and --container option to launch a Coiled notebook on a GPU instance inside a pre-built RAPIDS docker image:

# GPU enabled notebook running inside a RAPIDS docker image
coiled notebook start --gpu --container nvcr.io/nvidia/rapidsai/rapidsai:23.04-cuda11.8-runtime-ubuntu20.04-py3.10

Tip

The --gpu flag defaults to a g4dn.xlarge (NVIDIA T4) instance, but you can use the --vm-type option to specific any GPU instance you want to run on.

Try out and report

We’re excited about Coiled notebooks and think they can be useful to users. They’re also a work in progress. We value feedback from users to inform future development. Feel free to try Coiled notebooks and let us know about your experience in the Coiled feedback issue tracker or by emailing support@coiled.io.

Utilizing PyArrow to improve pandas and Dask workflows

2023-06-05T00:00:00+00:00

Jun 5, 2023

14m read

Patrick Hoefler

Get the most out of PyArrow support in pandas and Dask right now

Introduction

This post investigates where we can use PyArrow to improve our pandas and Dask workflows right now. General support for PyArrow dtypes was added with pandas 2.0 to pandas and Dask. This solves a bunch of long-standing pains for users of both libraries. pandas users often complain to me that pandas does not support missing values in arbitrary dtypes or that non-standard dtypes are not very well supported. A particularly annoying problem for Dask users is running out of memory with large datasets. PyArrow backed string columns consume up to 70% less memory compared to NumPy object columns and thus have the potential to mitigate this problem as well as providing a huge performance improvement.

Support for PyArrow dtypes in pandas, and by extension Dask, is still relatively new. I would recommend caution when opting into the PyArrow dtype_backend until at least pandas 2.1 is released. Not every part of both APIs is optimized yet. You should be able to get a big improvement in certain workflows though. This post will go over a couple of examples where I’d recommend switching to PyArrow right away, because it already provides huge benefits.

Dask itself can benefit in various ways from PyArrow dtypes. We will investigate how PyArrow backed strings can easily mitigate the pain point of running out of memory on Dask clusters and how we can improve performance through utilizing PyArrow.

I am part of the pandas core team and was heavily involved in implementing and improving PyArrow support in pandas. I’ve recently joined Coiled where I am working on Dask. One of my tasks is improving the PyArrow integration.

General overview of PyArrow support

PyArrow dtypes were initially introduced in pandas 1.5. The implementation was experimental and I wouldn’t recommend using it on pandas 1.5.x. Support for them is still relatively new. pandas 2.0 provides a huge improvement, including making opting into PyArrow backed DataFrames easy. We are still working on supporting them properly everywhere, and thus they should be used with caution until at least pandas 2.1 is released. Both projects work continuously to improve support throughout Dask and pandas.

We encourage users to try them out! This will help us to get a better idea of what is still lacking support or is not fast enough. Giving feedback helps us improve support and will drastically reduce the time that is necessary to create a smooth user experience.

Dataset

We will use the taxi dataset from New York City that contains all Uber and Lyft rides. It has some interesting attributes like price, tips, driver pay and many more. The dataset can be found here (see terms of service) and is stored in parquet files. When analyzing Dask queries, we will use a publicly available S3 bucket to simplify our queries: s3://coiled-datasets/uber-lyft-tlc/. We will use the dataset from December 2022 for our pandas queries, since this is the maximum that fits comfortably into memory on my machine (24GB of RAM). We have to avoid stressing our RAM usage, since this might introduce side effects when analyzing performance.

We will also investigate the performance of read_csv. We will use the Crimes in Chicago dataset that can be found here.

Dask cluster

There are various different options to set up a Dask cluster, see the Dask documentation for a non-exhaustive list of deployment options. I will use Coiled to create a cluster on AWS with 30 machines through:

import coiled

cluster = coiled.Cluster(
    n_workers=30,
    name="dask-performance-comparisons",
    region="us-east-2",  # this is the region of our dataset
    worker_vm_type="m6i.large",
)

Coiled is connected to my AWS account. It creates the cluster within my account and manages all resources for me. 30 machines are enough to operate on our dataset comfortably. We will investigate how we can reduce the required number of workers to 15 through some small modifications.

pandas StringDtype backed by PyArrow

We begin with a feature that was originally introduced over 3 years ago in pandas 1.0. Setting the dtype in pandas or Dask to string returns an object with StringDtype. This feature is relatively mature and should provide a smooth user experience.

Historically, pandas represented string data through NumPy arrays with dtype object. NumPy object data is stored as an array of pointers pointing to the actual data in memory. This makes iterating over an array containing strings very slow. pandas 1.0 initially introduced said StringDtype that allowed easier and consistent operations on strings. This dtype was still backed by Python strings and thus, wasn’t very performant either. Rather, it provided a clear abstraction of string data.

pandas 1.3 finally introduced an enhancement to create an efficient string dtype. This datatype is backed by PyArrow arrays. PyArrow provides a data structure that enables performant and memory efficient string operations. Starting from that point on, users could use a string dtype that was contiguous in memory and thus very fast. This dtype can be requested through string[pyarrow]. Alternatively, we can request it by specifying string as the dtype and setting:

pd.options.mode.string_storage = "pyarrow"

Since Dask builds on top of pandas, this string dtype is available here as well. On top of that, Dask offers a convenient option that automatically converts all string-data to string[pyarrow].

dask.config.set({"dataframe.convert-string": True})

This is a convenient way of avoiding NumPy object dtype for string columns. Additionally, it has the advantage that it creates PyArrow arrays natively for I/O methods that operate with Arrow objects. On top of providing huge performance improvements, PyArrow strings consume significantly less memory. An average Dask DataFrame with PyArrow strings consumes around 33-50% of the original memory compared to NumPy object. This solves the biggest pain point for Dask users that is running out of memory when operating on large datasets. The option enables global testing in Dask’s test suite. This ensures that PyArrow backed strings are mature enough to provide a smooth user experience.

Let’s look at a few operations that represent typical string operations. We will start with a couple of pandas examples before switching over to operations on our Dask cluster.

We will use df.convert_dtypes to convert our object columns to PyArrow string arrays. There are more efficient ways of getting PyArrow dtypes in pandas that we will explore later. We will use the Uber-Lyft dataset from December 2022, this file fits comfortably into memory on my machine.

import pandas as pd

pd.options.mode.string_storage = "pyarrow"

df = pd.read_parquet(
    "fhvhv_tripdata_2022-10.parquet",
    columns=[
        "tips", 
        "hvfhs_license_num", 
        "driver_pay", 
        "base_passenger_fare", 
        "dispatching_base_num",
    ],
)
df = df.convert_dtypes(
    convert_boolean=False, 
    convert_floating=False, 
    convert_integer=False,
)

Our DataFrame has NumPy dtypes for all non-string columns in this example. Let’s start with filtering for all rides that were operated by Uber.

df[df["hvfhs_license_num"] == "HV0003"]

This operation creates a mask with True/False values that specify whether Uber operated a ride. This doesn’t utilize any special string methods, but the equality comparison dispatches to PyArrow. Next, we will use the String accessor that is implemented in pandas and gives you access to all kinds of string operations on a per-element basis. We want to find all rides that were dispatched from a base starting with "B028".

df[df["dispatching_base_num"].str.startswith("B028")]

startswith iterates over our array and checks whether every string starts with the specified substring. The advantage of PyArrow is easy to see. The data are contiguous in memory, which means that we can efficiently iterate over them. Additionally, these arrays have a second array with pointers that point to the first memory address of every string, which makes computing the starting sequence even faster.

Finally, we look at a GroupBy operation that groups over PyArrow string columns. The calculation of the groups can dispatch to PyArrow as well, which is more efficient than factorizing over a NumPy object array.

df.groupby(
    ["hvfhs_license_num", "dispatching_base_num"]
).mean(numeric_only=True)

Let’s look at how these operations stack up against DataFrames where string columns are represented by NumPy object dtype.

The results are more or less as we expected. The string based comparisons are significantly faster when performed on PyArrow strings. Most string accessors should provide a huge performance improvement. Another interesting observation is memory usage, it is reduced by roughly 50% compared to NumPy object dtype. We will take a closer look at this with Dask.

Dask mirrors the pandas API and dispatches to pandas for most operations. Consequently, we can use the same API to access PyArrow strings. A convenient option to request these globally is the option mentioned above, which is what we will use here:

dask.config.set({"dataframe.convert-string": True})

One of the biggest benefits of this option during development is that it enables easy testing of PyArrow strings globally in Dask to make sure that everything works smoothly. We will utilize the Uber-Lyft dataset for our explorations. The dataset takes up around 240GB of memory on our cluster. Our initial cluster has 30 machines, which is enough to perform our computations comfortably.

import dask
import dask.dataframe as dd
from distributed import wait

dask.config.set({"dataframe.convert-string": True})

df = dd.read_parquet(
    "s3://coiled-datasets/uber-lyft-tlc/",
    storage_options={"anon": True},
)
df = df.persist()
wait(df)  # Wait till the computation is finished

We persist the data in memory so that I/O performance does not influence our performance measurements. Our data is now available in memory, which makes access fast. We will perform computations that are similar to our pandas computations. One of the main goals is to show that the benefits from pandas will translate to computations in a distributed environment with Dask.

One of the first observations is that the DataFrame with PyArrow backed string columns consumes only 130GB of memory, only half of what it consumed with NumPy object columns. We have only a few string columns in our DataFrame, which means that the memory savings for string columns are actually higher than around 50% when switching to PyArrow strings. Consequently, we will reduce the size of our cluster to 15 workers when performing our operations on PyArrow string columns.

cluster.scale(15)

We measure the performance of the mask-operation and one of the String accessors together through subsequent filtering of the DataFrame.

df = df[df["hvfhs_license_num"] == "HV0003"]
df = df[df["dispatching_base_num"].str.startswith("B028")]
df = df.persist()
wait(df)

We see that we can use the same methods as in our previous example. This makes transitioning from pandas to Dask relatively easy.

Additionally, we will again compute a GroupBy operation on our data. This is significantly harder in a distributed environment, which makes the results more interesting. The previous operations parallelize relatively easy onto a large cluster, while this is harder with GroupBy.

df = df.groupby(
    ["hvfhs_license_num", "dispatching_base_num"]
).mean(numeric_only=True)

df = df.persist()
wait(df)

We get nice improvements by factors of 2 and 3. This is especially intriguing since we reduced the size of our cluster from 30 machines to 15, reducing the cost by 50%. Subsequently, we also reduced our computational resources by a factor of 2, which makes our performance improvement even more impressive. Thus, the performance improved by a factor of 4 and 6 respectively. We can perform the same computations on a smaller cluster, which saves money and is more efficient in general, and still get a performance boost out of it.

Summarizing, we saw that PyArrow string-columns are a huge improvement when comparing them to NumPy object columns in DataFrames. Switching to PyArrow strings is a relatively small change that might improve the performance and efficiency of an average workflow that depends on string data immensely. These improvements are equally visible in pandas and Dask!

Engine keyword in I/O methods

We will now take a look at I/O functions in pandas and Dask. Some functions have custom implementations, like read_csv, while others dispatch to another library, like read_excel to openpyxl. Some of these functions gained a new engine keyword that enables us to dispatch to PyArrow. The PyArrow parsers are multithreaded by default and hence, can provide a significant performance improvement.

pd.read_csv("Crimes_-_2001_to_Present.csv", engine="pyarrow")

This configuration will return the same results as the other engines. The only difference is that PyArrow is used to read the data. The same option is available for read_json. The PyArrow-engines were added to provide a faster way of reading data. The improved speed is only one of the advantages. The PyArrow parsers return the data as a PyArrow Table. A PyArrow Table provides built-in functionality to convert to a pandas DataFrame. Depending on the data, this might require a copy while casting to NumPy (string, integers with missing values, …), which brings an unnecessary slowdown. This is where the PyArrow dtype_backend comes in. It is implemented as an ArrowExtensionArray class in pandas, which is backed by a PyArrow ChunkedArray. As a direct consequence, the conversion from a PyArrow Table to pandas is extremely cheap since it does not require any copies.

pd.read_csv("Crimes_-_2001_to_Present.csv", engine="pyarrow", dtype_backend="pyarrow")

This returns a DataFrame that is backed by PyArrow arrays. pandas isn’t optimized everywhere yet, so this can give you a slowdown in follow-up operations. It might be worth it if the workload is particularly I/O heavy. Let’s look at a direct comparison:

We can see that PyArrow-engine and PyArrow dtypes provide a 15x speedup compared to the C-engine.

The same advantages apply to Dask. Dask wraps the pandas csv reader and hence, gets the same features for free.

The comparison for Dask is a bit more complicated. Firstly, my example reads the data from my local machine while our Dask examples will read the data from a S3 bucket. Network speed will be a relevant component. Also, distributed computations have some overhead that we have to account for.

We are purely looking for speed here, so we will read some timeseries data from a public S3 bucket.

import dask.dataframe as dd
from distributed import wait

df = dd.read_csv(
    "s3://coiled-datasets/timeseries/20-years/csv/",
    storage_options={"anon": True},
    engine="pyarrow",
    parse_dates=["timestamp"],
)
df = df.persist()
wait(df)

We will execute this code-snippet for engine="c", engine="pyarrow" and additionally engine="pyarrow" with dtype_backend="pyarrow". Let’s look at some performance comparisons. Both examples were executed with 30 machines on the cluster.

The PyArrow-engine runs around 2 times as fast as the C-engine. Both implementations used the same number of machines. The memory usage was reduced by 50% with the PyArrow dtype_backend. The same reduction is available if only object columns are converted to PyArrow strings, which gives a better experience in follow-up operations.

We’ve seen that the Arrow-engines provide significant speedups over the custom C implementations. They don’t support all features of the custom implementations yet, but if your use-case is compatible with the supported options, you should get a significant speedup for free.

The case with the PyArrow dtype_backend is a bit more complicated. Not all areas of the API are optimized yet. If you spend a lot of time processing your data outside I/O functions, then this might not give you what you need. It will speed up your processing if your workflow spends a lot of time reading the data.

dtype_backend in PyArrow-native I/O readers

Some other I/O methods have an engine keyword as well. read_parquet is the most popular example. The situation is a bit different here though. These I/O methods were already using the PyArrow engine by default. So the parsing is as efficient as possible. One other potential performance benefit is the usage of the dtype_backend keyword. Normally, PyArrow will return a PyArrow table which is then converted to a pandas DataFrame. The PyArrow dtypes are converted to their NumPy equivalent. Setting dtype_backend="pyarrow" avoids this conversion. This gives a decent performance improvement and saves a lot of memory.

Let’s look at one pandas performance comparison. We read the Uber-Lyft taxi data from December 2022.

pd.read_parquet("fhvhv_tripdata_2022-10.parquet")

We read the data with and without dtype_backend="pyarrow".

We can easily see that the most time is taken up by the conversion after the reading of the Parquet file was finished. The function runs 3 times as fast when avoiding the conversion to NumPy dtypes.

Dask has a specialized implementation for read_parquet that has some advantages tailored to distributed workloads compared to the pandas implementation. The common denominator is that both functions dispatch to PyArrow to read the parquet file. Both have in common that the data are converted to NumPy dtypes after successfully reading the file. We are reading the whole Uber-Lyft dataset, which consumes around 240GB of memory on our cluster.

import dask.dataframe as dd
from distributed import wait

df = dd.read_parquet(
    "s3://coiled-datasets/uber-lyft-tlc/",
    storage_options={"anon": True},
)
df = df.persist()
wait(df)

We read the dataset in 3 different configurations. First with the default NumPy dtypes, then with the PyArrow string option turned on:

dask.config.set({"dataframe.convert-string": True})

And lastly with dtype_backend="pyarrow". Let’s look at what this means performance-wise:

Similar to our pandas example, we can see that converting to NumPy dtypes takes up a huge chunk of our runtime. The PyArrow dtypes give us a nice performance improvement. Both PyArrow configurations use half of the memory that the NumPy dtypes are using.

PyArrow-strings are a lot more mature than the general PyArrow dtype_backend. Based on the performance chart we got, we get roughly the same performance improvement when using PyArrow strings and NumPy dtypes for all other dtypes. If a workflow does not work well enough on PyArrow dtypes yet, I’d recommend enabling PyArrow strings only.

Conclusion

We have seen how we can leverage PyArrow in pandas in Dask right now. PyArrow backed string columns have the potential to impact most workflows in a positive way and provide a smooth user experience with pandas 2.0. Dask has a convenient option to globally avoid NumPy object dtype when possible, which makes opting into PyArrow backed strings even easier. PyArrow also provides huge speedups in other areas where available. The PyArrow dtype_backend is still pretty new and has the potential to cut I/O times significantly right now. It is certainly worth exploring whether it can solve performance bottlenecks. There is a lot of work going on to improve support for general PyArrow dtypes with the potential to speed up an average workflow in the near future.

There is a current proposal in pandas to start inferring strings as PyArrow backed strings by default starting from pandas 3.0. Additionally, it includes many more areas where leaning more onto PyArrow makes a lot of sense (e.g. Decimals, structured data, …). You can read up on the proposal here.

Thank you for reading. Feel free to reach out (patrick@coiled.io) to share your thoughts and feedback about PyArrow support in both libraries.

Distributed printing

2023-05-18T00:00:00+00:00

Dask makes it easy to print whether you’re running code locally on your laptop, or remotely on a cluster in the cloud.

One of the most basic things programmers do is print text to their screen.

Printing is often used for things like debugging:

# ...
print("Made it here...")
# ...

or to signal progress:

for i in range(10):
    print(f"Done with iteration {i}")

However, when running code at scale on a Dask cluster even simple print calls can become non-intuitive.

For example, this snippet runs two functions on a Dask/Coiled cluster in the cloud:

import coiled
import dask

# Spin up cluster
cluster = coiled.Cluster()
client = cluster.get_client()

@dask.delayed
def inc(x):
    print(f"Calling inc on {x}")
    return x + 1

@dask.delayed
def add(x, y):
    print(f"Calling add on {x} + {y}")
    return x + y

# Run computations on the cluster
x = inc(1)
y = add(x, 7)
y.compute()

Both the inc and add functions have print calls and, naively, I’d expect to see

Calling inc on 1
Calling add on 2 + 7

printed to my screen when I run this example. However, when you run this code snippet you don’t see any output. That’s because those prints happen on the remote Dask workers (in this case on AWS) where these tasks are run. We can see the result of these prints by looking our cluster logs.

It’s good that these prints work, but it’s not a very pleasant experience and probably doesn’t match most user’s expectations.

To make printing on a cluster similar to on programmer’s local machine, Dask has a dask.distributed.print function which mirrors Python’s built-in print function but, when called in a task running on a cluster, the output of the print will be forwarded back to the user’s client Python session and printed there.

Here’s the same example as above, but with one line added to use Dask’s print:

import coiled
import dask
from dask.distributed import print  # This line is the only difference

# Spin up cluster
cluster = coiled.Cluster()
client = cluster.get_client()

@dask.delayed
def inc(x):
    print(f"Calling inc on {x}")
    return x + 1

@dask.delayed
def add(x, y):
    print(f"Calling add on {x} + {y}")
    return x + y

# Run computations on the cluster
x = inc(1)
y = add(x, 7)
y.compute()

Running this on my laptop prints

Calling inc on 1
Calling add on 2 + 7

in my IPython session 🎉

When using prints inside tasks running on a Dask cluster, know that dask.distributed.print is available and probably what you want to use.

Tip

Dask also has a dask.distributed.warn utility, which mirrors Python’s built-in warnings.warn that forwards warnings to user’s client Python session.

Observability for Distributed Computing with Dask

2023-05-16T00:00:00+00:00

May 16, 2023

8m read

Hendrik Makait

Debugging is hard. Distributed debugging is hell.

When dealing with unexpected issues in a distributed system, you need to understand what and why it happened, how interactions between individual pieces contributed to the problems, and how to avoid them in the future. In other words, you need observability. This article explains what observability is, how Dask implements it, what pain points remain, and how Coiled helps you overcome these.

The Coiled metrics dashboard provides observability into a Dask cluster and its workloads.

This is based on a talk given by Hendrik Makait at PyCon DE & PyData Berlin 2023. You can also take a look at the original slides or watch the video recording.

The problem

Imagine receiving an alert: One of your daily computations failed after running for several thousands of CPU hours. By looking at the error message, you see that workers have died while attempting to run a particular task:

KilledWorker: Attempted to run task ('complex-computation-f970d1bdb34ba1bb5263425c7f356cf3', 0, 6) on 3 different workers, but all those workers died while running it. The last worker that attempt to run the task was tls://10.0.46.23:33455. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see https://distributed.dask.org/en/stable/killed.html.

“Weird,” you think, “this has never happened since I added it to production two months ago. So, why would it fail now? What has changed?”

Wouldn’t it be great if you had a way to answer these questions? Wouldn’t it be great if you could understand what had been going on (and going wrong!) with your Dask cluster when it failed? In other words, wouldn’t it be great if you had an observable system?

What is observability?

One particularly useful explanation of observability is provided by the OpenTelemetry project (emphasis added):

Observability lets us understand a system from the outside, by letting us ask questions about that system without knowing its inner workings. Furthermore, it allows us to easily troubleshoot and handle novel problems (i.e. “unknown unknowns”), and helps us answer the question, “Why is this happening?”

In essence, observability is the property of a system that allows users to understand why things are happening. By giving users a way to understand its behavior, an observable system enables users to take on different activities like monitoring the system’s health, root-causing issues, and optimizing performance.

To make a system observable, we need to collect data to describe the system’s behavior and performance. For Dask, two types of this telemetry data are the most useful:

Logs: A log is an (un-)structured text record about events emitted by system, application, or user code. Usually, logs are enriched with metadata, such as associated timestamps that provide the necessary context.

Metrics: A metric is a quantitative measurement that captures information about the behavior or performance of a system and its components. Metrics are collected at runtime and often presented as statistical aggregates and visualizations.

Having good telemetry data is necessary for a system to become observable. However, it is not enough: Users need diagnostic tools to leverage the telemetry data provided by the system and gain insights.

Observability with Dask

Logging

Dask and its components implement logging using Python’s standard logging module, which makes logging in Dask configurable. The difficulty with logging in Dask stems from its distributed nature:

How do you access logs in a cluster of hundreds of instances?
How do you correlate them to identify interactions?

While Dask contains the built-in methods Client.get_scheduler_logs() and Client.get_worker_logs() that fetch logs to the client, this approach has a major downside: If it’s dead, it doesn’t return logs. If you, like most developers, primarily use logs for debugging, the relevant instances are likely dead or struggling to answer requests, not to mention shutdown clusters when running a retroactive analysis.

The better solution is to rely on a centralized log management system to handle logs. This external system collects, stores, and analyzes logs from different sources. This approach has several benefits:

Persistent logs even after cluster shutdown
Simple log correlation between OS- and application-level logs across instances
Efficient and powerful querying

The snippet below shows how the Coiled CLI lets you filter for the cluster from our example and find all messages that mention the killed worker’s address:

$ coiled cluster logs --cluster 207560 --filter "10.0.46.23"
...
(10.0.46.23)    2023-05-04 15:01:25.668000 distributed.nanny.memory - WARNING - Worker tls://10.0.46.23:33455 (pid=52) exceeded 95% memory budget. Restarting...
(scheduler)     2023-05-04 15:01:25.697000 distributed.core - INFO - Connection to tls://10.0.46.23:43506 has been closed.
(scheduler)     2023-05-04 15:01:25.697000 distributed.scheduler - INFO - Remove worker <WorkerState 'tls://10.0.46.23:33455', name: hendrik-pycon-blog-worker-6660457730, status: paused, memory: 0, processing: 2>
(scheduler)     2023-05-04 15:01:25.697000 distributed.core - INFO - Removing comms to tls://10.0.46.23:33455
(10.0.43.162)   2023-05-04 15:01:25.703000 distributed.core - INFO - Removing comms to tls://10.0.46.23:33455
(10.0.35.110)   2023-05-04 15:01:25.704000 distributed.core - INFO - Removing comms to tls://10.0.46.23:33455
(10.0.33.60)    2023-05-04 15:01:25.705000 distributed.core - INFO - Removing comms to tls://10.0.46.23:33455
(10.0.36.173)   2023-05-04 15:01:25.706000 distributed.core - INFO - Removing comms to tls://10.0.46.23:33455
(10.0.35.15)    2023-05-04 15:01:25.707000 distributed.core - INFO - Removing comms to tls://10.0.46.23:33455
(10.0.46.123)   2023-05-04 15:01:25.714000 distributed.core - INFO - Removing comms to tls://10.0.46.23:33455
(10.0.42.85)    2023-05-04 15:01:25.714000 distributed.core - INFO - Removing comms to tls://10.0.46.23:33455
(10.0.39.241)   2023-05-04 15:01:25.721000 distributed.core - INFO - Removing comms to tls://10.0.46.23:33455
...

While logs contain rich information, sifting through them requires much work. Often, we want to gain an initial understanding of the state of our system at a glance. For this, metrics are more suitable.

Metrics

Metrics are what powers Dask. Internally, components like the scheduler or active memory manager rely on metrics to make data-driven decisions. To help users understand these decisions and observe what caused them, Dask provides two ways for users to access metrics:

Built-in dashboards
Integration with Prometheus

The integration with Prometheus adds an /metrics endpoint on the scheduler and workers. Prometheus can then periodically scrape these endpoints to collect the metrics. As with central log management, this has the benefit that Prometheus centrally collects, stores, and analyzes the metrics. You can then leverage Prometheus and build powerful dashboards for monitoring or drill-down analysis on top, as we did for our benchmark suite on our public Grafana server.

You can also configure rule-based alerts, e.g., if disk usage on a machine exceeds 95%. At Coiled, we use metrics to flag clusters with potentially performance-impacting behavior, such as high memory pressure on individual machines.

This Coiled cluster was flagged for cluster-wide memory pressure and disk usage.

The metrics exposed to Prometheus range from a global view of the system state

dask_scheduler_workers: Number of workers known by scheduler
dask_worker_threads: Number of worker threads

to a fine-grained performance focus

dask_scheduler_tasks_compute_seconds_total: Total time (per prefix) spent computing tasks
dask_worker_gil_contention_total: Value representing cumulative total GIL contention on worker, in the form of summed percentages.

You can find the complete list of supported metrics in the Dask docs.

To give you a better understanding of how you can use metrics to understand your system and make data-driven decisions, let’s briefly discuss some examples:

In Burstable vs. non-burstable AWS instance types for data engineering workloads, we look at runtime and other performance metrics to understand the practical difference between burstable and non-burstable instance types.

t3.large

m6i.large

As you can see in the plots above, it matters quite a bit, and we were able to explain the differences by looking at even more metrics.

In Performance testing at Coiled, we describe how we leverage performance metrics on an extensive suite of real-world use cases. Our approach allows us both to proactively gauge the performance impact of new changes to Dask through A/B testing as well as to identify and to root-cause performance regressions.

Diagnostic tools

Telemetry data like logs and metrics are only helpful if you can leverage them. To do so, you need good diagnostic tools that refine the available telemetry data into a format that helps the user understand what is happening inside the system and why.

We have already seen that centralized log management systems and metrics stores can be powerful diagnostic tools but require additional infrastructure work to set them up properly. Let us now look at the tools Dask has already built in.

Dashboards

The dashboards Dask offers provide a wealth of information. For example, the /status dashboard gives an overview of the cluster activity, including worker CPU and memory utilization as well as which tasks are being processed:

Other dashboards focus on details about individual aspects of the Dask cluster, e.g., providing more details on worker activity (/workers) or GIL contention (/individual-contention). See this walkthrough video for a tour of the available dashboards and how to leverage them.

When looking at how useful dashboards are for observability, the results are mixed:

👍 Powerful: Each dashboard is designed with a specific question in mind that it helps to answer.
👍 Pre-configured: Dashboards are pre-configured and require no additional setup.
👎 Ephemeral: Dashboards are only available as long as the scheduler/worker is alive, and they have no mechanism for playing back data from the past.

Performance reports

Performance reports are a built-in tool allowing users to store and share a subset of the information available in the dashboards. To create a performance report, users have to wrap their workload with the distributed.performance_report context manager:

from distributed import performance_report

with performance_report(filename="pycon_de.html"):
    ...

This will collect all the telemetry and compile the report into a shareable HTML file. For more, watch this 3-minute demo.

Profiling

When analyzing performance, you won’t get around profiling. Dask has a built-in statistical CPU profiler that you can access at /profile. This profiler has low overhead and is always running. It can also aggregate profiling information across the cluster or show individual profiles for each worker. However, the profiler is Python-only, so workloads that heavily use compiled packages will benefit less.

If you are interested in more detailed profiling of your Dask workloads or are heavily using compiled packages, I’d suggest taking a look at dask-pyspy. At a Dask Demo Day, Gabe Joseph gave an excellent demo.

Zooming back out

To summarize, I’d like to highlight a few key points:

An observable system lets us ask: “Why is this happening?” and answer questions about unknown unknowns.
Telemetry data such as logging, metrics, and tracing are necessary but not sufficient for an observable system.
Good diagnostic tooling is also needed to give users insights into the system and its components.
Dask offers powerful but ephemeral telemetry data and diagnostic tools.
To fully leverage the telemetry data available in Dask, you need to centrally and durably collect them.

If having good observability into Dask excites you, I’d encourage you to spin up your next Dask cluster on Coiled; we manage a centralized observability platform for you.

GIL monitoring in Dask

2023-05-15T00:00:00+00:00

May 15, 2023

2m read

Miles Granger

Added in version 2023.4.1: Support GIL contention monitoring.

With release 2023.4.1, dask.distributed adds support for monitoring Python’s Global Interpreter Lock (GIL).

New GIL “Contention” monitoring dashboard at <<dashboard host>>/individual-contention

Why should I care?

Well, if you are running a single-threaded Python application, odds are you don’t need to care much. In the context of Dask though, it can potentially make a significant difference in performance.

Code running on Dask runs in a threaded environment with other threads doing compression, disk I/O, network I/O, and other user code. Hopefully this code uses the GIL in a judicious manner.

GIL monitoring through the dashboard and Prometheus metrics enables users to identify workflows that can be improved to be more GIL friendly.

Alright fancy pants, how do I activate this feature?

Just have gilknocker installed. 🚀

By default the dask.config has distributed.admin.system-monitor.gil.enabled=True.

Tip

It’s likely this will cover your needs, however one can additionally adjust distributed.admin.system-monitor.gil.interval. At the time of this writing, it defaults to 1 ms but can be changed to be a higher or lower frequency. This will affect the frequency in which the GIL is polled to determine how busy it is.

Example

Here is an example adjusting the frequency for GIL polling to 10us.

distributed:
  admin:
    system-monitor:
      gil:
        enabled: true
        interval: 10us

Is there a performance impact by enabling this feature?

It depends. 😉

With that said, it’s quite likely you’ll experience no noticeable impact on performance using the default settings. Should you increase the polling interval to a higher frequency, then this will contribute to more GIL contention and potentially lead to decreased performance.

Let’s look at a comparison we ran before adding this feature, with GIL monitoring enabled (left) and disabled (right).

What am I looking at?

In an effort to preserve stable performance, we use Coiled’s A/B benchmarks to show the negligible performance impact of GIL monitoring.

Above we have a comparison of confidence intervals against the null-hypothesis, that is to say, not monitoring the GIL.

Each row represents a single benchmark test, the blue meaning improvement and red a regression. There is always variance to some extent, represented by the spread horizontally and shading. However, the more similar the two charts look, the less likely there is a performance difference.

Summary

With GIL monitoring you’ll be able to view how your code behaves in this context and enable you to make more informed decisions on performance improvements.

To enable monitoring the GIL, you only need to install gilknocker which is available on conda-forge and PyPI.

While these metrics are available with Prometheus, you’ll also get an updated dashboard called “Contention” located at <<dashboard host>>/individual-contention displaying the Event Loop and GIL contention metrics in realtime.

Performance testing at Coiled

2023-05-05T00:00:00+00:00

At Coiled we develop Dask and automatically deploy it to large clusters of cloud workers (sometimes 1000+ EC2 instances at once!). In order to avoid surprises when we publish a new release, Dask needs to be covered by a comprehensive battery of tests — both for functionality and performance.

Functional regression testing

Dask is thoroughly covered by functional unit tests which verify that, when you feed certain inputs into the system, you obtain the expected outputs. Similarly to most software development projects nowadays, we use a CI-assisted, automated workflow for it:

A developer opens a PR;
The PR code is automatically executed in CI against all functional regression tests;
If, and only if, all tests are green the PR can be merged into main.

The key features of this system are:

It re-runs on every change;
Nobody needs to explicitly remember to start the tests;
It runs in a stable, reproducible environment;
It re-tests everything, not just what the developers think may be impacted.

None of the above should surprise anyone these days. What we do, that is definitely more out of the ordinary, is we apply the exact same paradigms commonly used for functional regression testing also to performance testing.

What we did in the past

Historically, whenever we had to decide whether to merge a Dask PR with a likely performance impact, we ran ad hoc manual tests on one or two use cases that we thought would be relevant, which we found was not sustainable:

It was labor-intensive, as the developers would need to rebuild the use case and manually meter it every time;
PRs frequently impacted unexpected use cases: a PR that would produce a speed-up where expected sometimes caused a regression where nobody was looking;
A cluster with hundreds of workers very rarely exhibits deterministic performance; if you run the same test twice there will be some variance in the measured metrics. One-off manual tests tend to ignore variance and instead take the first measure as golden.

Clearly something had to change.

What about ASV?

Airspeed Velocity (ASV) is a great microbenchmarking tool. Emphasis however should be put on the micro bit in the name, as there are important limitations:

Everything must run on a single box;
You must trust performance to scale linearly;
You are only interested in end-to-end runtime.

When testing Dask:

Performance on a LocalCluster is completely different from that on AWS. Network latency and bandwidth play a major role, as do other bottlenecks such as I/O throughput. You can’t expect the behaviour of a single box to be representative of a cloud cluster, even with the same amount of CPUs and RAM.
You can’t expect a scaled down problem running on 4 CPUs to be representative of a production size problem running on 400 CPUs. Again, bottlenecks such as the scheduler, S3 throughput, and network bandwidth play a fundamental part in this.
There are many metrics other than end-to-end runtime such as memory and network usage, as well as the breakdown of the totals.

coiled/benchmarks

We created a GitHub repository and wrote a comprehensive suite of tests in it. Each test spawns a medium-sized Coiled cluster and runs a realistic use case, such as:

Each test case is measured on selected performance metrics; at the moment, end-to-end runtime, peak memory usage, and mean memory usage; we plan to add more in the future.

The test suite runs in two modes: nightly regression tests and A/B tests.

Nightly tests

Every night, CI will automatically:

Run all tests;
Detect regressions in test runtimes;
Open a GitHub issue if performance degrades from previous days;
Publish a report showing several months worth of results.

Case study

The test shown in the above image exhibited a very consistent end-to-end runtime throughout the months. Then, one fateful day in early March, somebody merged a PR into main that caused this test to slow down by a factor of 4x!

The next morning Dask maintainers woke up to a new GitHub issue highlighting the regression. The issue was promptly investigated and remediated shortly thereafter.

A/B tests

Before we merge major PRs we push to the coiled/benchmarks CI two or more conda environments:

A (a.k.a. baseline; what we know and love)

depedencies:
  ...
  - dask ==2023.4.1
  - distributed ==2023.4.1

B (a PR, the next release, whatever)

dependencies:
  ... # Everything same as baseline
  - pip:
    - git+https://github.com/dask/dask@b85bf5be72b02342222c8a0452596539fce19bce
    - git+https://github.com/dask/distributed@6dc4c24b7a01b93952bbfd95d3ecba043580c245

The conda environment files can be accompanied by Dask config files and/or parameters to the Coiled cluster; this allows, for example, to test how switching on/off work stealing affects performance or how 10 workers with 4 CPUs each compare against 20 workers with 2 CPUs each.

You can also increase cluster size on the fly; most test data scales dynamically with cluster size so that end-to-end runtime remains roughly constant:

# For all tests using the small_client fixture
small_cluster:
  n_workers: 10
  worker_vm_types: [m6i.large]  # 2 CPU, 8 GiB

def test_dot_product(small_client):
    memory = cluster_memory(small_client)  # Typically 76.66 GiB
    shape = scaled_array_shape_quadratic(memory // 17, "4.5 GiB", ("x", "x"))
    a = da.random.random(shape)
    b = (a @ a.T).sum()
    b.compute()

Once the test settings are uploaded CI will run all tests, multiple times, in the A and B configurations and will generate a statistical analysis of the differences.

Confidence report

This A/B report shows, for every test case, how likely it is that B is better/worse than A in end-to-end runtime, memory usage, etc. This is an aggregation of 5+ runs for B vs. 5+ runs for A; the report shows p values as transparency:

Case study

The above plot is the output of a preliminary performance test for a PR introducing Asynchronous disk access. The PR caused test_set_index (4^th row from the top) to become at least 40% faster (p=95%) and test_spilling (middle of the plot) to become at least 60% slower (p=84%). Both results were completely unexpected; the PR did not get merged.

Null hypothesis

Some tests are noisier than others. CI automatically runs tests for the baseline environment twice and compares them in A/B. If the null hypothesis shows changes, then information from that test should be treated as noise in the actual A/B tests.

Case study

In the above plot, test_single_future seems to be sometimes a bit faster, sometimes a bit slower in the B setup compared to A. Except that A and B are identical. This is because the test is so noisy that, out of sheer luck, the runs in B were faster on average than those in A. The test should be ignored when looking at actual A/B situations. If you want meaningful results out of it you should either run it many more times or redesign it to be less noisy.

Per test details

For each test we have a more detailed plot showing median (blue bar), mean (crosshair), and standard deviation (black horizontal line) of the observed metrics:

Case study

This plot is from the same A/B test run that produced the confidence reports shown before (the relevant bit is repeated right below for convenience). The runtime in the two B setups (async_spill and async_spill_nohack) is solidly less than that in the baseline and null hypothesis.

Case study

In this case, instead, you can clearly see that the noise-to-signal ratio is so high that there’s nothing useful to be gathered out of this test.

Introspection with Grafana

Finally, from the aggregated statistical reports, using Prometheus and Grafana we can drill down into each individual run that underlies them and observe what happened during it. This gives us information such as the runtime of each task group, the breakdown of memory usage, network throughput, etc.

FAQ

Q: Isn’t this very expensive?

A: Around $5 for the nightly tests and around $50 for a typical A/B test. The cost increases if you decide to increase the cluster size to hundreds of workers, but it’s not something needed every day.

Q: Could you simulate a slow network on localhost instead?

A: In theory, yes. In practice, we’d periodically need to make sure that the simulation behaves like the real thing. We’d also need to simulate actual task execution, as a single host would not be able to cope with the CPU and memory usage of a real cluster. We’d need to simulate how S3 behaves when multiple hosts request different files from it. There are many other pitfalls to consider.

Not saying that simulation is always a bad idea, just that it takes a substantial investment of time and effort to set up correctly and to ensure it keeps reflecting reality. Does each test run cost you $10 billion? Maybe you should figure out how to simulate it instead. Does it cost $50? Just not worth it.

What’s next?

The current toolkit is extremely powerful and light years ahead of the previous condition of testing things by hand. However, there are still large margins for improvement:

Make it more user-friendly;
Make it more interactive. For example, if you see an increase in end-to-end runtime, you probably want to know where the extra time is spent. Dask is already producing this information; we’re working on properly storing, aggregating, and visualizing it;
Observe more metrics such as disk and network usage.

Finally, we’re exploring the possibility of opening this system to Coiled customers! For example, imagine if you could tag some of your runs as “A”, some others as “B”, and then ask the Coiled dashboard to show you statistical insights on the differences between the two sets. You would be able to test not only different versions of Dask and other dependencies, but also different versions of your own algorithm, different cluster configurations, etc.

How well does Dask run on Graviton?

2023-05-05T00:00:00+00:00

May 5, 2023

4m read

Sarah Johnson, Nat Tabris

ARM-based processors are known for matching performance of x86-based instance types at a lower cost, since they consume far less energy for the same performance. It’s not surprising then that some companies, like Honeycomb, are switching their entire infrastructure to ARM.

We ran a number of Dask workloads on both x86- and ARM-based instance types and found costs were typically 20-30% lower when using ARM. We also tested out the latest generation of Amazon’s ARM-based processors Graviton3 instance types and looked at performance for a compute-heavy workload using Dask and XGBoost.

AWS Graviton

Intel designs all Intel chips (including Intel x86 processors, e.g.), but many different companies design ARM chips. This lets different companies tailor the design to specific use cases. Apple, for example, designs ARM processors that deliver a nice balance of performance and energy-efficiency for phones and laptops. Amazon has designed Graviton processors specifically for their cloud servers, and the latest generation Graviton3 is especially optimized for compute-intensive workloads. We launched Dask clusters on AWS using Intel, Graviton2, and Graviton3 processors to see how switching to ARM makes a difference.

Dask on Graviton2

Previously, we used Dask and Coiled to process ~ 1TB of data to see how popular Matplotlib has been in scientific articles over the past decade. The dataset is publicly available in S3, saved as over 5,000 tar files, each of which contain a directory of articles. We extract each directory and do some minimal processing, which takes ~ 7 minutes with a 100-worker Coiled cluster.

We ran this workload 5 times on both x86-based t3 instances and ARM-based t4g instances and found, on average, a 30% cost reduction.

This workload doesn’t use much CPU or memory, and is instead primarily relying on network bandwidth. The general-purpose x86-based t3 and ARM-based t4g instance types are a good fit for this type of workload, since they are readily available and provide more than enough performance without additional charges.

To see how robust these results were across a variety of workloads, we also ran a number of benchmarks for common Dask Array and Dask DataFrame operations. Across 19 different benchmarks, the ARM-based t4g instances were reliably cheaper for all but one, where there was a 2.3% cost increase. Cost savings were around 30% averaging across all tests, ranging from 15-55%.

Percent difference in estimated AWS costs for t4g (ARM-based) vs. t3 workers (x86-based)

On the high end, test_vorticity, for example, was 46% cheaper (on average across three trials). This test is a simplification of a Dask Array operation commonly used in geoscience workloads and is primarily compute-bound.

Dask on Graviton3

In February, AWS launched additional general-purpose Graviton3-based instance types, the m7g, with reportedly up to 25% better performance than their predecessors. For more details on Graviton3 performance improvements, we recommend this deep dive from Chips and Cheese.

We used XGBoost with Dask and Coiled to predict the trip duration for Uber and Lyft rides in NYC using the publicly available NYC Taxi dataset. In this workload we load the dataset, do some basic feature engineering, train and then cross-validate our XGBoost model. This takes ~5 minutes on a 45-worker Coiled cluster with 3 folds for cross-validation.

We ran this workload 5 times on both x86-based m6i instances and ARM-based m7g instances and found, on average, a 23% cost savings.

As opposed to the network-bound Matplotlib example, this workload is primarily computationally- and memory-bound, with total worker memory reaching just under 400 GB and workers reaching nearly 75% CPU utilization. For these types of workloads, instance type families that balance compute and memory are a good option.

Average worker percent CPU utilization (top) and average worker memory (bottom) for XGBoost workload

Again, to get a sense for how robust these cost savings are, we ran a the same benchmarks as above comparing m6i and m7g instance types. Across the 19 different benchmarks, the ARM-based m7g instances were reliably cheaper for all but three, where ARM was, at most 25% more expensive. Where ARM was cheaper, estimated cost savings ranged from 5-55%, with an average of ~ 20% cost savings across tests.

Percent difference in estimated AWS costs for m7g (ARM-based) vs. m6i workers (x86-based)

Focusing again at the primarily compute-bound test_vorticity, m7g workers were 48% cheaper than using a cluster with m6i workers.

Using ARM with Coiled

Graviton processors are designed by Amazon specifically with cloud performance in mind and they consume far less energy. So why not try them for your own workloads? One challenge is ensuring your compiled dependencies are ARM-compatible. In many cases, the dependencies you’ll need are prebuilt for both x86 and ARM. If you’re using custom dependencies or an older version of a library, however, you may run into issues.

As of Coiled version 0.6.0, your local Python environment will be replicated on your cluster, including any necessary platform-specific dependencies, making it easy to create a cluster with ARM-based instances.

cluster = coiled.Cluster(
    # ARM instance types for scheduler and workers
    arm=True,
    n_workers=10,
)

You can see for yourself whether switching to ARM-based clusters saves on AWS costs or improves performance (see our documentation for more details). Try it out and let us know how it goes!

If you’re not already using Coiled, it’s free to try out and easy to get started.

Upstream testing in Dask

2023-04-18T00:00:00+00:00

Dask has deep integrations with other libraries in the PyData ecosystem like NumPy, pandas, Zarr, PyArrow, and more. Part of providing a good experience for Dask users is making sure that Dask continues to work well with this community of libraries as they push out new releases. This post walks through how Dask maintainers proactively ensure Dask continuously works with its surrounding ecosystem.

Nightly testing

Dask has a dedicated CI build that runs Dask’s normal test suite once a day with unreleased, nightly versions of several libraries installed. This lets us check whether or not a recent change in a library like NumPy or pandas breaks some aspect of Dask’s functionality.

To increase visibility when such a breakage occurs, as part of the upstream CI build, an issue is automatically opened that provides a summary of what tests failed and links to the build logs for the corresponding failure (here’s an example issue).

This makes it less likely that a failing upstream build goes unnoticed.

How things can break and are fixed

There are usually two different ways in which things break. Either:

A library made an intentional change in behavior and a corresponding compatibility change needs to be made in Dask (the next section has an example of this case).
There was some unintentional consequence of a change made in a library that resulted in a breakage in Dask.

When the latter case occurs, Dask maintainers can then engage with other library maintainers to resolve the unintended breakage. This all happens before any libraries push out a new release, so no user code breaks.

Example: pandas 2.0

One specific example of this process in action is the recent pandas 2.0 release. This is a major version release and contains significant breaking changes like removing deprecated functionality.

As these breaking changes were merged into pandas, we started seeing related failures in Dask’s upstream CI build. Dask maintainers were then able to add a variety of compatibility changes so that Dask works well with pandas 2.0 immediately.

Acknowledgements

Special thanks to Justus Magin for his work on the xarray-contrib/issue-from-pytest-log GitHub action. We’ve found this to be really convenient for easily opening up GitHub issues when test failures occur.

Also, thanks to Irina Truong (Coiled), Patrick Hoefler (Coiled), and Matthew Roeschke (NVIDIA) for their efforts ensuring pandas and Dask continue to work together.

Burstable vs non-burstable AWS instance types for data engineering workloads

2023-04-04T00:00:00+00:00

Apr 4, 2023

8m read

Nat Tabris

There are many instance types to choose from on AWS. In this post, we’ll look at one choice you can make—burstable vs non-burstable instances—and show how the “cheaper” burstable option can end up being more expensive for data engineering workloads.

AWS offers burstable EC2 instances with a lower base price than the corresponding non-burstable instance types.

Burstable instances are perfect for workloads like web servers which have low CPU utilization most of the time, but occasionally fully utilize the CPU—web traffic tends to be bursty.

What about data science/engineering workloads? For these use-cases, burstable instances are worse than we expected. For a wide-range of data workloads, non-burstable instances are both cheaper per hour and faster.

Coiled offers managed Dask in the cloud. Based on what we’ve learned about burstable instances, we’re changing the default Intel-based instance types on AWS from t3 to m6i instances. This is one of the many ways that we help users run Python at scale efficiently (as well as cheaply and securely).

In the rest of the post, we’ll see the impact on overall costs across our benchmark suite, and then dig into a few workloads to see how non-burstable instances provide better performance for CPU, disk, and network.

Overall Benchmark Costs

Coiled runs a suite of Python/Dask benchmarks on AWS. This collection of ~100 benchmarks cover common data science, data engineering, and machine learning workloads. A typical run takes around 150 instance-hours (though as we’ll see, this depends on the choice of instance type).

When we switched from burstable to non-burstable instance types, we saw a significant reduction in cost. For instance, comparing two individual runs, we saw costs go from $7.13 (burstable) to $3.89 (non-burstable).

This cost difference is due to two things:

The burstable surcharge
Poor performance of burstable instances

We’ll go into more detail on each of these in the sections below.

Burstable surcharge

If you take a look at the prices that AWS lists for EC2 instances (here’s the on-demand pricing), you’ll see that t3.large instances are cheaper than m6i.large instances. So why is the total cost for m6i.large so much lower?

With burstable instances, the base price covers CPU utilization up to some limit: 30% for t3.large, 40% for t3.xlarge and t3.2xlarge. This base price is about 15% cheaper than the price of the corresponding non-burstable instances.

With burstable instances there are two types of costs you need to watch: the base price, and the extra CPU credit cost. What’s this extra cost about?

Burstable t3.large instances have a “baseline” of 30% CPU utilization. When your utilization is lower, you accumulate CPU credits. When your utilization is higher, you use the credits you’ve accumulated, or if you don’t have enough credits, you pay extra.

In the case of our benchmark suite, these extra CPU credits were 28% of the total non-burstable compute cost and made our burstable compute significantly more expensive than non-burstable compute.

You can see these costs for yourself on instances in your account as explained in our blog post about AWS Cost Explorer.

More generally, using m6i.large instances are cheaper than t3.large instances when you have sustained CPU utilization above 42.5%. See this AWS doc for details about how this is calculated.

What if I only want to pay the base price, nothing extra? If you never want to pay for extra CPU credits, you can use “standard” rather than “unlimited” mode with burstable instances. In this case, instead of paying for CPU credits, your performance is throttled. While this may make sense for bursty use-cases, this is a poor fit for compute-heavy workloads: you’d be paying 85% of the cost for 40% (or less) of the performance of a non-burstable instance.

Poor Performance of Burstable Instances

The benchmarks ran faster on m6i.large instances than on t3.large. At a high-level, we can see this by comparing the total number of instance-hours:

Let’s dig into why the non-burstable instances are faster than the burstable instances.

Modern Chipset

Part of the explanation is that m6i offers faster and newer Intel CPUs, as well as faster networking:

	Burstable (t3.large)	Non-burstable (m6i.large)
Processor	Up to 3.1 GHz Intel Xeon Scalable processor (Skylake 8175M or Cascade Lake 8259CL)	Up to 3.5 GHz 3rd Generation Intel Xeon Scalable processors (Ice Lake 8375C)
Network Performance	Up to 5 Gbps	Up to 10 Gbps

But that’s not the full story. To get a deeper understanding, let’s take a closer look at a few different workloads.

Workload 1: Network and disk with array workload

The following workload stresses network and disk. We construct a large random array (too large for memory) and then rechunk it twice across the network, forcing a full shuffle of all of the data:

x = da.random.random((50,000, 50,000))
x.rechunk((50,000, 20)).rechunk((20, 50,000)).sum().compute()

Coiled collects extensive metrics on workloads. We publicly host the metrics for the benchmark suite on a public Grafana server to aid community development. We can go and look at them there.

Let’s take a look at the network and disk metrics for the same tests on m6i.large and t3.large instances, moving 1.5TB over 25 minutes or 15 minutes respectively:

t3.large

m6i.large

The m6i.large instances ran the same workload in only 15 minutes, compared to 25 minutes on t3.large instances. This is explained by two facts:

First, peak speeds are higher on m6i.large instances:

Second, variance was higher across burstable instances. This is evident from the bands of color on the Grafana charts, which are narrower for the m6i.large run, but thicker for the t3.large run.

The m6i.large instances get faster network and disk speeds, and they get these speeds more consistently.

You can see these charts (and lots of others!) for yourself on our Grafana server:

Workload 2: XGBoost + Optuna with moderate CPU use

When we started to compare performance on burstable and non-burstable instances, one of the more surprising discoveries was that burstable instances showed much more “steal” CPU time. This came up in the XGBoost + Optuna benchmark:

def train_model(**study_params):
    model = xgboost.dask.train(
        None,
        {"tree_method": "hist", **study_params},
        d_train,
        num_boost_round=4,
        evals=[(d_train, "train")],
    )
    predictions = xgboost.dask.predict(None, model, x_test)
    return mean_squared_error(
        y_test.to_dask_array(),
        predictions.to_dask_array(),
        squared=False,
    )

def objective(trial):
    params = {
        "n_estimators": trial.suggest_int("n_estimators", 75, 125),
        "learning_rate": trial.suggest_float("learning_rate", 0.5, 0.7),
        "colsample_bytree": trial.suggest_float("colsample_bytree", 0.5, 1),
        "colsample_bynode": trial.suggest_float("colsample_bynode", 0.5, 1),
        "colsample_bylevel": trial.suggest_float("colsample_bylevel", 0.5, 1),
        "reg_lambda": trial.suggest_float("reg_lambda", 0, 1),
        "max_depth": trial.suggest_int("max_depth", 1, 6),
        "max_leaves": trial.suggest_int("max_leaves", 0, 2),
        "max_cat_to_onehot": trial.suggest_int("max_cat_to_onehot", 1, 10),
    }
    return train_model(**params)

study = optuna.create_study(study_name="nyc-travel-time-model")
study.optimize(objective, n_trials=3)

This benchmark uses little network or disk, and while it doesn’t saturate the CPU, it does have a consistent level of moderate CPU utilization. Here is the CPU utilization throughout the benchmark run:

t3.large (full Grafana dashboard)#

m6i.large (full Grafana dashboard)#

The m6i run is faster. Interestingly, there is also more “steal” (shown in red) on the t3 instances, while there’s almost none on the m6i instances. The non-burstable m6i instances have more of CPU use in user and system, where the t3.large instances had on average 5–12% steal.

What is steal CPU use you ask? According to the Linux man page for mpstat,

steal [is] the percentage of time spent in involuntary wait by the virtual CPU or CPUs while the hypervisor was servicing another virtual processor.

With a few exceptions, EC2 instances are run on much larger hosts. The physical server runs a hypervisor that virtualizes the resources and ensures that every EC2 instance gets its fair share of the actual hardware. Steal is the overhead of this sharing—it’s the amount of time your code spends waiting to get access to an actual CPU.

While you sometimes see steal on non-burstable instances, it’s very low. Burstable instances consistently show much higher steal. We haven’t found a detailed explanation from AWS but we suspect this has two possible causes:

Burstable instances are cheaper because AWS can “overbook” the hardware, much like an airline might overbook a flight. As long as enough people are underutilizing the CPU (or enough people are running late for the flight), everyone can get what they want. But sometimes more of the workloads want to utilize more of the CPU (or everyone wants the seat they paid for), and that’s when someone has to get bumped.
On burstable instances, the hypervisor itself also has more work to do, since it tracks CPU credits. While AWS doesn’t say this, it’s possible that tracking burst (and the associated costs) also adds overhead. Or maybe this isn’t the case, since the hypervisor on Nitro instances is actually offloaded to separate, dedicated hardware.

Regardless of the underlying cause, we’ve seen that burstable instances very frequently show 5–10% “steal” (or more), while non-burstable instances almost never have more than 1 or 2% steal.

The upshot is that for workloads with moderate CPU utilization—such as typical data engineering workloads—you’ll often end up paying more for burstable instances and getting worse performance.

Spot availability for burstable instances

Spot instances are another way to save money in the cloud. Spot instances are cheaper than on-demand instances: often providing 60–70% savings.

However, spot instances have more limited availability than on-demand instances. They’re offered at a discount because they’re unused capacity, and AWS can reclaim your spot instances if there’s more demand for full-priced instances.

It’s possible that even with all the caveats about burstable instances, burstable instances might be best all things considered if you had to choose between

t3 spot instances, and
m6i on-demand instances.

The key question here is whether it’s easier to get t3 spot instances than it is to get m6i spot instances.

Spot capacity varies over time by instance type and by region, so we can’t promise that m6i spot instances aren’t harder to get, but we can share what we see when running our benchmarks at scale.

With the “use best zone” feature (described in our blog post about spot instances), we’ve gotten 100% of the requested m6i.large workers in the last 14 days as spot, thousands of per day:

If you care about cost, you should definitely watch your ability to get spot vs on-demand instances (you can see your cost breakdowns as explained in this blog post). But in our experience, we’ve had no trouble getting non-burstable m6i.large spot instances.

Shuffling large data at constant memory in Dask

2023-03-15T00:00:00+00:00

Mar 15, 2023

5m read

Hendrik Makait

With release 2023.2.1, dask.dataframe introduces a new shuffling method called P2P, making sorts, merges, and joins faster and using constant memory. Benchmarks show impressive improvements:

P2P shuffling (blue) uses constant memory while task-based shuffling (orange) scales linearly.

This article describes the problem, the new solution, and the impact on performance.

What is Shuffling?

Shuffling is a key primitive in data processing systems. It is used whenever we move a dataset around in an all-to-all fashion, such as occurs in sorting, dataframe joins, or array rechunking. Shuffling is a challenging computation to run efficiently, with lots of small tasks sharding the data.

Task-based shuffling scales poorly

While systems like distributed databases and Apache Spark use a dedicated shuffle service to move data around the cluster, Dask builds on task-based scheduling. Task-based systems are more flexible for general-purpose parallel computing but less suitable for all-to-all shuffling. This results in three main issues:

Scheduler strain: The scheduler hangs due to the sheer amount of tasks required for shuffling.
Full dataset materialization: Workers materialize the entire dataset causing the cluster to run out of memory.
Many small operations: Intermediate tasks are too small and bring down hardware performance.

Together, these issues make large-scale shuffles inefficient, causing users to over-provision clusters. Fortunately, we can design a system to address these concerns. Early work on this started back in 2021 and has now matured into the P2P shuffling system:

P2P shuffling

With release 2023.2.1, Dask introduces a new shuffle method called P2P, making sorts, merges, and joins run faster and in constant memory.

This system is designed with three aspects, mirroring the problems listed above:

Peer-to-Peer communication: Reduce the involvement of the scheduler

Shuffling becomes an O(n) operation from the scheduler’s perspective, removing a key bottleneck.

Disk by default: Store data as it arrives on disk, efficiently.

Dask can now shuffle datasets of arbitrary size in small memory, reducing the need to fiddle with right-sizing clusters.
Buffer network and disk with memory: Avoid many small writes by buffering sensitive hardware with in-memory stores.

Shuffling involves CPU, network, and disk, each of which brings its own bottlenecks when dealing with many small pieces of data. We use memory judiciously to trade off between these bottlenecks, balancing network and disk I/O to maximize the overall throughput.

In addition to these three aspects, P2P shuffling implements many minor optimizations to improve performance, and it relies on pyarrow>=7.0 for efficient data handling and (de)serialization.

Results

To evaluate P2P shuffling, we benchmark against task-based shuffling on common workloads from our benchmark suite. For more information on this benchmark suite, see the GitHub repository or the latest test results.

Memory stability

The biggest benefit of P2P shuffling is constant memory usage. Memory usage drops and stays constant across all workloads:

For the tested workloads, we saw up to 10x lower memory. For even larger workloads, this gap only increases.

Performance and Speed

In the above plot, we can two performance improvements:

Faster execution: Workloads run up to 45% faster.
Quicker startup: Smaller graphs mean P2P shuffling starts sooner.

Comparing scheduler dashboards for P2P and task-based shuffling side-by-side helps to understand what causes these performance gains:

Task-based shuffling (left) shows bigger gaps in the task stream and spilling compared to P2P shuffling (right).

Graph size: 10x fewer tasks are faster to deploy and easier on the scheduler.
Controlled I/O: P2P shuffling handles I/O explicitly, which avoids less performant spilling by Dask.
Fewer interruptions: I/O is well balanced, so we see fewer gaps in work in the task stream.

Overall, P2P shuffling leverages our hardware far more effectively.

Changing defaults

These results benefit most users, so P2P shuffling is now the default starting from release 2023.2.1, as long as pyarrow>=7.0.0 is installed. Dataframe operations like the following will benefit:

df.set_index(...)
df.merge(...)
df.groupby(...).apply(...)

Keep old behavior

If P2P shuffling does not work for you, you can deactivate it by setting the dataframe.shuffle.method config value to "tasks" or explicitly setting a keyword argument, for example:

with the yaml config

dataframe:
  shuffle:
    method: tasks

or when starting a Coiled cluster

import dask
import coiled

# The dataframe.shuffle.method config is available since 2023.3.1
with dask.config.set({"dataframe.shuffle.method": "tasks"}):
    cluster = coiled.Cluster(...)  # coiled sends current dask config automatically

For more information on deactivating P2P shuffle, see the discussion #7509.

What about arrays?

While the original motivation was to optimize large-scale dataframe joins, the P2P system is useful for all problems requiring lots of communication between tasks. For array workloads, this often occurs when rechunking data, such as when organizing a matrix by rows when it was stored by columns. Similar to dataframe joins, array rechunking has been inefficient in the past, which has become such a problem that the array community built specialized tools like rechunker to avoid it entirely.

There is a naive implementation of array rechunking using the P2P system available for experimental use. Benchmarking this implementation shows mixed results:

👍 Constant memory use: As with dataframe operations, memory use is constant.
❓ Variable runtime: The runtime of workloads may increase with P2P.
👎 Memory overhead: There is a large memory overhead for many small partitions.

The constant memory use is a very promising result. There are several ways to tackle the current limitations. We expect this to improve as we work with collaborators from the array computing community.

Next steps

Development on P2P shuffling is not done yet. For the future we plan the following:

dask.array:

While the early prototype of array rechunking is promising, it’s not there yet. We plan to do the following:
- Intelligently select which algorithm to use (task-based rechunking is better sometimes)
- Work with collaborators on rechunking to improve performance
- Seek out other use cases, like map_overlap, where this might be helpful
Failure recovery:

Make P2P shuffling resilient to worker loss; currently, it has to restart entirely.
Performance tuning:

Performance today is good, but not yet at peak hardware speeds. We can improve this in a few ways:
- Hiding disk I/O
- Using more memory when appropriate for smaller shuffles
- Improved batching of small operations
- More efficient serialization

To follow along with the development, subscribe to the tracking issue on GitHub.

Summary

Shuffling data is a common operation in dataframe workloads. Since 2023.2.1, Dask defaults to P2P shuffling for distributed clusters and shuffles data faster and at constant memory. This improvement unlocks previously un-runnable workload sizes and efficiently uses your clusters. Finally, P2P shuffling demonstrates extending Dask to add new paradigms while leveraging the foundations of its distributed engine.

Share results in this discussion thread or follow development at this tracking issue.

Just in time Python environments

2023-02-23T00:00:00+00:00

Feb 23, 2023

4m read

Samantha Hughes, Sarah Johnson

Docker is a great tool for creating portable software environments, but we found it’s too slow for interactive exploration. We find that clusters depending on docker images often take 5+ minutes to launch. Ouch.

In the latest Coiled release, version 0.4.0, you can use a new system for creating software environments on the fly using only mamba instead. We’re seeing start times 3x faster, or about 1-2 minutes.

Why we dropped Docker for Python environments

Docker is a great tool for managing software environments, but we found that it’s just too slow, especially for exploratory data workflows where users change their python environments frequently.

This article goes into the challenges we faced, the solution we chose, and the performance impacts of that choice. Here are the results:

If you’re a Coiled user these changes are now deployed if you update to coiled>=0.4.

Background

In a previous blog post, we discussed the importance of matching your local Python packages and versions to those installed on remote machines running in the cloud. The easiest way to do this with a Coiled cluster is using package sync, which scans your local environment and replicates it on your cluster. This is faster, easier, and safer than building a Docker image each time.

However, some folks still want statically defined software environments stored in the cloud, like what Docker offers, and so we maintain an API create_software_environment , which used to make Docker images given conda/pip specifications, but now does something much faster.

coiled.create_software_environment(
    conda=["dask", "pandas", "pytorch", "arrow"],
    pip=["sqlalchemy", "dask-snowflake"],
)

Conda and Docker are slow

Previously, Coiled used conda to solve Python environments and Docker to build those environments. We then pushed those Docker images to users’ container registries and on cluster start, pulled these images onto the cloud instances that formed a Dask/Coiled cluster.

Benchmarking these steps using a complex/real-world software environment from the wild, we notice the first bottleneck is the environment solve. The customer reported anecdotally this could sometimes take 30 minutes! We fixed this simply by switching to Mamba, which is much faster.

That still leaves the build, push and pull steps. We noticed that this takes about 8 minutes; which is painfully slow for users. We break down that time cost below:

environment.yml

name: base
channels:
  - conda-forge
  - defaults
dependencies:
  - gcc
  - _libgcc_mutex=0.1=conda_forge
  -_openmp_mutex=4.5=2_gnu
  - abseil-cpp=20210324.2=h9c3ff4c_0
  - aiobotocore=2.3.4=pyhd8ed1ab_0
  - aiohttp=3.8.1=py39hb9d737c_1
  - aioitertools=0.10.0=pyhd8ed1ab_0
  - aiosignal=1.2.0=pyhd8ed1ab_0
  - alembic=1.8.1=pyhd8ed1ab_0
  - alsa-lib=1.2.6.1=h7f98852_0
  - anyio=3.6.1=py39hf3d152e_0
  - appdirs=1.4.4=pyh9f0ad1d_0
  - argon2-cffi=21.3.0=pyhd8ed1ab_0
  - argon2-cffi-bindings=21.2.0=py39hb9d737c_2
  - arrow-cpp=8.0.0=py39h811ffd7_4_cpu
  - asn1crypto=1.5.1=pyhd8ed1ab_0
  - asttokens=2.0.5=pyhd8ed1ab_0
  - async-timeout=4.0.2=pyhd8ed1ab_0
  - attr=2.5.1=h166bdaf_0
  - attrs=21.4.0=pyhd8ed1ab_0
  - aws-c-cal=0.5.11=h95a6274_0
  - aws-c-common=0.6.2=h7f98852_0
  - aws-c-event-stream=0.2.7=h3541f99_13
  - aws-c-io=0.10.5=hfb6a706_0
  - aws-checksums=0.1.11=ha31a3da_7
  - aws-sdk-cpp=1.8.186=hb4091e7_3
  - babel=2.10.3=pyhd8ed1ab_0
  - backcall=0.2.0=pyh9f0ad1d_0
  - backports=1.0=py_2
  - backports.functools_lru_cache=1.6.4=pyhd8ed1ab_0
  - beautifulsoup4=4.11.1=pyha770c72_0
  - bleach=5.0.1=pyhd8ed1ab_0
  - blinker=1.4=py_1
  - blosc=1.21.1=h83bc5f7_3
  - bokeh=2.4.3=py39hf3d152e_0
  - botocore=1.24.21=pyhd8ed1ab_1
  - brotli=1.0.9=h166bdaf_7
  - brotli-bin=1.0.9=h166bdaf_7
  - brotlipy=0.7.0=py39hb9d737c_1004
  - bzip2=1.0.8=h7f98852_4
  - c-ares=1.18.1=h7f98852_0
  - ca-certificates=2022.6.15=ha878542_0
  - certifi=2022.6.15=py39hf3d152e_0
  - cffi=1.15.1=py39he91dace_0
  - click=8.1.3=py39hf3d152e_0
  - cloudpickle=2.1.0=pyhd8ed1ab_0
  - colorama=0.4.5=pyhd8ed1ab_0
  - configparser=5.2.0=pyhd8ed1ab_0
  - cryptography=37.0.4=py39hd97740a_0
  - cycler=0.11.0=pyhd8ed1ab_0
  - cytoolz=0.12.0=py39hb9d737c_0
  - dask=2022.7.1=pyhd8ed1ab_0
  - dask-core=2022.7.1=pyhd8ed1ab_0
  - databricks-cli=0.17.0=pyhd8ed1ab_0
  - dbus=1.13.6=h5008d03_3
  - debugpy=1.6.0=py39h5a03fae_0
  - decorator=5.1.1=pyhd8ed1ab_0
  - defusedxml=0.7.1=pyhd8ed1ab_0
  - distributed=2022.7.1=pyhd8ed1ab_0
  - docker-py=5.0.3=py39hf3d152e_2
  - docker-pycreds=0.4.0=py_0
  - entrypoints=0.4=pyhd8ed1ab_0
  - executing=0.9.0=pyhd8ed1ab_0
  - expat=2.4.8=h27087fc_0
  - fftw=3.3.10=nompi_h77c792f_102
  - flask=2.1.3=pyhd8ed1ab_0
  - flit-core=3.7.1=pyhd8ed1ab_0
  - font-ttf-dejavu-sans-mono=2.37=hab24e00_0
  - font-ttf-inconsolata=3.000=h77eed37_0
  - font-ttf-source-code-pro=2.038=h77eed37_0
  - font-ttf-ubuntu=0.83=hab24e00_0
  - fontconfig=2.14.0=h8e229c2_0
  - fonts-conda-ecosystem=1=0
  - fonts-conda-forge=1=0
  - fonttools=4.34.4=py39hb9d737c_0
  - freetype=2.10.4=h0708190_1
  - frozenlist=1.3.0=py39hb9d737c_1
  - fsspec=2022.5.0=pyhd8ed1ab_0
  - gettext=0.19.8.1=h73d1719_1008
  - gflags=2.2.2=he1b5a44_1004
  - giflib=5.2.1=h36c2ea0_2
  - gitdb=4.0.9=pyhd8ed1ab_0
  - gitpython=3.1.27=pyhd8ed1ab_0
  - glib=2.72.1=h6239696_0
  - glib-tools=2.72.1=h6239696_0
  - glog=0.6.0=h6f12383_0
  - greenlet=1.1.2=py39h5a03fae_2
  - grpc-cpp=1.45.2=h3b8df00_4
  - gst-plugins-base=1.20.3=hf6a322e_0
  - gstreamer=1.20.3=hd4edc92_0
  - gunicorn=20.1.0=py39hf3d152e_2
  - heapdict=1.0.1=py_0
  - htmlmin=0.1.12=py_1
  - icu=70.1=h27087fc_0
  - idna=3.3=pyhd8ed1ab_0
  - imagehash=4.2.1=pyhd8ed1ab_0
  - importlib-metadata=4.11.4=py39hf3d152e_0
  - importlib_metadata=4.11.4=hd8ed1ab_0
  - importlib_resources=5.9.0=pyhd8ed1ab_0
  - ipykernel=6.15.1=pyh210e3f2_0
  - ipython=8.4.0=py39hf3d152e_0
  - ipython_genutils=0.2.0=py_1
  - itsdangerous=2.1.2=pyhd8ed1ab_0
  - jack=1.9.18=h8c3723f_1002
  - jedi=0.18.1=py39hf3d152e_1
  - jinja2=3.1.2=pyhd8ed1ab_1
  - jmespath=1.0.1=pyhd8ed1ab_0
  - joblib=1.1.0=pyhd8ed1ab_0
  - jpeg=9e=h166bdaf_2
  - json5=0.9.5=pyh9f0ad1d_0
  - jsonschema=4.7.2=pyhd8ed1ab_0
  - jupyter_client=7.3.4=pyhd8ed1ab_0
  - jupyter_core=4.11.1=py39hf3d152e_0
  - jupyter_server=1.18.1=pyhd8ed1ab_0
  - jupyterlab=3.4.4=pyhd8ed1ab_0
  - jupyterlab_pygments=0.2.2=pyhd8ed1ab_0
  - jupyterlab_server=2.15.0=pyhd8ed1ab_0
  - keyutils=1.6.1=h166bdaf_0
  - kiwisolver=1.4.4=py39hf939315_0
  - krb5=1.19.3=h3790be6_0
  - lcms2=2.12=hddcbb42_0
  - ld_impl_linux-64=2.36.1=hea4e1c9_2
  - lerc=3.0=h9c3ff4c_0
  - libblas=3.9.0=15_linux64_openblas
  - libbrotlicommon=1.0.9=h166bdaf_7
  - libbrotlidec=1.0.9=h166bdaf_7
  - libbrotlienc=1.0.9=h166bdaf_7
  - libcap=2.64=ha37c62d_0
  - libcblas=3.9.0=15_linux64_openblas
  - libclang=14.0.6=default_h2e3cab8_0
  - libclang13=14.0.6=default_h3a83d3e_0
  - libcrc32c=1.1.2=h9c3ff4c_0
  - libcups=2.3.3=hf5a7f15_1
  - libcurl=7.83.1=h7bff187_0
  - libdb=6.2.32=h9c3ff4c_0
  - libdeflate=1.12=h166bdaf_0
  - libedit=3.1.20191231=he28a2e2_2
  - libev=4.33=h516909a_1
  - libevent=2.1.10=h9b69904_4
  - libffi=3.4.2=h7f98852_5
  - libflac=1.3.4=h27087fc_0
  - libgcc-ng=12.1.0=h8d9b700_16
  - libgfortran-ng=12.1.0=h69a702a_16
  - libgfortran5=12.1.0=hdcd56e2_16
  - libglib=2.72.1=h2d90d5f_0
  - libgomp=12.1.0=h8d9b700_16
  - libgoogle-cloud=1.40.2=habd0e3a_0
  - libiconv=1.16=h516909a_0
  - liblapack=3.9.0=15_linux64_openblas
  - libllvm11=11.1.0=hf817b99_3
  - libllvm14=14.0.6=he0ac6c6_0
  - libnghttp2=1.47.0=h727a467_0
  - libnsl=2.0.0=h7f98852_0
  - libogg=1.3.4=h7f98852_1
  - libopenblas=0.3.20=pthreads_h78a6416_0
  - libopus=1.3.1=h7f98852_1
  - libpng=1.6.37=h753d276_3
  - libpq=14.4=hd77ab85_0
  - libprotobuf=3.20.1=h6239696_0
  - libsndfile=1.0.31=h9c3ff4c_1
  - libsodium=1.0.18=h36c2ea0_1
  - libssh2=1.10.0=ha56f1ee_2
  - libstdcxx-ng=12.1.0=ha89aaad_16
  - libthrift=0.16.0=h519c5ea_1
  - libtiff=4.4.0=hc85c160_1
  - libtool=2.4.6=h9c3ff4c_1008
  - libudev1=249=h166bdaf_4
  - libutf8proc=2.7.0=h7f98852_0
  - libuuid=2.32.1=h7f98852_1000
  - libvorbis=1.3.7=h9c3ff4c_0
  - libwebp=1.2.3=h522a892_1
  - libwebp-base=1.2.3=h166bdaf_2
  - libxcb=1.13=h7f98852_1004
  - libxkbcommon=1.0.3=he3ba5ed_0
  - libxml2=2.9.14=h22db469_3
  - libzlib=1.2.12=h166bdaf_2
  - llvmlite=0.38.1=py39h7d9a04d_0
  - locket=1.0.0=pyhd8ed1ab_0
  - lz4=4.0.0=py39h029007f_2
  - lz4-c=1.9.3=h9c3ff4c_1
  - mako=1.2.1=pyhd8ed1ab_0
  - markdown=3.4.1=pyhd8ed1ab_0
  - markupsafe=2.1.1=py39hb9d737c_1
  - matplotlib=3.5.2=py39hf3d152e_0
  - matplotlib-base=3.5.2=py39h700656a_0
  - matplotlib-inline=0.1.3=pyhd8ed1ab_0
  - missingno=0.4.2=py_1
  - mistune=0.8.4=py39h3811e60_1005
  - mlflow=1.27.0=py39ha39b057_0
  - msgpack-python=1.0.4=py39hf939315_0
  - multidict=6.0.2=py39hb9d737c_1
  - multimethod=1.4=py_0
  - munkres=1.1.4=pyh9f0ad1d_0
  - mysql-common=8.0.29=haf5c9bc_1
  - mysql-libs=8.0.29=h28c427c_1
  - nbclassic=0.4.3=pyhd8ed1ab_0
  - nbclient=0.6.6=pyhd8ed1ab_0
  - nbconvert=6.5.0=pyhd8ed1ab_0
  - nbconvert-core=6.5.0=pyhd8ed1ab_0
  - nbconvert-pandoc=6.5.0=pyhd8ed1ab_0
  - nbformat=5.4.0=pyhd8ed1ab_0
  - ncurses=6.3=h27087fc_1
  - nest-asyncio=1.5.5=pyhd8ed1ab_0
  - networkx=2.8.5=pyhd8ed1ab_0
  - notebook=6.4.12=pyha770c72_0
  - notebook-shim=0.1.0=pyhd8ed1ab_0
  - nspr=4.32=h9c3ff4c_1
  - nss=3.78=h2350873_0
  - numba=0.55.2=py39h66db6d7_0
  - numpy=1.22.4=py39hc58783e_0
  - oauthlib=3.2.0=pyhd8ed1ab_0
  - openjpeg=2.4.0=hb52868f_1
  - openssl=1.1.1q=h166bdaf_0
  - orc=1.7.5=h6c59b99_0
  - packaging=21.3=pyhd8ed1ab_0
  - pandas=1.4.3=py39h1832856_0
  - pandas-profiling=3.2.0=pyhd8ed1ab_0
  - pandoc=2.18=ha770c72_0
  - pandocfilters=1.5.0=pyhd8ed1ab_0
  - parquet-cpp=1.5.1=2
  - parso=0.8.3=pyhd8ed1ab_0
  - partd=1.2.0=pyhd8ed1ab_0
  - patsy=0.5.2=pyhd8ed1ab_0
  - pcre=8.45=h9c3ff4c_0
  - pexpect=4.8.0=pyh9f0ad1d_2
  - phik=0.12.2=py39ha791e8c_0
  - pickleshare=0.7.5=py_1003
  - pillow=9.2.0=py39hae2aec6_0
  - pip=22.2=pyhd8ed1ab_0
  - ply=3.11=py_1
  - portaudio=19.6.0=h57a0ea0_5
  - prometheus_client=0.14.1=pyhd8ed1ab_0
  - prometheus_flask_exporter=0.20.2=pyhd8ed1ab_0
  - prompt-toolkit=3.0.30=pyha770c72_0
  - protobuf=3.20.1=py39h5a03fae_0
  - psutil=5.9.1=py39hb9d737c_0
  - pthread-stubs=0.4=h36c2ea0_1001
  - ptyprocess=0.7.0=pyhd3deb0d_0
  - pulseaudio=14.0=h7f54b18_8
  - pure_eval=0.2.2=pyhd8ed1ab_0
  - pybind11-abi=4=hd8ed1ab_3
  - pycparser=2.21=pyhd8ed1ab_0
  - pydantic=1.9.1=py39hb9d737c_0
  - pygments=2.12.0=pyhd8ed1ab_0
  - pyjwt=2.4.0=pyhd8ed1ab_0
  - pyopenssl=22.0.0=pyhd8ed1ab_0
  - pyparsing=3.0.9=pyhd8ed1ab_0
  - pyqt=5.15.7=py39h18e9c17_0
  - pyqt5-sip=12.11.0=py39h5a03fae_0
  - pyrsistent=0.18.1=py39hb9d737c_1
  - pysocks=1.7.1=py39hf3d152e_5
  - python=3.9.13=h9a8a25e_0_cpython
  - python-dateutil=2.8.2=pyhd8ed1ab_0
  - python-fastjsonschema=2.16.1=pyhd8ed1ab_0
  - python_abi=3.9=2_cp39
  - python~=3.9
  - pytz=2022.1=pyhd8ed1ab_0
  - pywavelets=1.3.0=py39hd257fcd_1
  - pyyaml=6.0=py39hb9d737c_4
  - pyzmq=23.2.0=py39headdf64_0
  - qt-main=5.15.4=ha5833f6_2
  - querystring_parser=1.2.4=py_0
  - re2=2022.06.01=h27087fc_0
  - readline=8.1.2=h0f457ee_0
  - s2n=1.0.10=h9b69904_0
  - s3fs=2022.5.0=pyhd8ed1ab_0
  - scikit-learn=1.1.1=py39h4037b75_0
  - scipy=1.8.1=py39he49c0e8_0
  - seaborn=0.11.2=hd8ed1ab_0
  - seaborn-base=0.11.2=pyhd8ed1ab_0
  - send2trash=1.8.0=pyhd8ed1ab_0
  - setuptools=63.2.0=py39hf3d152e_0
  - shap=0.41.0=py39h1832856_0
  - sip=6.6.2=py39h5a03fae_0
  - six=1.16.0=pyh6c4a22f_0
  - slicer=0.0.7=pyhd8ed1ab_0
  - smmap=3.0.5=pyh44b312d_0
  - snappy=1.1.9=hbd366e4_1
  - sniffio=1.2.0=py39hf3d152e_3
  - sortedcontainers=2.4.0=pyhd8ed1ab_0
  - soupsieve=2.3.2.post1=pyhd8ed1ab_0
  - sqlalchemy=1.4.39=py39hb9d737c_0
  - sqlite=3.39.2=h4ff8645_0
  - sqlparse=0.4.2=pyhd8ed1ab_0
  - stack_data=0.3.0=pyhd8ed1ab_0
  - statsmodels=0.13.2=py39hd257fcd_0
  - tabulate=0.8.10=pyhd8ed1ab_0
  - tangled-up-in-unicode=0.2.0=pyhd8ed1ab_0
  - tblib=1.7.0=pyhd8ed1ab_0
  - terminado=0.15.0=py39hf3d152e_0
  - threadpoolctl=3.1.0=pyh8a188c0_0
  - tinycss2=1.1.1=pyhd8ed1ab_0
  - tk=8.6.12=h27826a3_0
  - toml=0.10.2=pyhd8ed1ab_0
  - toolz=0.12.0=pyhd8ed1ab_0
  - tornado=6.1=py39hb9d737c_3
  - tqdm=4.64.0=pyhd8ed1ab_0
  - traitlets=5.3.0=pyhd8ed1ab_0
  - typing-extensions=4.3.0=hd8ed1ab_0
  - typing_extensions=4.3.0=pyha770c72_0
  - unicodedata2=14.0.0=py39hb9d737c_1
  - urllib3=1.26.10=pyhd8ed1ab_0
  - visions=0.7.4=pyhd8ed1ab_0
  - wcwidth=0.2.5=pyh9f0ad1d_2
  - webencodings=0.5.1=py_1
  - websocket-client=1.3.3=pyhd8ed1ab_0
  - werkzeug=2.1.2=pyhd8ed1ab_1
  - wheel=0.37.1=pyhd8ed1ab_0
  - wrapt=1.14.1=py39hb9d737c_0
  - xcb-util=0.4.0=h166bdaf_0
  - xcb-util-image=0.4.0=h166bdaf_0
  - xcb-util-keysyms=0.4.0=h166bdaf_0
  - xcb-util-renderutil=0.3.9=h166bdaf_0
  - xcb-util-wm=0.4.1=h166bdaf_0
  - xorg-libxau=1.0.9=h7f98852_0
  - xorg-libxdmcp=1.1.3=h7f98852_0
  - xz=5.2.5=h516909a_1
  - yaml=0.2.5=h7f98852_2
  - yarl=1.7.2=py39hb9d737c_2
  - zeromq=4.3.4=h9c3ff4c_1
  - zict=2.2.0=pyhd8ed1ab_0
  - zipp=3.8.0=pyhd8ed1ab_0
  - zlib=1.2.12=h166bdaf_2
  - zstd=1.5.2=h8a70e8d_2
  - pip:
      - aenum==3.1.11
      - affinegap==1.12
      - altair==4.2.0
      - awswrangler==2.16.1
      - backoff==1.11.1
      - bcrypt==3.2.2
      - black==22.6.0
      - blis==0.7.8
      - boto3==1.21.21
      - btrees==4.10.0
      - cachetools==5.2.0
      - catalogue==2.0.7
      - categorical-distance==1.9
      - charset-normalizer==2.0.12
      - coiled==0.2.15
      - colored==1.4.3
      - commonmark==0.9.1
      - coverage==6.4.2
      - cymem==2.0.6
      - dateparser==1.1.1
      - datetime-distance==0.1.3
      - dedupe==2.0.13
      - dedupe-hcluster==0.3.9
      - dedupe-variable-datetime==0.1.5
      - dedupe-variable-name==0.0.14
      - doublemetaphone==1.1
      - elasticsearch==7.13.4
      - et-xmlfile==1.1.0
      - faiss-cpu==1.7.2
      - fastcluster==1.2.6
      - flake8==4.0.1
      - floret==0.10.2
      - future==0.18.2
      - gender-guesser==0.4.0
      - google-auth==2.9.1
      - google-auth-oauthlib==0.5.2
      - gremlinpython==3.6.0
      - gspread==5.4.0
      - gspread-pandas==3.2.2
      - haversine==2.6.0
      - highered==0.2.1
      - iniconfig==1.1.1
      - ipywidgets==7.7.1
      - isodate==0.6.1
      - jsondiff==2.0.0
      - jsonpath-ng==1.5.3
      - jupyterlab-widgets==1.1.1
      - langcodes==3.3.0
      - levenshtein-search==1.4.5
      - lxml==4.9.1
      - mccabe==0.6.1
      - murmurhash==1.0.7
      - mypy-extensions==0.4.3
      - openpyxl==3.0.10
      - opensearch-py==1.1.0
      - paramiko==2.11.0
      - parseratorvariable==0.0.18
      - pathspec==0.9.0
      - pathy==0.6.2
      - pdfkit==1.0.0
      - persistent==4.9.0
      - pg8000==1.29.1
      - platformdirs==2.5.2
      - pluggy==1.0.0
      - preshed==3.0.6
      - probableparsing==0.0.1
      - probablepeople==0.5.4
      - progressbar2==4.0.0
      - psycopg2-binary==2.9.3
      - py==1.11.0
      - pyarrow==7.0.0
      - pyasn1==0.4.8
      - pyasn1-modules==0.2.8
      - pycodestyle==2.8.0
      - pydeck==0.7.1
      - pyflakes==2.4.0
      - pyhacrf-datamade==0.2.6
      - pylbfgs==0.2.0.14
      - pympler==1.0.1
      - pymysql==1.0.2
      - pynacl==1.5.0
      - pytest==7.1.2
      - pytest-cov==3.0.0
      - python-crfsuite==0.9.8
      - python-dotenv==0.20.0
      - python-utils==3.3.3
      - pytz-deprecation-shim==0.1.0.post0
      - ratelimit==2.2.1
      - redshift-connector==2.0.908
      - regex==2022.3.2
      - requests==2.28.0
      - requests-aws4auth==1.1.2
      - requests-oauthlib==1.3.1
      - rich==12.5.1
      - rlr==2.4.6
      - rsa==4.9
      - s3transfer==0.5.2
      - scramp==1.4.1
      - semver==2.13.0
      - simplecosine==1.2
      - smart-open==5.2.1
      - spacy==3.4.0
      - spacy-legacy==3.0.9
      - spacy-loggers==1.0.3
      - srsly==2.4.4
      - sshtunnel==0.4.0
      - streamlit==1.11.0
      - syrupy==2.3.1
      - textual==0.1.18
      - thinc==8.1.0
      - tomli==2.0.1
      - ttictoc==0.5.6
      - typer==0.4.2
      - tzdata==2022.1
      - tzlocal==4.2
      - unidecode==1.3.4
      - validators==0.20.0
      - wasabi==0.9.1
      - watchdog==2.1.9
      - widgetsnbextension==3.6.1
      - zope-index==5.2.0
      - zope-interface==5.4.0

Step	Time in seconds (mean ± σ)
Build image	155 ± 1.7
Push image to registry	126
Pull image to cloud instance (t3.medium)	163 ± 21

Why is this slow? Mostly compression.

Docker images are compressed with bzip2 which is efficient at creating smaller files, but slow as it trades CPU time for smaller image size. Caching helps a little, but most of each image is Python packages, which vary wildly and are hard to cache effectively. Additionally, each image layer is decompressed sequentially, which also slows things down despite using a multithreaded Gzip implementation.

So we dropped Docker.

A brief word from our sponsors: Coiled makes it very easy to scale up your Python code in the cloud. We deploy in your AWS or Google Cloud account, and we configure the infrastructure so that you get sensible security and cost-savings by default. If you have Python code that you need to run at scale, give us a try!

Switching to mamba

Mamba environments normally aren’t portable (and in fact has had to pull some tricks to allow package installation in the first place). We can cheat though because Mamba environments are portable if you keep them in the same location on the filesystem, and we control the filesystem on our clusters. Our new build system instead does the following:

Creates an environment on a large machine
Turns it into lz4-compressed tarballs
Loads it into Amazon S3
Downloads/extracts it from s3 directly onto the cluster

We can tune each of these stages using more modern techniques than Docker provides.

Tuning Mamba performance

First, let’s understand how Mamba works:

Download metadata.json for each channel in use (40MB+)
Solve package dependencies (single-core, RAM hungry)
Downloading and extract the required packages
Symlink those packages into the new environment
Run any post_install steps for packages that have them
Run pip install on any additional pip packages

This process presents some challenges to overcome:

Solving the environment can be slow and require lots of memory (up to 10 GB+)
Downloading packages can be limited by network and disk speeds

To address these we perform a few basic optimizations:

Download packages concurrently
Use a big machine and ask conda/pip to store packages to a RAM disk (we’re using S3 as persistent storage eventually, so local disk is slow and not that useful anyway)
Compress with LZ4 rather than bzip2
Stream into S3 in 16 MB chunks, achieving multiple gigabit upload and download speeds

Results: Faster builds

Our new system provides the same user experience, but is significantly faster:

That’s under 2 minutes total, over three times faster than using a Docker image.

If you’re using coiled>=0.4, you’ll automatically be using the new environment build system, and you can keep using the old system by pinning an older version of Coiled. Though we’re no longer using Docker to build software environments, you can keep using pre-built Docker images in your clusters. See our docs for more details.

If you’re not using Coiled then what are you waiting for? The Coiled free tier is enough for most people’s workloads and is easy to get started.

How many PEPs does it take to install a package?

2023-01-17T00:00:00+00:00

A few months ago we released package sync, a feature that takes your Python environment and replicates it in the cloud with zero effort.

The reaction has been very positive; people enable the feature and then mostly forget about their previous woes with docker or version mismatches.

However to make this feature really stand out, it needs to work the first time, every time. The feature worked great when you used Python exactly like we did, but there are so many ways of using Python it was rare for this to be true.

The first and most critical step for package sync, is to identify every installed package, the installed version, and where it came from. This is critical information if we’re going to reverse engineer your environment! Some existing tools try this, like pip list, but none provided all the information we needed.

To achieve this we needed to do a deep dive on not only current PEP standards for Python environments, but every previous iteration of them (including some from before the time of PEPs) and every operating system we intended to support.

There’s a surprisingly large number of ways of getting a Python package in position so import foo works.

Eggs
PEP 376 Dist Info
Egg Links
.pth files
PEP 660 Editable Installs
Python files or modules in your current working directory
Manually dropping Python files or module in your site-packages folder (hands up if you kept site-packages in VCS back in the day!)

Some of these also have different variants. For example, PEP 376 has an addendum that allows a JSON file describing where the package was installed from.

While Python has a relatively powerful tool to discover the names of all the packages installed, we have to find out where you got those packages from and how we might be able to repeat what you did. Did you install dask from conda, Git, PyPI or perhaps a local editable directory?

Every package manager also implements its own version of these standards. For example, Poetry prior to version 1.2.0 implemented an incorrect standard of PEP 610. PDM uses PEP 582, which is still a draft, and installs everything to a __pypackages__ folder.

Each operating system also has its own quirks, we recently had to deal with importlib.metadata reporting paths to files in egg packages with / separators instead of \, as they are technically in a zip file.

So how do we make sure package sync works and keeps working in all of these cases?

With all this knowledge we built out a large integration test suite, where we install a dummy version of every package variant we know of using a bunch of different package managers on a bunch of a different operating systems.

It currently has 29 different combinations of operating systems, versions, and package managers.

OS:

Ubuntu
Windows
macOS
Debian
amazonlinux

Package managers:

apt
yum
venv
pyenv
mamba
venv

Python versions:

3.11
3.10
3.9
3.8
3.7

During development of the test suite we squashed 7 different bugs, some minor and some show stoppers.

Our test suite will grow as we spot more ways of managing package and as the Python community creates even more. I created package sync with a firm belief that our solution should “just work” for everyone.

Scaling Hyperparameter Optimization With XGBoost, Optuna, and Dask

2023-01-06T00:00:00+00:00

Scaling XGBoost and Optuna with Dask

XGBoost is one of the most well-known libraries among data scientists, having become one of the top choices among Kaggle competitors. It is performant in a wide of array of supervised machine learning problems, implements scalable training through the rabit library, and integrates with many big data processing tools, including Dask.

While XGBoost enables distributed model training, hyperparameter optimization (HPO) adds a layer of complexity which only grows with the size of the dataset. Tools like Optuna have emerged to reduce this friction by dynamically optimizing the HPO search space providing much-needed visibility into this process, and enabling parallel HPO, but in this paradigm, the distributed HPO step remains disconnected from data engineering. Examples of recent integrations between Dask and Optuna show how bringing distributed feature engineering together with HPO simplifies the workflow for machine learning engineering workflows.

Here, we demonstrate one way to connect these pieces of the data science workflow on a dataset with hundreds of millions of rows. This is accomplished by leveraging Dask for exploration and feature engineering. Distributed hyperparameter optimization and distributed model training are enabled through native dask integration with Optuna and XGBoost, respectively. We bring all of these tools together by spinning up dask clusters on-demand with Coiled. Code for this work is located at this GitHub repository.

The Problem To Solve

For this example, we extend on our work demonstrating the use of Coiled and Prefect to automate data engineering workloads, by leveraging the initially prepared New York City “For Hire Vehicle” taxi dataset as a source for exploratory analysis. Cleaning the data consists of declaring datatypes and partitioning files to manageable partition sizes.

One motivating problem with the New York City Taxi dataset is to predict expected trip duration. This has been the topic of both Kaggle competitions and the subject of research papers. We will use this as a motivating example, as it allows us to explore many steps in the data science workflow. Much of this previous work was based on the New York City Taxi data. In our example, we will focus on the New York City For Hire Vehicle dataset (e.g. Uber/Lyft).

Given this, our goal is to train an XGBoost model to predict trip duration (in seconds), given a projected route, projected pickup time and known pickup/dropoff locations as well as current traffic conditions. Given those variables, we expect variables like tolls, congestion_surcharge and airport_fee are available, whereas sales_tax and tips can only be known after the trip has been completed. We will also take inspiration from the geospatial research paper above, by engineering categorical features based on the concept of Superboroughs, which are derived from domain knowledge about trips that require bridge crossings.

Getting Started

Since it’s faster to perform exploratory analysis with data that fits in memory, we start by creating a distributed compute environment and select a representative subset of the data. Using Coiled, we spin up a Dask cluster, and after some exploration, sample 2.5% of the data (or 1.6MM rows), and write this to a single local Parquet file for ease of use.

Exploratory analysis with Pandas shows our travel time is (unsurprisingly) not normally distributed, and almost certainly contains outliers. We remove outliers using the IQR Method. Given the knowledge no trip should have a duration with a negative value, we floor the lower bound at 0 seconds. Boundaries are shown on the histogram plot, and filtering results in a removal of 4.6% of the data.

Building on our research paper’s definition of a Superborough with the taxi+_zone_lookup dataset, we map each Pickup and Dropoff LocationID to a Borough, and subsequently to a Superborough, deriving a categorical feature CrossSuperborough that defines whether a trip requires crossing between Superboroughs.

We drop columns containing data that is expected to be unavailable at the time of prediction and clean up features related to route planning, with the result being a mixture of float, int, and categorical variables.

Finally, for the in-memory example, we prepare three models using the following estimators: DummyRegressor(), LinearRegression(), and XGBRegressor(), and compare their performance when evaluating against a test dataset using rmse metric. Results are summarized below, recalling that our travel durations are expressed in seconds:

Estimator	RMSE
DummyRegressor()	556.7
LinearRegression()	324.4
XGBRegressor()	236.2

These results, while sligthly better than those reported in the research article, mirror the trend reported by the authors.

Scaling XGBoost and Optuna with Dask

One of the enabling attributes about working in Dask is the tight integration between the pandas and Dask. Scaling our pandas feature engineering to the entire dataset with Dask is more natural than if we were to rewrite it to a different language, which accelerates the process of scaling our workload from a subset of the data to the entire dataset. Executing on this vision is accomplished by breaking our feature engineering and model training steps into two discrete pipelines, as shown in feature engineering and model training notebooks.

Our feature engineering occurs in two steps. The first loads our minimally cleaned dataset into a Dask cluster, creates a handful of columns that will ultimately become categorical variables, drops unneeded columns, and removes outliers. Our second function merges the data with the taxi+_zone_lookup data to create the Superborough categorical feature.

A few observations about scaling this work to the entire dataset:

There’s active work underway to fully implement pandas Nullable datatypes, but at the time of this writing, using these datatypes across pandas, Dask, Parquet Datasets and XGBoost is not recommended.
The default method for calculating percentiles with dask is known to give potentially incorrect results. For more information, see here, here, and here. This manifests in our outlier detection method, such that <0.25% of the data is discarded when performing the percentile calculation with dask on the entire dataset, as opposed to the 4.6% discarded our sample. The suggested solution is to use the tdigest method, which requires crick, which itself is currently being updated for compatibility with recent version of python and scipy. We attempt to mitigate this by using the 75th percentile calculated by pandas (under the assumption that our randomly sampled dataset is representative of the entire dataset), but this approach only discards 1.6% of the data, suggesting our assumption is incorrect.
Filtering a Dask DataFrame can lead to empty, or sparsely populated partitions. See the following for best practices around repartitioning DataFrames for optimal performance.

Tuning the XGBoost Model

Recent work by Dask and Optuna developers enabled native Dask integration with Optuna. This work, combined with the enabling of study.optimize in multiple threads allows parallel hyperparameter optimizaton by creating a Python function that will:

Create a remote dask cluster
Load the feature_table from a Parquet Dataset
Train an XGBoost model using cross-validation

Executing this function in a thread allows running many dask clients locally, with each client connecting to a separate remote dask cluster. Optuna enables hyperparameter optimization across these dask clusters by independently training xgboost models.

A few notes and observations about performance:

Using categorical dtypes requires that the categories are known. Writing a DataFrame with categorical datatypes to a Parquet Dataset causes known categories to be lost.
scikit-learn implements cross-validation by generating indices for splitting a dataset into train and test sets, but positional indexing into a Dask DataFrame is inefficient. Here, we tackle this by making use of DataFrame.random_split().
Creating a DMatrix forces lazy computation to evaluated, so we need to be mindful of when and how these are created.
We use the dask-ml implementation of mean_squared_error to keep this computation on the cluster.
We periodically observe concurrent.futures._base.CancelledError when performing our cv_split. This seems to be connected to dask/distributed #4612. We can avoid this by adding an explicit wait when calculating the rmse score from each training iteration.

Why Does The Dask Cluster Seem Idle When Training A XGBoost Model?

When training XGBoost models on a dask cluster, parallelization of the training step is handed off to rabit, which uses all available threads on each dask worker. From the user’s perspective, the dask dashboard may seem idle during this phase of the workflow, but watching the CPU Utilization tab will give visibility to the training step, even though the dask task stream is empty. By default, xgboost uses all available threads on each worker for training. Here, we train on a dask cluster with 10 workers, and compare the performance of multiple EC2 Instance Types, with the following execution time / behavior. Later in the post, we will explore this further.

EC2 Instance Type	# Cores	Memory (GiB)	Training time / iteration
m6i.xlarge	4	16	Training doesn’t start after 5 minutes, killed
m6i.2xlarge	8	32	Training succeeds, 90–180-seconds – some trials fail with `../rabit/include/rabit/internal/utils.h:86: Allreduce failed`
m6i.4xlarge	16	64	Training succeeds, 90–180-seconds

Model Training Outcomes

Below, we can see the outcome of our hyperparameter optimization experiment. This shows the progression of the rmse during the optimization process. We trained 129 separate models using 5-Fold Cross-Validation, with the best model having a rmse=286.4 seconds.

We can also inspect the impact of different hyperparameters on the overall model performance, as shown in the Hyperparameter Importances plot below.

At this point, we have the option of training a model using the best_params from our Optuna Study to save for deplopyment, or optionally, doing additional experimentation based on the most current set of results.

Diving Deeper Into Training XGBoost on Dask

In the above look at the impact of Dask workers based on different EC2 Instance types on model training, we observed dramatically different performance during the XGBoost model training phase of our workflow. We can leverage recent work with Grafana by the team at Coiled to delve into these differences.

The above Grafana plot show cluster-wide memory-usage during the dispatch train step for ourour workflow. The white line is total memory usage, the green line is total dask memory usage, and the yellow line is total_occupied_memory. Normally, we expect dask’s view of memory and the host’s view of memory to be very close. Recall however, that training the XGBoost model is dispatched to the rabit library.

At about the 09:30:10 mark, the Grafana dashboard shows memory pressure falling, which is consistent with what is reported on the dask dashboard. But host memory pressure spikes at 09:30:15. This happens when model training is dispatched to rabit. Dask has no visibility to this process, so everything looks fine on the dask dashboard, but Grafana sees this on the host. In fact, the memory pressure was so high that at least one EC2 instance was unable to report metrics for about 20 seconds, as evidenced by the apparent dip in total_cluster_memory. We also see that, during this 17-second time window, no tasks are completed.

By comparison, the following plots show the same views during the dispatch step, when run on a cluster provisioned with m6i.4xlarge workers.

Notice the much smaller gap between the host’s view of memory usage and dask’s view of memory usage on the cluster.

To extend on this further, the following plot shows total memory usage during an AllReduce error event. Observe the gap between the host’s view of memory and dask’s view.

While we have no idea how this will (or will not) manifest when deploying Dask to different deployment environments, it’s a potential pitfall when training XGBoost models with Dask. It also presents a potential opportunity for enhancements in future work.

Handling Unexpected AWS IAM Changes

2023-01-06T00:00:00+00:00

The cloud is tricky! You might think the rules that determine which IAM permissions are required for which actions will continue to apply in the same way. You might think they’d apply the same way to different AWS accounts. Or that if these things aren’t true, at least AWS will let you know. (I did.) You’d be wrong!

Coiled operates in our users’ AWS (or Google Cloud) accounts, with IAM permissions they assign to us. We have a CLI that lets users get set up by assigning us a minimal set of IAM permissions to do stuff in their account (like run EC2 instances), and we spin up Dask clusters when they ask us to. Or they can manually set things up by creating roles with the permissions we document that we need.

We also “tag” the resources we create, to help our users track down what Coiled is up to (for cost tracking or other auditing).

Someone seems to be on a rampage to fix some bugs in how IAM permissions apply to tags, and it’s been a bit rough on some of our users.

Act 1: a `logs:TagLogGroup` problem in November

When a user runs coiled setup aws, we create a CloudWatch log group in their account, where we’ll send the logs from their Dask cluster.

Back in November, a new user in a small AWS region reported a problem setting up Coiled:

User with account ID: <redacted> is not authorized
to perform CreateLogGroup with Tags.

We’d never seen this before and couldn’t reproduce it in any of our larger regions, but could reproduce in the user’s region.

This was a bit of a mystery, but not that much of one. I guess someone realized that they hadn’t been requiring the TagLogGroup permission to create a log group with tags (only to tag an existing log group?), and fixed it. Maybe the “fix” appeared first in this small alphabetically-early region, and broke us there.

We added logs:TagLogGroup to the list of permissions we ask users to give us, and all was well.

Act 2: a warning for January!

On December 12, I received an email (well, many emails across several languages) that seemed related:

As of October 30, 2022, tagging is supported for the “Destination” resource. Previously, CloudWatch Logs supported tagging only for the “Log Group” resource. We recommend that, for your IAM policies that are used to access the CreateLogGroup API, you add logs:TagResource permission to your IAM policies by January 31, 2023. The new logs:TagResource permission will not be required for existing accounts that previously used CreateLogGroup API with tags.

In order to tag new log groups using the CreateLogGroup API, we recommend you add logs:TagResource permission to your IAM policies …

At first I read this as a belated warning about the behavior from November. But no, it’s not actually about a change from needing logs:TagLogGroup to needing logs:TagResource. Thanks for the warning! That would have been nice back in November. (I still haven’t seen any acknowledgement about the behavior change requiring logs:TagLogGroup.)

Act 3: an `ecr:TagResource` problem in December

This morning a new user reported another IAM problem. They were trying to create a Coiled “software environment” (just a Docker image we build and push to ECR for them), and saw:

botocore.exceptions.ClientError: An error occurred (AccessDeniedException) when calling the CreateRepository operation: User:
arn:aws:iam::<redacted>:user/coiled is not authorized to perform: ecr:TagResource on resource: <redacted> because no identity-based
 policy allows the ecr:TagResource action

I quickly reproduced the problem and helped the user past it by giving us the ecr:TagResource permission. We’ve seen this before! Seems like they should have been requiring that all along, and broke us by requiring it now.

My next question was: Why isn’t this affecting more users?

Assisted by the hint in AWS’s December 12th email, I realized that maybe accounts which had previously created an ECR repo (tagged) would be unaffected. After confirming this hypothesis (my first repro was lucky to use a disposable account, which had never created an ECR repo, or at least not in that region), I exited “emergency” mode since I now knew existing users’ workflows aren’t at risk.

I haven’t found any public discussion of this change.

We deployed a fix and set up notifications to alert us if this ever happens again so we can individually help affected users (which it might, since a few users “set up” Coiled with the old permissions, but have not yet tried to create a software environment).

Act 4: Test in new accounts…?

It seems like AWS’s strategy is to make these changes in a way that will only affect accounts which haven’t used those permissions before.

That keeps these changes from breaking existing users, and helps folks choosing a collection of IAM permissions for the first time. But we’re asking for a fixed collection of permissions from AWS accounts that may have never taken any of these actions before.

We already do testing in accounts that are wiped clean nightly, but that’s not enough to catch these issues.

I don’t know how much more of this to expect (or how much will be announced to us ahead of time). If this is going to keep happening, it would probably pay off to set up a process to programmatically create completely brand new accounts in which to test Coiled.

Obligatory pitch: Use Coiled!

At Coiled we love these absurd details, and we love dealing with them so you don’t have to. We pay close attention to strange cloud behaviors and have fixes in place before most (or often any) users are aware.

AWS Cost Explorer Tips and Tricks

2023-01-06T00:00:00+00:00

Jan 6, 2023

7m read

Nat Tabris

Spending time in AWS Cost Explorer is one of the best ways to understand what’s going on in your AWS account. It’s one of the few places in the AWS Console where you can get a global view of your account or even of your entire organization.

In this post, I’m going to share some views that I’ve found especially useful when digging in to our costs for running data-science workloads on EC2 instances. Along the way, you might also learn some ways to reduce those costs.

Before we dig in, let’s talk about some basics. You can access Cost Explorer by logging in to the AWS Console and following this link (or going to the “AWS Cost Management” service and clicking “Cost Explorer” in the side menu).

Cost Explorer provides ways to filter and group your cost data. You can select time ranges, and you can see hourly, daily, and monthly costs, although hourly data (and other fine-grained views) are only available for more recent costs.

For a high-level view, I often start by disabling all filters and grouping by “Service”:

and then by “Region”:

Depending on what you see, you might want to explore in more depth by (e.g.) filtering to a specific region and looking at the breakdown by service for just that region. Or filter to a service and then grouping by region to see where the service is being used.

AWS Cost Explorer has powerful filtering and grouping functions. Powerful—but also not entirely intuitive.

At Coiled we make lots of clusters of EC2 instances for running big-data computations at scale. I personally spend a lot of time in AWS Cost Explorer trying to understand and optimize costs related to these clusters. In this post, I’ll explain the most helpful ways I’ve found for finding and understanding these costs. I’ll also share some tips I’ve learned about how to reduce these costs.

A brief word from our sponsors: Coiled makes it very easy to scale up your Python code in the cloud. We deploy in your AWS or Google Cloud account, and we configure the infrastructure so that you get sensible security and cost-savings by default. If you have Python code that you need to run at scale, give us a try!

Spot vs On Demand Instance Usage

Group by	Filters
Purchase option	Service: EC2-Instances

If you want to see cost and hours, add “Usage type” filter to everything that includes “Box” or “Spot”. (Open the Usage type filter dropdown, type “box” and click select all. Then type “spot” and click select all again. Now apply and you should get usage in hours for both on demand “box” and spot instances.)

Cost savings tip: As you can see, spot is a great way to save money! In December, our usage hours for on demand is very low, yet this accounts for over half of our cost. See this blog post for our approach to using spot instances well.

Once you’ve logged in to the AWS Console, here’s a link to bring up this view for your own AWS account: Spot vs On Demand (by month, July–December 2022)

Instance Types

Same as the Spot vs On Demand, but group by Instance type.

Group by	Filters
Instance type	Service: EC2-Instances

If you want to see cost and hours, add “Usage type” filter to everything that includes “Box” or “Spot”.

See this in your own AWS account: Instance types (by month, July–December 2022)

Burstable Instances

In addition to the base cost for burstable instances (t2/t3/t4g), you might pay extra if you use more than the “baseline” of CPU util.

When you have burstable instances and are using less than the baseline, you’ll accumulate credits that you use when you go above the baseline. But if you burst more than your accumulated credits, you’ll pay extra, and that’s what we’re looking at here.

Group by	Filters
Usage type	Service: EC2 - Other Usage type: include everything with “CPUCredits”

See this in your own AWS account: Extra burst credits (by month, July–December 2022)

Cost savings tip: As you can see, extra burst credits can be expensive for CPU-intensive workloads! We’ve found that by switching from t3 to m6i instance types we’ve significantly reduced our AWS bill. See this blog post for more details.

EBS Disks

Group by	Filters
Usage type	Usage type: include everything with “EBS”

This will show both cost and usage in GB-Months. (See AWS docs for details about EBS pricing.)

Note that if you have Amazon Machine Images (AMIs), these often are backed by EBS snapshots so will be shown here under “SnapshotUsage”.

See this in your own AWS account: EBS (by month, July–December 2022)

Cost savings tip: EBS is relatively cheap compared to EC2 but can still add up. In September our cost for compute was about $2200 and our cost for EBS disks was about $250. We’ve reduced these costs by using smaller disks as the default (with the option to increase disk size for specific workloads).

Network traffic for EC2 instances

You pay for moving data from AWS to the internet, you pay for moving data between regions, and you pay for moving data between zones inside a region.

For reasons unknown to me, most types of traffic are listed under “EC2-Instances” service but interzone traffic shows under the “EC2 - Other” service.

To see all of this traffic, broken down by type:

Group by	Filters
Usage type	Service: EC2-Instances and EC2 - Other Usage type group: include everything with “EC2: Data”

See this in your own AWS account: EC2 traffic breakdown (by month, July–December 2022)

As you can see, there’s lots of “Other” traffic that didn’t cost anything. Presumably this is traffic to your EC2 instances, which is free:

If you want to see just one type of traffic, you can be more specific in the Usage type group filter. For example, here’s cost and usage for traffic between zones:

See this in your own AWS account: Interzone traffic (by month, July–December 2022)

Cost savings tip: We’re running distributed Python computations at scale and for some workloads there’s lots of data moved between compute nodes. For a while we experimented with clusters that were in multiple availability zones. As you can see, this was expensive. We now make sure that all the instances for a cluster are in the same zone.

NAT Gateways

We don’t use NAT Gateways for our compute clusters, but we do use them for our control-plane infrastructure.

NAT Gateways have a fixed hourly cost and a per GB traffic cost.

You can see both together like this:

Group by	Filters
Usage type	Service: EC2 - Other Usage type group: “EC2: NAT Gateway - Data Processed” and “EC2: NAT Gateway - Running Hours”

See this in your own AWS account: NAT Gateway total cost

If you restrict the usage type group filter to just “EC2: NAT Gateway - Data Processed”, you’ll get both cost and GB usage:

See this in your own AWS account: NAT Gateway data transfer

Cost savings tip: If you can, avoid using NAT Gateways. For our compute clusters, we use public IP addresses (for outbound connections) and then use security groups to block inbound connections to the workers (we only want inbound connections to the scheduler node).

Other costs

There’s a good chance you’ll have costs for other services associated with your workloads. In our case, we have

S3 storage costs (transfer costs are minimal because we keep our EC2 instances in the same region as the relevant S3 buckets)
ECR storage and transfer (we build container images with our code and store these images in ECR)
CloudWatch logs (usually very low cost)

The best way to discover additional costs in your own account is by disabling all filters and grouping by “Service”. When there’s a service with a non-trivial (or higher than expected cost), I’ll usually dig in by filtering to that service and then grouping by various dimensions: region, usage type, maybe even by resource or by any tags that we use to distinguish costs for different projects.

Cost Explorer Links

Use these views to start your own cost exploration.

Log in to the AWS Console, and if you use AWS roles, switch to the role associated with the account you want to explore. The links below should then take you directly to Cost Explorer with the filters and grouping already selected for you.

Once the view is loaded, you can change the date range, add filters (to, e.g., look at a specific region), see hourly or daily costs, or try other ways of grouping the data.

Automated Data Pipelines On Dask With Coiled & Prefect

2022-12-19T00:00:00+00:00

Operationalizing Pure Python Data Engineering Workloads

Dask is widely used among data scientists and engineers proficient in Python for interacting with big data, doing statistical analysis, and developing machine learning models. Operationalizing this work has traditionally required lengthy code rewrites, which makes moving from development and production hard. This gap slows business progress and increases risk for data science and data engineering projects in an enterprise setting. The need to remove this bottleneck has prompted the emergence of production deployment solutions that allow code written by data scientists and engineers to be directly deployed to production, unlocking the power of continuous deployment for pure Python data science and engineers.

Here, we demonstrate the use of Coiled and Prefect to automate a common component of a data engineering pipeline. The motivating example is a file ingestion pipeline that a) runs on a scheduled basis, b) leverages distributed computing capability of Dask to process highly variable and potentially large volumes of data, and c) allows for reprocessing the data when a change in business logic or incoming data mandates a change.

The Problem To Solve

The New York City Taxi Dataset is well-known in the big data community. This dataset is now being further enriched by including For-Hire Vehicle records (Uber, Lyft, etc). In May 2022, data stewards moved from a bi-annual to monthly update cadence, and switched to Parquet format. These changes open the door to using the data to experiment with modern data engineering and data science workloads.

Experimentation with the data demonstrates that, while we can take advantage of the Parquet file format, row groups are nonuniform and at times very large, which makes working with the data cumbersome. Additionally, there are no guarantees about data quality or uniformity of datatypes. This creates unnecessary barriers and risk to using the data in a real-world business setting, where data quality is of paramount importance.

Here, we take the first step in a common data engineering pipeline from a raw data source by automating the ingestion and initial cleaning of the data for downstream consumers. This preprocessing step assures the data is properly formatted for downstream users. We do this in a way that minimizes cost by waiting until we know there are files to process before creating a Dask cluster on Coiled, and automatically tearing down the cluster when the workload is complete.

Getting Started

The following is an architecture diagram of the solution to build.

We assume users have both Coiled and Prefect accounts, are using Prefect >=2.0, and have a private S3 bucket with credentials for landing the output our our data transformation pipeline. Clone the coiled-from-prefect repository, create the Python environment, and collect the following information:

Coiled

Coiled User and Team Account information, and a Coiled API token This allows creating Coiled clusters from Prefect.

The NYC Taxi Dataset is in a public S3 bucket (labeled Source) above)
An AWS Access Key Id and AWS Secret Access Token to write to your private S3 bucket (labeled Sink above). The bucket must be configured for write access

Prefect

Create the following Prefect Blocks in your Prefect Workspace for storing secrets:
1. AWS Credentials Block named prefectexample0. Add the Aws Access Key Id and Aws Secret Access Key for your S3 sink here. Dask workers need these credentials to write the output of your data transformation pipeline
2. S3 Block named prefect-s3-storage. This is where user code is stored for Prefect to run in automation. In our example, this is the same S3 bucket as our target, but we specify a Bucket Path of bucket/prefect-s3-storage. This Block also expects an AWS Access Key ID and AWS Secret Access Key
3. Three separate Secret Blocks, named: coiled-account, coiled-user, and coiled-token. These enable authentication to your Coiled account.

In a terminal, and from your active Python environment running python flows/coiled_flow.py should trigger the execution of your Prefect flow with the following output:

Prefect creates a flow from the check_for_files Python function, and generates a random name tuscan-cormorant. The Flow is parameterized for testing purposes to only choose the first file from the list, and so returns the file s3://nyc-tlc/trip data/fhvhf_tripdata_2019-02.parquet. Since we have a file that needs to be processed, this triggers execution of second Prefect Flow that dynamically uses Coiled to create a Dask and submits the file processing job for execution.

Insepcting the logs shows a few items that are noteworthy:

The cluster comes up quickly; the Dask dashboard is available about 90-seconds after being requested.
Package Sync makes it easy to replicate our local software environment on the remote cluster

Create a Prefect Deployment, Work Queue, and Agent

In the terminal, run:

prefect deployment build flows/coiled_flow.py:check_for_files -n check_for_files  -q test -sb s3/prefect-s3-storage

This will create check_for_files-deployment.yaml (example below), provision a Prefect Work Queue and push all the code needed for remote execution to the S3 bucket defined in Prefect S3 Block.

###
### A complete description of a Prefect Deployment for flow 'Check for files'
###
name: check_for_files
description: |-
  We're going to run this on a weekly schedule.  Our intention is, as a default behavior,
  to find files in the S3 bucket that match `fhvhv_tripdata_*.parquet` and evaluate
  their modification timestamp.  We only want to process files that were
  modified in the last seven days.

  We also have the option to either: a) `reprocess` the entire dataset, which will delete
  the existing data and rewrite it, or b) `test_subset`, which will process only the first file
  in the list.  This is for testing only.

  Parameters
  ----------

  intent: str: one of `test_subset`, `reprocess`, or ""
version: ae6b4b09dd5bdfeb4e0de153ef60b7c6
# The work queue that will handle this deployment's runs
work_queue_name: test
tags: []
parameters:
  intent: ""  <--- Verify this line is populated
schedule: null
infra_overrides: {}

###
### DO NOT EDIT BELOW THIS LINE
###
flow_name: Check for files
manifest_path: null
infrastructure:
  type: process
  env: {}
  labels: {}
  name: null
  command: null
  stream_output: true
  working_dir: null
  _block_document_id: <XXX>
  _block_document_name: <YYY>
  _is_anonymous: true
  block_type_slug: process
  _block_type_slug: process
storage:
  bucket_path: <S3 Bucket>/prefect-s3-storage
  aws_access_key_id: '**********'
  aws_secret_access_key: '**********'
  _block_document_id: <ZZZ>
  _block_document_name: prefect-s3-storage
  _is_anonymous: false
  block_type_slug: s3
  _block_type_slug: s3
path: ''
entrypoint: flows/coiled_flow.py:check_for_files
parameter_openapi_schema:
  title: Parameters
  type: object
  properties:
    intent:
      title: intent
      type: string
  required:
  - intent
  definitions: null

Let’s verify that parameters: {"intent": ""}, which sets the default behavior for our Prefect Flow to incrementally process data. If necessary, update this setting, then push the changes to our Deployment with:

prefect deployment apply ./check_for_files-deployment.yaml

And then start a Prefect Agent on our local machine, which watches the Prefect Work Queue.

prefect agent start -q 'test'

Now we can see the creation of both Deployment, named Check for files/my_run and the Work Queue and in our Prefect workspace.

Our first test of the pipeline is to create a custom run named test_subset. This validates our pipeline writes data to the target S3 bucket.

Now, we run the pipeline against the full dataset, and confirm new data lands as expected.

There we have it. An automated, pure Python, big data engineering pipeline running Dask on Coiled and Prefect. From here, users should follow Prefect guidelines for operationalizing their Agent and scheduling execution of their deployment.

Reading CSV files into Dask DataFrames with read_csv

2022-02-09T00:00:00+00:00

Feb 9, 2022

6m read

Matthew Powers

Here’s how this post is organized:

Reading a single small CSV file
Reading a large CSV file
Reading multiple CSV files
Reading files from remote data stores like S3
Limitations of CSV files

Lots of datasets are stored with the CSV file format, so it’s important to understand the Dask read_csv API in detail.

Dask `read_csv`: single small file

Dask makes it easy to read a small file into a Dask DataFrame. Suppose you have a dogs.csv file with the following contents:

first_name,age
fido,3
lucky,4
gus,8

Here’s how to read the CSV file into a Dask DataFrame.

>>> import dask.dataframe as dd
>>> ddf = dd.read_csv("dogs.csv")

You can inspect the content of the Dask DataFrame with the compute() method.

>>> ddf.compute()
  first_name  age
0       fido    2
1      lucky    8
2     copper    3

This is quite similar to the syntax for reading CSV files into pandas DataFrames.

>>> import pandas as pd
>>> df = pd.read_csv("dogs.csv")

The Dask DataFrame API was intentionally designed to look and feel just like the pandas API.

For a single small file, Dask may be overkill and you can probably just use pandas. Dask starts to gain a competitive advantage when dealing with large CSV files. Rule-of-thumb for working with pandas is to have at least 5x the size of your dataset as available RAM. Use Dask whenever you exceed this limit. For example, when working on a 16 GB RAM machine, consider switching over to Dask when your dataset exceeds 3 GB in size.

Dask `read_csv`: single large file

Dask DataFrames are composed of multiple partitions, each of which is a pandas DataFrame. Dask intentionally splits up the data into multiple pandas DataFrames so operations can be performed on multiple slices of the data in parallel.

Let’s read in a 5.19 GB file (5,190 MB) into a Dask DataFrame. This file is hosted in a public S3 bucket at s3://coiled-datasets/h2o/G1_1e8_1e2_0_0/csv/G1_1e8_1e2_0_0.csv if you’d like to download it yourself.

>>> ddf = dd.read_csv("data/G1_1e8_1e2_0_0.csv")

We can run ddf.npartitions to see how many partitions the data is divided into.

>>> ddf.npartitions
81

You can customize the number of partitions that the DataFrame will contain by setting the blocksize parameter when invoking read_csv.

Dask `read_csv`: blocksize

The number of partitions depends on the value of the blocksize argument. If you don’t supply a value to the blocksize keyword, it is set to “default” and the blocksize is computed based on the available memory and number of cores on the machine, up to a maximum of 64 MB. In the example above, Dask automatically splits up the 5,190 MB data file into ~64 MB chunks when run on a MacBook Air with 8 GB of RAM and 4 cores.

We can also manually set the blocksize parameter when reading CSV files to make the partitions larger or smaller.Let’s read this data file with a blocksize of 16 MB.

>>> ddf = dd.read_csv("data/G1_1e8_1e2_0_0.csv", blocksize="16MB")
>>> ddf.npartitions
324

The Dask DataFrame consists of 325 partitions when the blocksize is 16 MB. The number of partitions goes up when the blocksize decreases. Let’s read in this data file with a blocksize of 128 MB.

>>> ddf = dd.read_csv("data/G1_1e8_1e2_0_0.csv", blocksize="128MB")
>>> ddf.npartitions
40

The Dask DataFrame has 40 partitions when the blocksize is set to 128 MB. There are fewer partitions when the blocksize increases.

Let’s take a look at how Dask infers the data types of each column when a CSV file is read. The rule of thumb when working with Dask DataFrames is to keep your partitions under 100 MB in size.

Dask `read_csv`: inferring dtypes

CSV is a text-based file format and does not contain metadata information about the data types or columns. When reading a CSV file, Dask needs to infer the column data types if they’re not explicitly set by the user.

Let’s look at the data types that Dask has inferred for our DataFrame.

>>> ddf.dtypes
id1    string[pyarrow]
id2    string[pyarrow]
id3    string[pyarrow]
id4              int64
id5              int64
id6              int64
v1               int64
v2               int64
v3             float64
dtype: object

Dask automatically uses PyArrow strings, which consume up to 70% less memory versus NumPy objects and can therefore provide a huge performance improvement (more details in Utilizing PyArrow to improve pandas and Dask workflows).

Dask infers the dtype based on a sample of the data. It doesn’t look at every row in the dataset to infer the data types because that would be prohibitively slow for large datasets.

You can increase the number of rows that are sampled by setting the sample parameter, though for this dataset, this does not change the inferred data types.

ddf = dd.read_csv("data/G1_1e8_1e2_0_0.csv", sample=5000)

Dask `read_csv`: manually specifying `dtype`

Inferring data types based on a sample of the rows can be error-prone if the dtype is not the same for all column values. A common example of this is when a column contains NaN values. You can avoid dtype inference by explicitly specifying the dtype when reading CSV files.

Let’s manually set the id4, id5, and id6 columns to a Python float, which would avoid any errors on computing our results if these columns contained NaN values.

>>> ddf = dd.read_csv(
    "data/G1_1e8_1e2_0_0.csv",
    dtype={
        "id4": float,
        "id5": float,
        "id6": float,
    },
)
>>> ddf.dtypes
id1    string[pyarrow]
id2    string[pyarrow]
id3    string[pyarrow]
id4            float64
id5            float64
id6            float64
v1               int64
v2               int64
v3             float64
dtype: object

Dask will infer the data types for the columns that you don’t manually specify. If you specify the dtype for all the columns, then Dask won’t do any data type inference.

Dask `read_csv`: multiple files

Dask can read data from a single file, but it’s even faster for Dask to read multiple files in parallel.

Let’s write out the large 5.19 GB CSV file from earlier examples as multiple CSV files so we can see how to read multiple CSV files into a Dask DataFrame. Start by writing out the single CSV file as multiple CSV files.

>>> ddf = dd.read_csv("data/G1_1e8_1e2_0_0.csv")
>>> ddf.npartitions
81
>>> ddf.to_csv("data/csvs")
['data/csvs/00.part',
 'data/csvs/01.part',
 'data/csvs/02.part',
 ...
 'data/csvs/79.part',
 'data/csvs/80.part']

This will write out 81 files, one for each partition. Let’s read these 81 CSV files into a Dask DataFrame.

ddf = dd.read_csv("data/csvs/*.part")

Reading multiple files into a pandas DataFrame must be done sequentially and requires more code. Here’s the pandas syntax:

import glob
import pandas as pd

all_files = glob.glob("./data/csvs/*.part")
df = pd.concat((pd.read_csv(f) for f in all_files))

Parallel I/O is a huge strength of Dask compared to pandas. pandas is designed for read / write operations with single files. pandas does not scale well with multiple files. Dask is designed to perform I/O in parallel and is more performant than pandas for operations with multiple files (learn more).

The same blocksize and dtype arguments we discussed earlier for reading a single file also apply when reading multiple files.

Dask readers also make it easy to read data that’s stored in remote object data stores, like AWS S3.

Dask `read_csv`: analyzing remote files with localhost compute

You can easily read a CSV file that’s stored in S3 to your local machine. Here’s how to read a public S3 file.

ddf = dd.read_csv("s3://coiled-datasets/timeseries/20-years/csv/0000.part")

You normally should not analyze remote data on your localhost machine because it’s slow to download the data locally. It’s more natural to process cloud data with cloud compute power. Files stored in AWS S3, for example, should be processed with AWS EC2 cloud provisioned compute instances.

Let’s look at a S3 folder with a lot of data. The s3://coiled-datasets/timeseries/20-years/csv/ S3 bucket has 1,095 files and requires 100 GB of memory when loaded into a DataFrame. Let’s try to run a computation on the entire dataset and see if it works.

ddf = dd.read_csv("s3://coiled-datasets/timeseries/20-years/csv/*.part")
ddf.describe().compute()

I let this computation run for 30 minutes before canceling the query. Running this locally is way too slow.

Let’s see how to read in this large dataset of CSVs to a Dask cluster that contains multiple compute nodes, so we can execute this query faster.

Dask `read_csv`: read remote files in cluster environment

Let’s spin up a 5 node cluster with Coiled and try to run the same computation with more computing power.

import coiled

cluster = coiled.Cluster(
    n_workers=5,
    region="us-east-2" # same region as data, faster + cheaper
)
client = cluster.get_client()

ddf = dd.read_csv("s3://coiled-datasets/timeseries/20-years/csv/*.part")
ddf.describe().compute()

This computation runs in just under 2 minutes and costs ~$0.08 in AWS costs. Running this computation on a cluster is certainly faster than running on localhost. It would have taken a very long time for the query to finish on localhost.

Storing the data in a different file format could make the query run even faster.

Limitations of CSV file format

CSV files are commonly used because they’re human readable and support columns with messy data types (like a column with both string and integer values, for example). However, they are usually not the best file format for a data analysis.

Here are the five reasons:

Parquet files don’t require schema inference / manual schema specification
Parquet files are easier to compress
Columnar nature of Parquet files allows for column pruning, which often yields big query performance gains
Row group metadata in Parquet files allows for predicate pushdown filtering
Parquet files are immutable

See the Dask docs for more details on the advantages of Parquet files.

Learn more

In this post we went through how to work with CSV files using Dask. For more realistic use cases you might consider:

Writing Parquet Files with Dask using to_parquet

2022-01-01T00:00:00+00:00

This blog post explains how to write Parquet files with Dask using the to_parquet method.

The Parquet file format allows users to enjoy several performance optimizations when reading data, benefiting downstream users of the data.

This blog post will teach you the basics of writing Parquet files and advanced options like customizing the filenames and controlling if the metadata file gets written. You need to be careful to avoid writing the metadata file when it can become a performance bottleneck.

Let’s dive in with a simple example.

Dask write Parquet: Small example

Let’s create a small Dask DataFrame and then write it out to disk with to_parquet.

import dask.dataframe as dd
import pandas as pd

df = pd.DataFrame(
    {"nums": [1, 2, 3, 4, 5, 6], "letters": ["a", "b", "c", "d", "e", "f"]}
)
ddf = dd.from_pandas(df, npartitions=2)

ddf.to_parquet("data/something", engine="pyarrow")

Here are the files that are output to disk.

data/something/
  _common_metadata
  _metadata
  part.0.parquet
  part.1.parquet

Dask DataFrame’s file I/O writer methods output one file per partition. This DataFrame has two partitions, so two Parquet files are written. The Parquet writer also outputs _common_metadata and _metadata files by default.

Parquet files contain metadata about the file contents including the schema of the data, the column names, the data types, and min/max values for every column in each row group. This metadata is part of what allows for Parquet files to be read in a more performant manner compared to file formats like CSV.

Dask writes out the metadata in the Parquet file footers to a single file to bypass the need to fetch the metadata from each Parquet file that’s being read. Footers are a section of Parquet files where metadata is stored. It can just look at the metadata file to get the metadata for each file in the directory. Creating the metadata file is great for small datasets, but can become a performance bottleneck when lots of files are written.

Dask to_parquet: write_metadata_file

The write_metadata_file argument is set to True by default. For large writes, it’s good to set write_metadata_file to False, so collecting all the Parquet metadata in a single file is not a performance bottleneck.

Here’s how to write out the same data as before, but without a _metadata file.

ddf.to_parquet(
    "data/something2",
    engine="pyarrow",
    write_metadata_file=False,
)

Here are the files that are output to disk.

data/something2/
  part.0.parquet
  part.1.parquet

When Dask reads a folder with Parquet files that do not have a _metadata file, then Dask needs to read the file footers of all the individual Parquet files to gather the statistics. This is slower than reading the data directly from a _metadata file. But, as mentioned previously, trying to write too much data to a single _metadata file can be prohibitively slow when lots of files are written.

In short, set write_metadata_file to False when doing large Parquet writes.

Dask to_parquet: name_function

The Parquet writes will output files with names like part.0.parquet and part.1.parquet by default. You can set the name_function parameter to customize the filenames that are written with to_parquet.

Let’s output a universally unique identifier (UUID) in the filename that’s output so we can see what files are written per each batch.

import uuid

id = uuid.uuid4()
def batch_id(n):
    return f"part-{n}-{id}.parquet"

ddf.to_parquet(
    "data/something3",
    engine="pyarrow",
    write_metadata_file=False,
    name_function=batch_id,
)

Here are the files that are output to disk.

data/something3/
part-0-09a19442-309e-485f-b006-fd9cb10f9cc7.parquet
part-1-09a19442-309e-485f-b006-fd9cb10f9cc7.parquet

It’s really nice to write files like this in your production applications. This lets you easily identify the files that are written, per batch.

You can also use the name_function to write out files with a timestamp, so they’re ordered chronologically.

Dask to_parquet: compression

Dask writes out Parquet files with Snappy compression by default. Snappy compression is typically the best for files used in distributed compute contexts. Snappy doesn’t usually compress files as much as other compression algorithms like gzip, but it’s faster when decompressing files (aka inflating files). You’ll often read compressed files from disk in big data systems, so fast inflation is important.

You can specify the Snappy compression algorithm, but this doesn’t change anything because Snappy is used by default.

ddf.to_parquet(
    "data/something4",
    engine="pyarrow",
    write_metadata_file=False,
    compression="snappy",
)

You can write out the files with the compression algorithm in the filename, which can be convenient.

def with_snappy(n):
    return f"part-{n}.snappy.parquet"

ddf.to_parquet(
    "data/something5",
    engine="pyarrow",
    write_metadata_file=False,
    compression="snappy",
    name_function=with_snappy,
)

Here are the files that are output to disk.

data/something5/
  part-0.snappy.parquet
  part-1.snappy.parquet

This is a nice way to make it clear to other humans what compression algorithm is being used by the file. Note that you will need to manually update the name_function parameter to reflect any changes made to the compression parameter.

Different columns of a Parquet file can actually be compressed with different compression algorithms. Here’s how to compress the nums column with Snappy and the letters column with gzip.

ddf.to_parquet(
    "data/something6",
    engine="pyarrow",
    write_metadata_file=False,
    compression={"nums": "snappy", "letters": "gzip"},
)

You don’t normally see different compression algorithms for different columns in a Parquet file, but it’s a cool feature that’s incredibly handy for niche use cases.

Dask write parquet: FastParquet vs PyArrow Engines

You can write Parquet files with both the FastParquet and PyArrow engines. For most use cases, either selection is fine. There are some subtle differences between the two engines, but they don’t matter for most users.

PyArrow’s popularity has taken off and it’s well supported in Dask, so PyArrow is a good option if you’re not sure which engine to use.

Dask write parquet: partition_on

You can write Parquet files in a Hive-compliant directory structure that allows for readers to leverage disk partition filtering by setting the partition_on argument.

Disk partition filtering can be a significant performance optimization for certain types of queries, but it will make other types of queries run slower. For example, queries that filter or group-by on the partition key will run a lot faster than queries operating on different keys. Queries that don’t operate on the partition key will face the drag of extra files and globbing nested directories. Globbing is particularly expensive in cloud based object stores like AWS S3.

See the post on Creating Dask Partitioned Lakes with Dask using partition_on for more information. Disk partitioning allows for huge performance gains for some query patterns, so you should know how it works.

Dask to_parquet Conclusion

Dask makes it easy to write Parquet files and provides several options allowing you to customize the operation.

You can specify the compression algorithm, choose if the metadata file should be written, customize the filename, and write the files in a nested directory structure for disk partition filtering.

Consider omitting the metadata file when you must write a lot of partitions because it can be a performance bottleneck. Including a UUID and the compression algorithm in the filename is nice when you’re building production pipelines that repeatedly write to the same directory.

You can always repartition before writing to change the number of files that are written. See this blog post for more information about the repartition method.

See this blog post to learn more about the performance benefits you can get from column pruning and predicate pushdown filtering, two optimization techniques that are afforded by Parquet files but not available in other file formats like CSV or JSON.

Why we passed on Kubernetes

2022-01-01T00:00:00+00:00

Kubernetes is great if you need to organize many always-on services and have in-house expertise, but can add an extra burden and abstraction when deploying a single bursty service like Dask, especially in a user environment with quickly changing needs.

At Coiled we chose raw cloud APIs for flexibility and simplicity, and we encourage this choice for others in similar situations. This article explains this recommendation.

Motivation

Coiled deploys Dask clusters in the cloud. We do this using raw cloud APIs rather than with Kubernetes. When new Dask users embark on resolving the deployment question they are often curious about using raw Cloud SaaS (like AWS EC2) vs Kubernetes. This article goes into some of the pros and cons behind this question.

For context, I’ve deployed Dask on many different Kubernetes systems. Kubernetes is fantastic and I often recommend it, but I’ve run into enough downstream pain that it’s no longer my default recommendation.

Deploying Dask on Kubernetes is Great

Kubernetes is great to deploy Dask for lots of different reasons:

Great Dask deployment technologies exist including dask-kubernetes, dask-gateway, the helm chart, and QHub.
Kubernetes is the de-facto standard deployment technology today. You’re buying into a system that has strong buy-in and universal support from many different technologies.
Kubernetes is cloud agnostic, and so it’s theoretically easy to shift your deployment to a different cloud vendor and avoid lock-in.
Kubernetes is battle tested and so you know that it’ll just work, or that any problems will be your fault.
Kubernetes allows you to share warm nodes across services, allowing for faster uptime when you have a large userbase.
Kubernetes can provide a single system for many services, allowing you to amortize some of your organizational cost.

In general we find that when a large organization goes all-in on Kubernetes it can provide a lot of benefits.

Dask Kubernetes is a Pain in the Butt

My experience is that setting up Dask clusters on Kubernetes initially is delightful, but that over time the Kubernetes abstraction layer provides more pain than benefit, at least when all you’re trying to do is deploy Dask. Let’s dig into why.

Excess Abstraction

Most of the upcoming issues arise because Kubernetes tries to abstract away a system (the cloud) that you don’t actually want to abstract away. Eventually you find yourself modeling much of the cloud VM / Networking stack in Kubernetes, which gets pretty painful.

Instance Types

As an example, when you set up a Kubernetes cluster you’ll back it with a node pool for which you will select a standard instance type, say, m5.4xlarge. This instance type determines:

The maximum size of a pod
The general ratio between CPU cores and memory
Local storage / disk technology used (local SSD or not)
Whether or not you have GPUs
Advanced networking that may exist

Early on, this will be fine. You don’t really care about your instance type. Dask abstracts that away for you… kind of.

Invariably you will have a sharp user who wants to use GPUs, or SSDs, or large-memory workers for cost savings, or spot instances, or ARM instances, or other aspects.

That’s ok though! You can still use Kubernetes, you’ll just set up a suite of different node pools, and use various Kubernetes labeling schemes to select between them based on the worker type requested. You can absolutely model out the complexity of cloud instance types in Kubernetes; it’s just work.

Scale Up / Scale Down

Dask workloads are bursty. Often you want to scale from zero to 100 machines as soon as possible, use them for twenty minutes, and then scale back down to zero. With cloud APIs this is easy. You ask for 100 instances, use them for a while, and then destroy them.

With Dask Kubernetes there is some interplay between Kubernetes and the underlying auto-scaling node pools. These systems are often defined by the underlying cloud infrastructure, and have various policies like “If there are excess nodes, wait a minute and then remove one”. These policies make sense for the slowly varying services for which Kubernetes was designed, but can be catastrophically costly for bursty computational services like Dask.

Again, we can reconfigure this, but we find ourselves un-engineering a system that didn’t need engineering in the first place.

Always-on

Less relevant for larger users, but for smaller groups just running the Kubernetes cluster sometimes does have some non-trivial cost. EKS has about a baseline cost of $0.10 per hour, or roughly $1,000 per year (GKE has a similar base cost). This is often less than computational costs for any corporate user, but isn’t great for individuals or smaller groups.

Maintenance

Kubernetes is great, but it does need tending. Invariably we see data professionals start down the Kubernetes path, have good early success, and then within about six months be spending about half of their time tending and modifying the cluster. This is, frankly, how many of the Dask engineers who work on distributed deployment changed careers.

The human cost is by far the largest monetary cost that we see for leveraging Kubernetes for Dask deployment.

Multi-Region

When dealing with larger organizations, or multiple organizations, one tends to need to manage multiple regions. When using raw cloud APIs it’s pretty easy to say “give me 100 workers in us-west-1 or eu-west-2”. When using Kubernetes you tend to create new Kubernetes clusters in each region, which compounds the issues above.

Do I Need Kubernetes?

In general, we find that when organizations have bought into Kubernetes, have qualified professionals maintaining it, and run many always-on services on it then it makes sense to continue using it for bursty computational services like Dask. However when organizations are first starting, don’t have lots of in-house experience, or can’t amortize the maintenance cost across many similar services then Kubernetes is rarely the correct choice long-term.

A common challenge arises because Kubernetes is pretty easy to get started with, and so people start out, but then the complexities arise late and they’re stuck with a sunken cost.

Coiled’s Decision

Given the experience above, it was obvious for Coiled, an always-on SaaS service provider that we would choose to host our computation with raw cloud APIs. Coiled provides a flexible service for a very wide variety of use cases. Our users want different regions, GPUs, Spot, ARM chips, etc. and they want it all with rapid scale-up / scale-down properties without an ongoing cost.

Of course, making a highly robust service wasn’t trivial, but now that this work is done it’s pretty easy for the world to leverage it, a fact of which we are quite proud.

Use Mambaforge to Conda Install PyData Stack on your Apple M1 Silicon Machine

2022-01-01T00:00:00+00:00

Use Mambaforge

Running PyData libraries on an Apple M1 machine requires you to use an ARM64-tailored version of conda-forge. This article provides a step-by-step guide of how to set that up on your machine using mambaforge.

Note: this article was first published in August 2021. Some of the issues with conda described below have since been resolved. I still find that mambaforge builds environments and installs packages faster and recommend trying it out.

The Apple M1 silicon chip is a fantastic innovation: it’s low-power, high-efficiency, has great battery life, and is inexpensive. However, if you’re running the PyData Stack with default builds, you may encounter some strange behaviour.

When you install a library written for Intel CPU architecture, the code will be run through the Rosetta emulator. This often works great but may cause unexpected problems for specific PyData libraries. For example, I personally encountered this with Dask’s memory usage blowing up when importing a <100MB dataset.

Since the internals of Rosetta is proprietary, there’s no way to know exactly why this is happening.

Installing Python Libraries using Mambaforge

But fear not, there is a way forward. We can resolve much of this strange behaviour by installing ARM64-native versions of the PyData libraries. While we typically recommend Anaconda/miniconda for most installations, currently, the mambaforge deployment seems to have the smoothest ARM installation process.

Below is a step-by-step guide to installing this version of conda and using it to install your favorite PyData libraries.

1. Check your version of conda

To check your conda installation, type conda info in your terminal and check the value of the platform key. This should say osx-arm64.
If it doesn’t, go ahead and remove the folder containing your conda installation. Note that this will remove your software environments, so be sure to have those saved as .yml or .txt files. You can export the ‘recipe’ for an active software environment by typing conda list — export > env-name.txt OR environment.yml into your terminal.

2. Download the mambaforge installer

Navigate to this Github repository and select the Mambaforge-MacOSX-arm64 installer. You can, of course, also opt for one of the other deployment versions listed on that page.

3. Run the installer.

This should ask you if you want it to run conda init for you. Select “yes”.
If for some reason it doesn’t do so (which happened in my case), you will have to run conda init yourself using /path/to/mambaforge/bin/conda init zsh. Note the zsh is only if you’re working with the zsh shell. Use other shell flags as needed.

After restarting your terminal, you should now be able to use conda as you always have. You can, for example, begin by rebuilding your software environments using conda env create -f .

Specific issues when running Dask, Coiled and XGBoost

Below are some further notes on issues I encountered while trying to run a workload using Dask and XGBoost on Coiled:

import coiled
from dask import Client
import dask.dataframe as dd
cluster = coiled.Cluster()
client = Client(cluster)
ddf = dd.read_parquet("<path to 100mb file on s3>")
ddf = ddf.persist()

I encountered three issues and would like to share my solutions for them as I can imagine others might be running into this as well:

Dask’s memory usage was blowing up when importing a small (<100MB) dataset with an error like the one below:

distributed.nanny — WARNING — Worker exceeded 95% memory budget. 
Restarting

In some software environments, my kernel was crashing when calling import coiled
After switching over to the mambaforge deployment, I could not install XGBoost.

1. Dask

The first issue was resolved by switching over to mambaforge and installing all the required PyData libraries from there, as outlined in the step-by-step guide above.

2. Coiled

The second issue was due to a conflict caused by the python-blosc library. For some unknown reason, kernels crash on Apple M1 machines when importing python-blosc (e.g. as part of import coiled) with a segmentation fault. To avoid this, make sure that python-blosc is not installed in your current software environment. If it is, run conda uninstall python-blosc to remove it.

3. XGBoost

Finally, If you’re looking to run XGBoost on your Apple M1 machine, you’ll need to jump through some extra hoops as there is currently no arm64-compatible version of XGBoost available on conda-forge. This is not a problem because XGBoost itself can actually run on your Apple M1 machine; the issue is with the dependencies that XGBoost installs, specifically numpy, scipy, and scikit-learn. These dependencies need to be arm64-compatible, which we’ll ensure by completing the following steps:

Create a new conda environment containing Python, numpy, scikit-learn, and scipy. It’s crucial to install these dependencies using conda and not pip here.

conda create -n boost
conda activate boost
conda install python=3.8.8 numpy scipy scikit-learn

Install the necessary libraries to compile XGBoost

conda install cmake llvm-openmp compilers

Now that we have the right dependencies in place, we can install XGBoost from pip. This is fine since we have all the right arm-64 dependencies installed already.

pip install xgboost

Take it for a spin!

Note: the steps above were summarised from this great tutorial by Fabrice Daniel and updated after helpful input from Isuru Fernando.

Moving Forward

As with any technological breakthrough, there are going to be some birthing pains as the industry recalibrates and adjusts to ripple-waves caused by Apple’s switch to M1 silicon. We’re excited to see companies like Anaconda and QuanStack and the Conda-Forge community continue to improve support for ARM hardware and stoked for the powerful future of PyData this promises!

We hope this post helps get your PyData workflows running smoothly and effectively on your Apple M1 machine. If you have any feedback or suggestions for future material, please do not hesitate to reach out to us on Twitter or our Coiled Community Slack channel.

If you’re interested in trying out Coiled, which provides hosted Dask clusters, docker-less managed software, and one-click deployments, you can do so for free today when you click below.

Understanding Managed Dask (Dask as a Service)

2022-01-01T00:00:00+00:00

‍Dask is a flexible Python library for parallel and distributed computing. There are a number of ways you can create Dask clusters, each with their own benefits. In this article, we explore how Coiled provides a managed cloud infrastructure solution for Dask users, addressing:

Deployment challenges
Software environment and credential management
Hardware flexibility
Historical tracking and observability
Cost management
Cost optimization

‍

How to deploy a Dask cluster

There are a number of tools you can use to manage Dask clusters, including:

In principle, Coiled is just like these other Dask as a service products, in that you can also use Coiled to create a Dask cluster. While this might sound straightforward, there’s significant nuance to deploying a Dask cluster effectively.

Install packages and sync credentials

On cluster start, Coiled examines your local machine, identifying pip and conda packages and any development packages you have locally. It packages these libraries and deploys them dynamically on cloud-based Linux machines in your cluster. This approach solves the software environment matching problem between your local and remote software environments. Additionally, Coiled securely transfers your local cloud credentials to the workers, ensuring they can access your cloud resources, like S3 buckets, safely and securely.

Use GPUs, ARM, and other hardware

Coiled provides a range of options for configuring your Dask cluster, making it highly flexible. You can specify the number of virtual machines with n_workers and choose a region close to your remote data with the region option. Set the number of CPU cores with worker_cpu, access GPUs with worker_gpu, or opt for more efficient ARM instances using arm=True.

import coiled

cluster = coiled.Cluster(
    n_workers=100,
    worker_cpu=8,
    arm=True,
    region="us-east-1",
    spot_policy="spot_with_fallback",
)

By contrast, if you’re using Kubernetes, you’d have to deploy a new Kubernetes cluster in this region with the node pool attached to a certain node type that has these kinds of machines in it. It can be kind of daunting.

Coiled’s cluster configuration flexibility extends to AWS, GCP, and Azure allowing you to dynamically start machines of any type in any region, a powerful feature for hardware experimentation.

Dask on AWS

Create Dask clusters on AWS by connecting Coiled to your AWS account. Coiled starts EC2 instances when you need them and turns them off when you’re done. It’s easy to set up Coiled to run from your AWS account:

$ pip install coiled
$ coiled login
$ coiled setup aws

Dask on GCP

You can also create Dask clusters on GCP by connecting Coiled to your Google Cloud account. Coiled will start virtual machines when you need them and clean up VMs when you’re done.

$ pip install coiled
$ coiled login
$ coiled setup gcp

Dask on Azure

You can also create Dask clusters on Azure by connecting Coiled to your Azure account. Coiled will start virtual machines when you need them and clean up VMs when you’re done.

$ pip install coiled
$ coiled login
$ coiled setup azure

Observability into your Dask cluster

Coiled provides visibility and historical tracking, offering insights into all activities for your team. You can monitor currently running and past clusters, with flags marking specific workload characteristics. For each cluster, you gain access to information about executed code, errors, and metrics. This visibility lets you optimize work across your team, even for members who are less familiar with the Dask codebase.

Managing Dask cloud costs

Effective cost management is crucial for organizations using cloud resources. Coiled offers team management features that enable you to limit the number of CPU cores running concurrently and on a monthly basis. You can also control team spending by monitoring usage. By default, all Coiled clusters will shutdown after 20 minutes of inactivity. These features are especially valuable for teams supporting users with varying skill levels, ensuring costs remain within budgetary constraints.

Optimizing Dask cloud costs

Coiled goes beyond cost management and helps optimize expenses in several ways, including minimizing network costs, using cost-effective non-burstable instance types as the default, and making it easy to use efficient hardware like ARM instances. There’s nuance to offering this level of configurability and also creating a seamless user experience, though. One example of this is how Coiled helps you take advantage of spot (or preemptible) instances.

Though spot instances are heavily discounted, most users are willing to pay more to avoid interruptions to their workflows as spot instances become unavailable. With Coiled, when you request spot instances, we look at your region and select the data center in that region with the best spot availability. Once we get the warning that these spot instances are being deprovisioned, we move your tasks onto new spot or on-demand instances. This ensures the continuity of your workloads without interruption. Learn more in our blog post.

In Conclusion

Coiled is a seamless managed Dask solution. It offers a smooth and secure cloud infrastructure experience for both individual users and organizations. Notably, you can use Coiled for free for the first 500 CPU hours per month, making it accessible to many users (you’ll still be responsible for your cloud provider costs).

If you are currently using other Dask as a service products, like Dask Cloud Provider or Dask Kubernetes, Coiled offers cost savings and an improved user experience, making it worth considering to manage your Dask workflows in the cloud.

It’s easy to get started with Dask on AWS, GCP, or Azure.

Tackling unmanaged memory with Dask

2022-01-01T00:00:00+00:00

TL;DR: unmanaged memory is RAM that the Dask scheduler is not directly aware of and which can cause workers to run out of memory and cause computations to hang and crash.

In this article we:

Shed light on the common error message “Memory use is high but worker has no data to store to disk. Perhaps some other process is leaking memory?”
Identify the most common causes of unmanaged memory
Use the Dask dashboard to monitor it
Learn techniques to deflate it in some of the most typical cases
Get an overview of what’s changing in how Dask deals with unmanaged memory

What is managed memory?

At any given moment, the Dask scheduler tracks chunks of data scattered over the workers of the cluster, which is typically either the input or output of Dask tasks. The RAM it occupies is called managed memory.

For example, if we write:

from dask.distributed import Client
import numpy
client = Client()
future = client.submit(numpy.random.random, 1_000_000)

Now one of the workers is holding 8 MB of managed memory, which will be released as soon as the future will be dereferenced.

What is unmanaged memory?

Managed memory never fully accounts for all of the process memory, that is the RAM occupied by the worker processes and observed by the OS.

The difference between process memory and managed memory is called unmanaged memory and it is the sum of:

The Python interpreter code, loaded modules, and global variables
sizeof() method of objects in managed memory returning incorrect results
Memory temporarily used by running tasks (sometimes improperly referred to as heap)
Dereferenced Python objects that have not been garbage-collected yet
Unused memory that the Python memory allocator did not return to libc through free() yet
Unused memory that the user-space libc free() function did not release to the OS yet (read Memory trimming below)
Memory leaks (unreachable memory in C space and/or Python objects that are never released for some reason)
Memory fragmentation (large contiguous objects - e.g. numpy arrays - failing to fit in the empty space left by deallocated smaller objects; large memory pages that are mostly free but can’t be released because a small part of them is still in use)

In an ideal situation, all of the above are negligible in size. This is not always the case, however, and unmanaged memory can accumulate to unsustainable levels. This can cause workers to run out of memory and cause whole computations to hang or crash. You may have experienced this yourself. For example, when using Dask you may have seen error reports like the following:

Memory use is high but worker has no data to store to disk.

Perhaps some other process is leaking memory?

Process memory: 61.4GiB – Worker memory limit: 64 GiB

Monitor unmanaged memory with the Dask dashboard

Since distributed 2021.04.1, the Dask dashboard breaks down the memory usage of each worker and of the cluster total:

In the graph we can see:

Managed memory in solid color (blue or, if the process memory is close to the limit, orange)
Unmanaged memory in a lighter shade
Unmanaged recent memory in an even lighter shade (read below)
Spilled memory (managed memory that has been moved to disk and no longer occupies RAM - see https://distributed.dask.org/en/latest/worker.html#spill-data-to-disk) in grey

Unmanaged recent is the portion of unmanaged memory that has appeared over the last 30 seconds. The idea is that, hopefully, most of it is due to garbage collection lag or tasks heap and should disappear soon. If it doesn’t, it will transition into unmanaged memory. Unmanaged recent is tracked separately so that it can be ignored by heuristics - namely, rebalance() - for increased stability.

By construction, process memory = managed + unmanaged + unmanaged recent.

The same information is available in tabular format on the Workers tab:

In the screenshots above, we see large amounts of unmanaged recent memory - which should hopefully help us identify what event caused it.

Investigate and deflate unmanaged memory

If you see a large amount of unmanaged memory, it would be a mistake to blindly assume it is a leak. As listed earlier, there are several other frequent causes:

Task heap

How much RAM is the execution of your tasks taking? Do large amounts of unmanaged memory appear only when running a task and disappear immediately afterwards? If so, you should consider breaking your data into smaller chunks.

Garbage collection

Do you get a substantial reduction in memory usage if you manually trigger the garbage collector?

import gc
client.run(gc.collect)  # collect garbage on all workers

In CPython (and unlike, namely, pypy and Java) the garbage collector is only involved in the case of circular references. If calling gc.collect() on your workers makes your unmanaged memory drop, you should investigate your data for circular references and/or tweak the gc settings on the workers through a Worker Plugin [https://distributed.dask.org/en/latest/plugins.html#worker-plugins].

Memory trimming

Another important cause of unmanaged memory on Linux and MacOSX, which is not widely known about, derives from the fact that the libc malloc()/free() manage a user-space memory pool, so free() won’t necessarily release memory back to the OS. This is particularly likely to affect you if you deal with large amounts of small Python objects - basically, all non-NumPy data as well as small NumPy chunks (which underlie both dask.array and dask.dataframe). What constitutes a “small” NumPy chunk varies by OS, distribution, and possibly the amount of installed RAM; it’s been empirically observed on chunks sized 1 MiB each or smaller.

To see if trimming helps, you can run (Linux workers only):

import ctypes
def trim_memory() -> int:
    libc = ctypes.CDLL("libc.so.6")
    return libc.malloc_trim(0)
client.run(trim_memory)

Here’s an example of how the dashboard can look like before and after running malloc_trim:

If you observe a reduction in unmanaged memory like in the above images, read https://distributed.dask.org/en/latest/worker.html#memtrim for how you can robustly tackle the problem in production.

What’s changing in how Dask handles unmanaged memory

Starting from distributed 2021.06.0, rebalance()[https://distributed.dask.org/en/latest/api.html#distributed.Client.rebalance] takes unmanaged memory into consideration, so if e.g. you have a memory leak on a single worker, that worker will hold less managed data than its peers after rebalancing. This is something to keep in mind if you are both a rebalance() user and are affected by trimming issues (as described in the previous section).

In the future, we will revise other heuristics that measure worker memory - namely, those underlying scatter() [https://distributed.dask.org/en/latest/locality.html#data-scatter]

and replicate() [https://distributed.dask.org/en/latest/api.html#distributed.Client.replicate] to ensure that they consider unmanaged memory.

Conclusions

Recent developments in Dask made memory usage on the cluster a lot more transparent to the user. Unexplained memory usage is not necessarily caused by a memory leak, and hopefully, this article added some tools to your kit when you need to debug it.

To learn more about Dask, visit the Dask website.

Speed up a pandas query 10x with these 6 Dask DataFrame tricks

2022-01-01T00:00:00+00:00

Updated April 18th, 2024: For Dask versions >= 2024.3.0, Dask will perform a number of the optimizations discussed in this blog post automatically. See this demo for more details.

This post demonstrates how to speed up a pandas query to run 10 times faster with Dask using six performance optimizations. You’ll often want to incorporate these tricks into your workflows so your data analyses can run faster.

After learning how to optimize the pandas queries on your local machine, you’ll see how to query much larger datasets with cloud resources using Dask clusters that are managed by Coiled.

Here are the 6 strategies covered in this post:

Parallelize the query with Dask
Avoid object columns
Split up single large file into 100 MB files
Use Parquet instead of CSV
Compress files with Snappy
Leverage column pruning

This diagram illustrates how the query time decreases as each performance optimization is applied.

After seeing how to optimize this query on a 5 GB dataset, we’ll check out how to run the same query on a 50 GB dataset with a Dask cluster. The cluster computation will show that these performance optimizations also help when querying large datasets in the cloud.

Let’s start by running a groupby operation on a 5 GB dataset with pandas to establish a runtime baseline.

Pandas query baselines

The queries in this post will be run on a 5 GB dataset that contains 9 columns. Here are the first five rows of data, so you can get a feel of the data contents.

‍

You can run this AWS CLI command to download the dataset.

aws s3 cp s3://coiled-datasets/h2o-benchmark/N_1e8_K_1e2_single.csv data/

Let’s use pandas to run a groupby computation and establish a performance baseline.

import pandas as pd

df = pd.read_csv("data/N_1e8_K_1e2_single
.csv")
df.groupby("id1", dropna=False, observed=True).agg({"v1": "sum"})

This query takes 182 seconds to run.

Here’s the query result:

‍

Let’s see how Dask can make this query run faster, even when the Dask computation is not well structured.

If you don’t want to download this dataset locally, you can also run these computations on the cloud. Later in this blog post, we’ll provide a demo that’ll show you how to use Coiled to easily spin up a Dask cluster and run these computations with machines in the cloud.

Dask query baseline

Let’s run the same groupby query with Dask. We’re going to intentionally type all the columns as object columns, which should give us a good worst case Dask performance benchmark. object columns are notoriously memory hungry and inefficient.

dtypes = {
    "id1": "object",
    "id2": "object",
    "id3": "object",
    "id4": "object",
    "id5": "object",
    "id6": "object",
    "v1": "object",
    "v2": "object",
    "v3": "object",
}

import dask.dataframe as dd

ddf = dd.read_csv("data/N_1e8_K_1e2_single.csv", dtype=dtypes)

Let’s type v1 as an integer column and then run the same groupby query as before.

ddf["v1"] = ddf["v1"].astype("int64")

ddf.groupby("id1", dropna=False, observed=True).agg({"v1": "sum"}).compute()

That query runs in 122 seconds.

Dask runs faster than pandas for this query, even when the most inefficient column type is used, because it parallelizes the computations. pandas only uses 1 CPU core to run the query. My computer has 4 cores and Dask uses all the cores to run the computation.

Let’s recreate the DataFrame with more efficient data types and see how that improves the query runtime.

Avoiding object columns with Dask

We can type id1, id2, and id3 as string type columns, which are more efficient as described in the following video.

‍

id4, id5, id6, v1, and v2 can be typed as integers. v3 can be typed as a floating point number. Here is the revised computation.

better_dtypes = {
    "id1": "string[pyarrow]",
    "id2": "string[pyarrow]",
    "id3": "string[pyarrow]",
    "id4": "int64",
    "id5": "int64",
    "id6": "int64",
    "v1": "int64",
    "v2": "int64",
    "v3": "float64",
}

ddf = dd.read_csv("data/N_1e8_K_1e2_single.csv", dtype=better_dtypes)

ddf.groupby("id1", dropna=False, observed=True).agg({"v1": "sum"}).compute()

This query runs in 67 seconds. Avoiding object type columns allows for a significant performance boost.

Split data in multiple files

Let’s split up the data into multiple files instead of a single 5 GB CSV file. Here’s code that’ll split up the data into 100 MB CSV files.

ddf.repartition(partition_size="100MB").to_csv("data/csvs")

Let’s rerun the query on the smaller CSV files.

ddf = dd.read_csv("data/csvs/*.part", dtype=better_dtypes)

ddf.groupby("id1", dropna=False, observed=True).agg({"v1": "sum"}).compute()

The query now runs in 61 seconds, which is a bit faster. Dask can read and write multiple files in parallel. Parallel I/O is normally a lot faster than working with a single large file.

Let’s try to use the Parquet file format and see if that helps.

Use Parquet instead of CSV

Here’s the code that’ll convert the small CSV files into small Parquet files.

ddf.to_parquet("data/parquet", engine="pyarrow", compression=None)

Let’s rerun the same query off the Parquet files.

ddf = dd.read_parquet("data/parquet", engine="pyarrow")

ddf.groupby("id1", dropna=False, observed=True).agg({"v1": "sum"}).compute()

This query executes in 39 seconds, so Parquet provides a nice performance boost. Columnar file formats that are stored as binary usually perform better than row-based, text file formats like CSV. Compressing the files to create smaller file sizes also helps.

Compress Parquet files with Snappy

Let’s recreate the Parquet files with Snappy compression and see if that helps.

ddf.to_parquet("data/snappy-parquet", engine="pyarrow", compression="snappy")

Let’s rerun the same query off the Snappy compressed Parquet files.

ddf = dd.read_parquet("data/snappy-parquet", engine="pyarrow")

ddf.groupby("id1", dropna=False, observed=True).agg({"v1": "sum"}).compute()

This query runs in 38 seconds, which is a bit of a performance boost.

In this case, the files are stored on my hard drive and are relatively small. When data is stored in cloud based object stores and sent over the wire, file compression can result in an even bigger performance boost.

Let’s leverage the columnar nature of the Parquet file format to make the query even faster.

Leverage column pruning

Parquet is a columnar file format, which means you can selectively grab certain columns from the file. This is commonly referred to as column pruning. Column pruning isn’t possible for row based file formats like CSV.

Let’s grab the columns that are relevant for our query and run the query again.

ddf = dd.read_parquet("data/snappy-parquet", engine="pyarrow", columns=["id1", "v1"])

ddf.groupby("id1", dropna=False, observed=True).agg({"v1": "sum"}).compute()

This query runs in 19 seconds.

Overall performance improvement

The original pandas query took 182 seconds and the optimized Dask query took 19 seconds, which is about 10 times faster.

Dask can provide performance boosts over pandas because it can execute common operations in parallel, where pandas is limited to a single core. For this example, I configured Dask to use the available hardware resources on my laptop to execute the query faster. Dask scales with your hardware to optimize performance.

My machine only has 4 cores and Dask already gives great results. A laptop that has 14 cores would be able to run the Dask operations even faster.

Parquet files allow this query to run a lot faster than CSV files. The query speeds get even better when Snappy compression and column pruning are used.

The query optimization patterns outlined in this blog are applicable to a wide variety of use cases. Splitting your data into optimally sized files on disk makes it easier for Dask to read the files in parallel easier. Strategically leveraging file compression and column pruning almost always helps. There are often easy tweaks that’ll let you run your pandas code a lot faster with Dask.

Querying a large dataset with the cloud

Let’s take a look at how to run this same query on a 50 GB dataset that’s stored in the cloud.

Downloading the 50 GB dataset to our laptop and querying it in a localhost environment would be tedious and slow. You’re limited to the computing power of your laptop when running analysis locally.

Dask is designed to either be run on a laptop or with a cluster of computers that process the data in parallel. Your laptop may only have 8GB or 32GB of RAM, so its computation power is limited. Cloud clusters can be constructed with as many workers as you’d like, so they can be made quite powerful.

Cloud devops is challenging, so let’s use Coiled to build our Dask cluster. This will make it a lot easier than building a production grade Dask solution ourselves.

Here’s how to make a Dask cluster with 15 nodes (each with 16 GiB of RAM) to query the 50 GB dataset.

Here’s how to spin up the Dask cluster with Coiled:

import coiled

cluster = coiled.Cluster(
    n_workers=15,                     
    region="us-east-2",               # run in the same region as data
    spot_policy="spot_with_fallback", # save money with spot instances
)

client = cluster.get_client()

Each instance has 16 GiB of RAM. The cluster has 15 nodes, so the entire cluster has 240 GiB of RAM. That’s a lot of computing resources to perform an analysis on a large dataset!

Once the cluster is created, we can read in data from S3 into a Dask DataFrame.

ddf = dd.read_parquet(
    "s3://coiled-datasets/h2o-benchmark/N_1e9_K_1e2_parquet",
    storage_options={"anon": True, "use_ssl": True},
    engine="pyarrow",
)

Let’s see how many rows of data the DataFrame contains.

len(ddf) # 1000000000

Wow, this DataFrame contains a billion rows of data. Let’s run the same query as before and see how long it takes:

ddf.groupby("id1", dropna=False, observed=True).agg({"v1": "sum"}).compute()

The query runs in 10 seconds. We can see that Coiled allows us to easily spin up a Dask cluster and query a large dataset quickly. It’s free and easy to get started with Coiled and run the query yourself.

Part of the reason this query is running so fast is because we’re using the query optimizations covered earlier. We’re querying snappy compressed Parquet files that each contain 100 MB of data and Dask DataFrame automatically uses column pruning to grab the relevant columns for the query.

Conclusion

pandas is a fine technology for smaller datasets, but when the data volume is large, you’ll want to use a cluster computing technology like Dask that can scale beyond the limits of a single machine.

You can manually create a Dask cluster yourself, but it’s a lot of devops work. Most data scientists and data engineers don’t want to be bothered with a challenging devops project like productionalizing and tuning a Dask environment. It’s generally best to use Coiled, a service that lets you focus on your analysis without worrying about Dask devops.

In this post you’ve seen how to significantly speed up a pandas query and you’ve also taken a glimpse of harnessing the power of the cloud to run a computation in a cluster. These tactics will serve you well when you’re running analyses on larger datasets.

Spark to Dask: The Good, Bad, and Ugly of Moving from Spark to Dask

2022-01-01T00:00:00+00:00

Apache Spark has long been a popular tool for handling petabytes of data in analytics applications. It’s often used for big data and machine learning, and most organizations use it with cloud infrastructure to run models and build algorithms. Spark is no doubt a fast analytical tool that provides high-speed queries for large datasets, but recent client testimonials tell us that Dask is even faster. So, what should you keep in mind when moving from Spark to Dask?

We interviewed Steppingblocks, an analytic organization that provides valuable insights on workforce and education outcomes based on 130 million individuals registered in US universities, employed at private organizations, and working at government agencies. Their data journey started with basic linear SQL, they moved to Spark, and later they migrated to Dask. By moving to Dask, Steppingblocks shortened their processing times from 24 hours in Spark to 6 hours in Dask, but they learned several lessons in the process. They’ve provided us (and you!) with some valuable insights on the good, bad, and ugly should you decide to migrate your organization’s machine learning and big data processing over to Dask.

Spark to Dask: What Prompted the Move?

For most organizations, migration is a huge undertaking that isn’t decided lightly. You need a return on your investment, and Steppingblocks was no different. They saw some improvements when moving from standard relational SQL to Spark, but they also found several challenges with Spark development that motivated their team to look elsewhere.

The biggest challenge was debugging code in Spark. Developers knew Python but did not know Spark, so it was a learning experience that didn’t come with much support. They found that traceback errors were not very intuitive, and they spent too much time on Stack Overflow to find answers to their problems. Support from Spark was too expensive, so any problems were left to their own devices. Finally, Spark felt like overkill for what they were doing, and it showed in their computing costs. Costs and resources were high for what Steppingblocks needed.

Also, Steppingblocks was purely a Python development team. They found that they were trying to force-fit their Python into the Spark framework. What attracted them to Dask was the more native way it worked with Python, and it made development simpler.

The Good That Led to Migration Success

Every migration takes planning and work, but not every migration is quick and painless. Before getting into Steppingblocks challenges, they revealed several successes that can be attributed to Dask. First, let’s talk about the migration successes.

The biggest win was the improvement in ETL speed. Steppingblocks processing time went from 24 hours to run their ETL procedures to 6 hours. Integration of MLFlow models with Dask was much more efficient with better feedback from the Dask system should they run into any errors from their Python code. The Coiled team was supportive and helped them work through difficult errors that developers could not figure out on their own.

Running ETLs in Spark requires cloud resources, but Dask can run locally. With Spark, you need a powerful machine only available in the cloud even for small datasets and computing tasks. Developing on Spark with datasets of any size requires access to cloud resources. By contrast, the ability for developers to run Dask locally accelerated the development lifecycle.

Overall, Steppingblocks was able to migrate 90% of their codebase from Spark to Dask in only a few weeks with a few extra months needed for optimization. The entire initial migration was done in 3 months. They shared a few success metrics below:

Timeframe: Migration took three months with two engineers and 9 months for optimization, but optimization delays were not related to Dask.
Cost: The cost to run Dask is 40% less than Spark.
Run time: Dask tasks run three times faster than Spark ETL queries and use less CPU resources.
Codebase: The main ETL codebase took three months to build with 13,000 lines of code. Developers then built the codebase to 33,000 lines of code in nine months of optimization, much of which was external library integration.

The Bad and Ugly, or Challenges during Spark to Dask Migration

Scaling Dask was a challenge at times. Steppingblocks developers could scale Dask locally but then production Dask servers would return errors. This is where support helped Steppingblocks work through their errors and figure out what was wrong. This type of support is costly for commercial companies developing with Spark for enterprise applications.

The biggest challenge was porting Spark SQL code to Dask. It’s what took the most time, including the time to optimize it. As with any new environment, developers need time to learn the new framework and grow accustomed to changes.

Migration Aftermath and the Dask Journey Since Deployment

As any developer knows, initial migration doesn’t stop when code is finally ported. You still need to maintain and reliably work with the new framework. Moving away from linear relational SQL is always a benefit for machine learning and modeling, but choosing the right platform is the challenge to ensure that you have the right tools and environment.

Steppingblocks found it much easier to migrate machine learning models and iterate them more rapidly. Their developers work in a fast-paced environment, and Dask supports their efforts. Coiled also works quickly, so developers must ensure that they are aware of the latest updates and changes to Coiled versions during testing and deployment. Whenever they have issues, however, the Coiled team helps support them.

Because Dask can work locally, developers can test their Python code in all three stages – local, local Dask, and Coiled. Moving Python modules rather than notebooks makes development iterations much faster and more reliable. Developers can build software following industry best practices including linting, unit testing, and integration testing.

Moving from Spark to Dask? Here are Lessons Learned and Tips and Tricks

If you’re thinking about moving from Spark to Dask, the Steppingblocks developers have some tips and tricks from their own lessons learned. You might find that you have your own challenges based on your unique business requirements, but one consistent story from developers is that the Coiled team was helpful in migration success and will be there to help work through any issues.

Here are a few tips before you plan out your migration from Spark to Dask:

You must run mirror development environments locally and on the cluster, or you may run into errors from version differences, but the newly released Coiled Runtime helps alleviate this issue.
Beware of defaults around partitioning and prepare to split out, but it’s useful to take advantage of map_partitions to help reduce issues.
Use worker.data in Dask for lookup tables and don’t forget to leverage caching for speed.
Don’t rely on the cloud for 100% uptime. Instead, build defensive programming techniques for retries in case of timeouts.
Leverage distributed.wait to break down large computer graphs if memory usage becomes an issue.

Watch the recording of our webinar with Sébastien Arnaud of Steppingblocks: “Why We Switched from Spark to Dask | Our Data Journey”.

Snowflake and Dask: a Python Connector for Faster Data Transfer

2022-01-01T00:00:00+00:00

Snowflake is a leading cloud data platform and SQL database. Many companies store their data in a Snowflake database.

Dask is a leading framework for scalable data science and machine learning in Python.

This article discusses why and how to use both together, and dives into the challenges of bulk parallel reads and writes into data warehouses, which is necessary for smooth interoperability. It presents the new dask-snowflake connector as the best possible solution for fast, reliable data transfer between Snowflake and Python.

Here is a video walkthrough of what that looks like:

Motivation

Data warehouses are great at SQL queries but less great at more advanced operations, like machine learning or the free-form analyses provided by general-purpose languages like Python.

That’s ok. If you’re working as a Python data scientist or data engineer, you’re probably no stranger to extracting copies of data out of databases and using Python for complex data transformation and/or analytics.

It’s common to use a database and/or data warehouse to filter and join different datasets, and then pass that result off to Python for more custom computations. We use each tool where it excels. This works well as long as it is easy to perform the handoff of large results.

However, as we work on larger datasets and our results grow beyond the scale of a single machine, passing results between a database and a general purpose computation system is challenging.

This article describes a few ways in which we can move data between Snowflake and Dask today to help explain why this is both hard and important to do well, namely:

Use pandas and the Snowflake connector
Break your query up into subqueries
Perform a bulk export to Apache Parquet

At the end, we’ll preview the new functionality, which provides the best possible solution, and show an example of use.

Three ways to pass data from Snowflake to Dask

There are different ways to accomplish this task, each with different benefits and drawbacks. By looking at all three, we’ll gain a better understanding of the problem, and what can go wrong.

1. Just use Pandas

First, if your data is small enough to fit in memory then you should just use pandas and Python. There is no reason to use a distributed computing framework like Dask if you don’t need it.

Snowflake publishes a Snowflake Python connector with pandas compatibility. The following should work fine:

pip install snowflake-connector-python

import snowflake
df = snowflake.fetch_pandas_all(...)

2. Break into many subqueries

For larger datasets, we can also break one large table-scan query into many smaller queries and then run those in parallel to fill the partitions of a Dask dataframe.

This is the approach implemented in the dask.dataframe.read_sql_table function. This function does the following steps:

Query the length and schema of the table to determine the expected size, and determine an ideal number of partitions
Query for the minimum and maximum of a column on which to split
Submit many small queries that segment that column linearly between the minimum and maximum values.

For example, if we’re reading from a time series then we might submit several queries which each pull off a single month of data. These each lazily return a pandas dataframe. We construct a Dask dataframe from these lazily generated subqueries.

import dask.dataframe as dd
df = dd.read_sql_table(
    'accounts', 
    'snowflake://user:pass@...warehouse=...role=...', 
    npartitions=10, 
    index_col='id'
)

This scales and works with any database that supports a SQLAlchemy Python connector (which Snowflake does). However, there are a few problems:

It only works for very simple queries, really just full table scans that can optionally select down columns and filter out rows. For example, you couldn’t run a big join query and then feed that result to Dask without first writing out to a temporary table
It’s kinda slow. SQLAlchemy/ODBC connectors are rarely well optimized (see Turbodbc for a good counter-example)
It’s inconsistent. Those subqueries aren’t guaranteed to run at the same time, and so if other processes are writing to your database while you’re reading you could get incorrect/inconsistent results. In some applications, this doesn’t matter. In some, it matters a great deal

For more information see this guide to Getting Started with Dask and SQL.

3. Bulk Export to Parquet

We can also perform a bulk export. Both Snowflake and Dask (and really any distributed system these days) can read and write Parquet data on cloud object stores. So we can perform a query with Snowflake and then write the output to Parquet, and then read in that data with Dask dataframe.

import dask.dataframe as dd
import snowflake
query = "
    COPY INTO 's3://my_storage_location'
    from <Table name>
    file_format = (type = parquet)
    credentials = (aws_key_id='xxxx' aws_secret_key='xxxxx' aws_token='xxxxxx');
"
con = snowflake.connector.connect(user='XXXX', passwoard='XXXX', account='XXXX',)
con.curson().execute(query)
df = dd.read_parquet('s3://my_storage_location', ...)

This is great because now we can perform large complex queries and export those results at high speed to Dask in a fully consistent way. All of the challenges of the previous approach are removed.

However, this creates new challenges with data management. We now have two copies of our data. One in Snowflake and one in S3. This is technically suboptimal in many ways (cost, storage, …) but the major flaw here is that of data organization. Inevitably these copies persist over time and get passed around the organization. This results in many copies of old data in various states of disrepair. We’ve lost all the benefits of centralized data warehousing, including cost efficiency, correct results, data security, governance, and more.

Bulk exports provide performance at the cost of the organization.

Although note, users interested in this option could also look into Snowflake Stages.

New: Managed Bulk Transfers

Recently, Snowflake has added a capability for staged bulk transfers to external computation systems like Dask. It combines the raw performance and support for complex queries of bulk exports, with the central management of directly reading SQL queries from the database.

Using this, we’re able to provide an excellent Snowflake <-> Dask data movement tool.

Understanding Snowflake’s Parallel Data Exports

Now Snowflake can take the result of any query, and stage that result for external access.

After running the query, Snowflake gives us back a list of pointers to chunks of data that we can read. We can then send pointers to our workers and they can directly pull out those chunks of data. This is very similar to the bulk export option described above, except that now rather than us being responsible for that new Parquet-on-S3 dataset, Snowflake manages it. This gives us all of the efficiency of a bulk export, while still letting Snowflake maintain control of the resulting artifact. Snowflake can clean up that artifact for us automatically so that we’re always going back to the database as the single source of truth as we should.

Snowflake + Dask Integration

Given this new capability, we then wrapped Dask around it, resulting in the following experience:

import dask_snowflake
import snowflake
with snowflake.connector.connect(...) as conn:
   ddf = dask_snowflake.from_snowflake(
      query="
      SELECT * FROM TableA JOIN TableB ON ...
      ",
      conn=conn,
   )

This is a smooth and simple approach that combines the best of all previous methods. It is easy to use, performs well, and maturely handles any query.

For more details on how to use distributed XGBoost training with Dask, see this GitHub repo.

Snowflake/Dask Fine Print

There are a few things to keep in mind:

We recommend using Dask on the same cloud and region as your Snowflake deployment for optimal speed and to avoid data ingress/egress charges.
The performance of your Dask queries will be impacted by the size and concurrency settings of your Snowflake Warehouse.
Your driving Python session (where you’re importing dask-snowflake) needs to be properly authenticated with Snowflake. Your remote workers do not. Dask-snowflake doesn’t require any additional considerations for authentication beyond what you do normally. However, because authentication information is included in the messages sent to each worker we encourage you to ensure that you use secure network connections.
The Snowflake Python connector has some known issues on Apple M1 machines. See this Github issue for more details.

Give it a spin!

Download the dask-snowflake Jupyter Notebook to get started. If you run into any issues or have thoughts about further improvements to the package, come find us on the Dask Discourse.

The technical work described in this post was done by various engineers at Snowflake, and James Bourbeau and Florian Jetter from Coiled. Watch James Bourbeau’s talk at PyData Global 2021 together with Mark Keller and Miles Adkinds from Snowflake for a detailed walkthrough.

Seven Stages of Open Software

2022-01-01T00:00:00+00:00

This post lays out the different stages of openness in Open Source Software (OSS) and the benefits and costs of each.

Motivation

“Open Source Software” is a hot term today.

As a result, people are reasonably encouraged to open up their software. This is great, but means that the term “open source” can get a bit confusing. Is Linux as open as TensorFlow? How about my personal project? Is that the same?

To help give depth to this topic, this post structures opening software into a sequence of stages of openness.

Publicly visible source code: We uploaded our code to GitHub
Licensed for reuse: And let people use it for free
Accepting contributions: And if they submit a patch, we’ll take the time to look at it, and work with them to merge it in
Open development: And when we work we’ll make sure that all of our communication happens in the open as well, so that others can see what we’re doing and why
Open decision making: And that communication will be open to the public, so that everyone can weigh in, vote, and determine what happens to the project
Multi-institution engagement: So much so that no single institution or individual has control over the project
Retirement: So now we can retire, and know that the software will live on forever

To be clear, I’m not advocating that going deeper into this hierarchy is a good thing. Often it’s more productive to stop somewhere around 3 to 5 (this is what most companies do). There are costs and benefits to every layer. We’ll now go into each layer more deeply.

1 - Publicly visible source code

Technically all that “open sourcing” your code means is that you’ve made the source code publicly visible. This might be throwing a tarball up on some website, but today it mostly means pushing a repository to GitHub.

Benefits

This is great because your code is now auditable. For example if your code handles personal data or communications people can see what is happening internally to gain confidence that everything you’re doing is secure and ethical. Or if your code runs some complex mathematical algorithm your users can inspect inside your black box and see why things behave the way that they do.

Opening up your code creates trust among your users.

Costs

But of course if you are trying to make money on this software then your competition can now gain insight.

2 - Licensed for reuse

Publishing your code doesn’t actually legally enable anyone to use it. Code is still legally protected by copyright unless the author provides a license. For proprietary software you purchase this license. Often for open source software a license is distributed with the software which lays out the ways in which anyone can use it without ever asking you for permission.

Benefits

Yay, free things

Costs

You no longer control this software. If someone makes a bunch of money off it it then you can’t complain. If someone uses it to build bombs then you also can’t complain.

3 - Accepting Contributions

So you get an e-mail, and someone who you’ve never met has found and fixed a bug for you. Great, right?

Benefits

People do work for you for free. Assuming that you depend on this software this is great, you’re getting free engineering time and a new perspective.

Costs

As this happens more and more it becomes a major time sink. You now have to educate these bug fixers, integrate all of their work, and make sure that things stay on the right track. Often you will have to reject contributions and handle the fallout of that.

These activities require a level of social engineering that you may not be familiar with or want.

4 - Open development

You spend enough time educating the people contributing to your code that you realize that it would be better if your internal team started having all of their conversations in the open as well. You push them away from in-person conversations (this is hard), out of the company Slack (this is harder), and force everyone to have all of your conversations in the open on public and searchable forums like GitHub or mailing lists.

Benefits

Yay transparency!

The community outside of your organization now has much more insight on why certain choices have been made. They really appreciate this and it builds a lot of trust around your product. Additionally, this attracts developers of a higher calibre to the project because they’re now able to operate with the full context.

Costs

This is really hard to do. It’s also pretty inefficient in the short term, especially if you already have good internal communication practices. Fully open and public communication is really time consuming to do correctly.

You’ll also be fighting an uphill battle to get your teammates to move off of internal communication, particularly if the rest of the company is there.

5 - Open decision making

Now that you’ve opened development, you have attracted some developers who fully understand the project, and have thoughts on where it should go. You listen to them and make sure that the aggregate needs are met, and not just the needs of your company. Sometimes you disagree with them, and sometimes they win these disagreements. This ends up being mostly a good thing.

You give them commit rights, the ability to push releases, and generally the ability to change things that you might not agree with.

Benefits

You get additional developer power and perspective at a completely different level. This isn’t just the occasional bugfix, this is people pouring their work-life into your project, and using it to do things that you never thought possible.

Other people in the community see that, and stop seeing your project as just a bit of corporate software, but as something that might be infrastructure some day.

Costs

You’ve given up a critical amount of control. Certainly, you still have the ability to make changes and make sure that the software meets your needs, but now you may have to negotiate from time to time.

6 - Multi-institution engagement

Those core maintainers grow to the point where they dwarf your original team of contributors. There are now multiple companies/universities/institutions that build and maintain this software. You’ve given up the majority of control in order to enable a large and diverse workforce.

Benefits

This thing has grown in maturity and scale beyond what you could have predicted

Costs

It may not grow in the way that you want. You no longer have control

7 - Retirement

You can die now. Your software will survive without you :)

Conclusion

I wrote this in response to a variety of people saying “Oh, we open sourced our software”

Setting a Dask DataFrame index

2022-01-01T00:00:00+00:00

This post demonstrates how to change a DataFrame index with set_index and explains when you should perform this operation.

Setting an index can make downstream queries faster, but usually requires an expensive data shuffle, so it should be used strategically.

This post demonstrates the set_index syntax on small datasets and then shows some query benchmarks on larger datasets.

What is an index?

Indexes are used by normal pandas DataFrames, Dask DataFrames, and many databases in general.

Indexes let you efficiently find rows that have a certain value, without having to scan each row. In plain pandas, it means that after a set_index(“col”), df.loc[“foo”] is faster than df[df.col == “foo”] was before. The loc uses an efficient data structure that only has to check a couple rows to figure out where “foo” is. Whereas df[df.col == “foo”] has to scan every single row to see which ones match.

The thing is, computers are really really fast at scanning memory, so when you’re running pandas computations on one machine, index optimizations aren’t as important.

But scanning memory across many distributed machines is not fast. So index optimizations that you don’t notice much with pandas makes an enormous difference with Dask.

A Dask DataFrame is composed of many pandas DataFrames, potentially living on different machines, each one of which we call a partition. See the Dask docs for an illustration and in-depth explanation.

The Dask client has its own version of an index for the distributed DataFrame as a whole, called divisions. divisions is like an index for the indexes—it tracks which partition will contain a given value (just like pandas’s index tracks which row will contain a given value). When there are millions of rows spread across hundreds of machines, it’s too much to track every single row. We just want to get in the right ballpark—which machine will hold this value?—and then tell that machine to go find the row itself.

So divisions is just a simple list giving the lower and upper bounds of values that each partition contains. Using this, Dask does a quick binary search locally to figure out which partition contains a given value.

Just like with a pandas index, having known divisions let us change a search that would scan every row (df[df.col == “foo”]) to one that quickly goes straight to the right place (df.loc[“foo”]). It’s just that scanning every row is much, much slower in a distributed context with Dask than on one machine with plain pandas.

When to set an index

Setting an index in Dask is a lot slower than setting an index in pandas. You shouldn’t always set indexes in Dask like you do in pandas.

In pandas, you can get away with calling set_index even when it’s not necessary. Sometimes, you’ll see pandas code like set_index(“foo”).reset_index() in the wild, which sorts the data for no reason. In pandas, you can get away with this, because it’s cheap. In Dask it’s very, very not cheap. So you really should understand what set_index does and why you’re doing it before you use it.

set_index syntax

Let’s create a small DataFrame and set one of the columns as the index to gain some intuition about how Dask leverages an index.

Create a pandas DataFrame with two columns of data, and a 2-partition Dask DataFrame from it.

import dask.dataframe as dd
import pandas as pd

df = pd.DataFrame(
     {"col1": ["aa", "dd", "cc", "bb", "aa"], "col2": ["a1", "d", "c", "b", "a2"]}
)
ddf = dd.from_pandas(df, npartitions=2)

Print the DataFrame and see that it has one index column that was created by default by pandas and two columns with data.

print(ddf.compute())

  col1 col2
 aa   a1
 dd    d
 cc    c
 bb    b
 aa   a2

Take a look at the divisions of ddf.

ddf.divisions

(0, 3, 4)

ddf has two divisions. Let’s create a print_partitions() helper method and print ddf to better illustrate the divisions.

def print_partitions(ddf):
     for i in range(ddf.npartitions):
         print(ddf.partitions[i].compute())

print_partitions(ddf)

  col1 col2
0   aa   a1
1   dd    d
2   cc    c
  col1 col2
3   bb    b
4   aa   a2

The first partition contains values with indexes from 0 to 2. The second division contains values with indexes 3 and above.

Create a new DataFrame with col1 as the index and print the result.

ddf2 = ddf.set_index("col1", divisions=["aa", "cc", "dd"])

print(ddf2.compute())


col1  col2     
aa     a1
aa     a2
bb      b
cc      c
dd      d

Let’s confirm that the divisions of ddf2 are set as we specified.

ddf2.divisions

('aa', 'cc', 'dd')

Let’s look at how the data is distributed in ddf2 by partition.

print_partitions(ddf2)

col1  col2     
aa     a1
aa     a2
bb      b

col1  col2     
cc      c
dd      d

Notice that the set_index method sorted the DataFrame based on col1. Take note of two important changes after set_index is run:

The divisions of ddf2 are now set as specified
The data is repartitioned, as per our division specifications. bb has moved to the first partition, and dd to the second, because that’s what our divisions specified.
Repeated values are all in the same partition. That is, both aa’s are in the first partition.

Sorting a Dask DataFrame can be a slow computation, see this blog post for more detail.

You will only want to run set_index when you can leverage it to make downstream queries faster.

Let’s look at a big dataset and see how much faster some queries execute when an index is set.

Reading DataFrames with index

Let’s grab a few rows of data from a 662 million row DataFrame when filtering on the index. Then, let’s run the same query when the index isn’t set to quantify the performance speedup offered by the index.

Read in the DataFrame using a 5 node Coiled cluster and verify that the timestamp is the index.

import coiled
import dask

cluster = coiled.Cluster(name="demo-cluster", n_workers=5)

client = dask.distributed.Client(cluster)

ddf = dd.read_parquet(
     "s3://coiled-datasets/timeseries/20-years/parquet",
     storage_options={"anon": True, "use_ssl": True},
     engine="pyarrow",
)

ddf.head()

Let’s take a look at what Dask knows about the divisions of this DataFrame.

ddf.known_divisions

True

You can print out all the divisions too if you’d like to inspect them.

ddf.divisions

(Timestamp('2000-01-01 00:00:00'),
 Timestamp('2000-01-08 00:00:00'),
 Timestamp('2000-01-15 00:00:00'),
 Timestamp('2000-01-22 00:00:00'),
 Timestamp('2000-01-29 00:00:00'),
 Timestamp('2000-02-05 00:00:00'))

This DataFrame has 1,096 divisions.

Here’s a query that’ll grab three records from the DataFrame and takes 2 seconds to execute.

len(ddf.loc["2000-01-01 00:00:02":"2000-01-01 00:00:04"])

Read in the same dataset without an index, so we can see how much slower the same query runs on a regular column that’s not the index.

ddf2 = dd.read_parquet(
    "s3://coiled-datasets/timeseries/20-years/parquet",
    storage_options={"anon": True, "use_ssl": True},
    engine="pyarrow",
    index=False,
)

len(
    ddf2.loc[
        (ddf2["timestamp"] >= "2000-01-01 00:00:02")
        & (ddf2["timestamp"] <= "2000-01-01 00:00:04")
    ]
)

That query takes 115 seconds to execute because it had to load and scan all of the data looking for those values, instead of jumping straight to the correct partition. The query without the index is 58 times slower in this example.

Filtering on an index value is clearly a lot quicker than filtering on non-index columns.

How long does it take to set the index?

Let’s look at a query to compute the number of unique names in the data when the id equals 1001. We’ll see if setting an index can improve the performance of the query.

Start by reading in the data and running a query to establish a baseline.

ddf3 = dd.read_parquet(
     "s3://coiled-datasets/timeseries/20-years/parquet",
     storage_options={"anon": True, "use_ssl": True},
     engine="pyarrow",
     index=False,
)

ddf3.loc[ddf3["id"] == 1001].name.nunique().compute()

This query takes 75 seconds to execute.

dd3 has 1,095 partitions, which you can see by running ddf3.npartitions.

Setting an index before running the query

Let’s set an index and run the same query to see if that’ll make the query run faster.

dask_computed_divisions = ddf3.set_index("id").divisions
unique_divisions = list(dict.fromkeys(list(dask_computed_divisions)))
len(unique_divisions) # 149
ddf3u = ddf3.set_index("id", divisions=unique_divisions)
ddf3u.npartitions # 148
ddf3u.loc[[1001]].name.nunique().compute()

This query takes 130 seconds to run. Setting the index before running the query is slower than simply running the query. It’s not surprising that this runs slower. Setting the index requires a data sort, which is expensive.

This example just shows one way you can create the divisions that are used when the index is set. You can create the divisions however you’d like. Choosing the optimal divisions to make downstream queries faster is challenging.

It’s not surprising that this particular index does not improve the overall query time: we’re sorting all of the values, but we only needed to extract just one. It’s faster to just scan the un-sorted data and find the single value we’re interested in.

Let’s look at some more examples to gather some intuition and then look and when setting an index is more likely to improve query times.

Setting an index without specifying divisions

Let’s look at setting an index without explicitly setting divisions and running the same count query.

ddf3a = ddf3.set_index("id") # 134 seconds
ddf3a.loc[[1001]].name.nunique().compute() # 118 seconds

This takes a total of 252 seconds to compute and is slower than the original query. In this case, setting an index doesn’t help our query run faster.

Notice that ddf3.set_index(“id”) does not specify divisions, so Dask needs to go figure them out. To do this, it has to load and compute all of ddf3 immediately to look at the distribution of its values. Then, when you later call .compute(), it’ll load and compute ddf3 a second time. This is slow in general, and particularly bad if the DataFrame already has lots of operations applied to it—all those operations also have to run twice.

When divisions are passed, Dask doesn’t need to compute the whole DataFrame to figure them out, which is obviously a lot faster.

Setting an index and specifying divisions poorly

Now let’s try to manually set the divisions and run the same query.

ddf3b = ddf3.set_index("id", divisions=[800, 900, 1000, 1100, 1200]) # 0.3 seconds
ddf3b.loc[[1001]].name.nunique().compute() # didn't work

This didn’t work because we’re taking a DataFrame that used to have 1,095 partitions and trying to smash all that data into just 4 partitions. This runs out of memory. Clearly, using bad divisions is worse than the extra time for Dask to compute them for you.

Not setting the divisions well can cause problems. You want to pick divisions that:

Separate the data into approximately equally-sized groups. If some values are more common than others, adjust partitions accordingly.
Result in reasonably-sized partitions—around 100MB is typically a good target, which often means row counts around 100,000 per partition.

To pick good divisions, you must use your knowledge of the dataset. What range of values is there for the column? What sort of general distribution do they follow—a normal bell curve, a continuous distribution? Are there known outliers?

Another strategy is to let Dask compute the divisions once, then copy-paste them to reuse later:

dask_computed_divisions = ddf3.set_index("id").divisions
unique_divisions = list(dict.fromkeys(list(dask_computed_divisions)))
print(repr(unique_divisions))
# ^ copy this and reuse

This is especially helpful if you’ll be rerunning a script or notebook on the same (or similar) data many times.

However, you shouldn’t set divisions if the data you’re processing is very unpredictable. In that case, it’s better to spend the extra time and let Dask re-compute good divisions each time.

High level index strategies

Since set_index is an expensive operation, you should only run the computation when it’ll help subsequent computations run faster.

Here are the main operations that’ll run faster when an index is set:

Filtering on the index
Joining on the index
Groupby on the index (soon! dask/dask#8442)
Sophisticated custom code in map_partitions (advanced use case)

By all means, set indexes whenever it’ll make your analysis faster. Just don’t run these expensive computations unnecessarily.

For example, you shouldn’t always set_index before you .loc. If you just need to pull a value out once, it’s not worth the cost of a whole shuffle. But if you need to pull lots of values out, then it is. Same with a merge: if you’re just merging a DataFrame to another, don’t set_index first (the merge will do this internally anyway). But if you’re merging the same DataFrame multiple times, then the set_index is worth it.

As a rule of thumb, you should only set_index if you’ll do a merge, groupby(df.index), or .loc on the re-indexed DataFrame more than once.

You may also want to re-index your data before writing it to storage in a partitioned format like Parquet. That way, when you read the data later, it’s already partitioned the way you want, and you don’t have to re-index it every time.

Conclusion

set_index changes the index of a DataFrame, sorts the underlying data, and ensures the data is partitioned in a manner consistent with the divisions. Divisions allow for you to cheaply look up values without having to process every partition to search for them.

Sorting is an expensive operation when data is distributed, but it can be worth it if you’re performing multiple operations that run faster on data with an index.

You may also want to sort data before persisting it on disk, so Dask can read it with an index and you don’t need to perform the sort later.

Keep this powerful performance optimization in mind when you’re performing analyses with Dask.

Search at Grubhub and User Intent

2022-01-01T00:00:00+00:00

Alex Egg, Senior Data Scientist at Grubhub, joins Matt Rocklin and Hugo Bowne-Anderson to talk and code about how Dask and distributed compute are used throughout the user intent classification pipeline at Grubhub!

Context: Search is the top-of-funnel at Grubhub. That means when a user interacts with the Grubhub search engine, they want to be able to service their request with high precision and recall. One way to do that is to understand the intent of the user. Do they have a favourite restaurant in mind or are they just browsing cuisines? Are they looking for an obscure dish? Depending on the intent of the user, Grubhub can optimize the backend query or apply UI treatments to help the user find just what they had in mind. In this session, we’ll take a look at how to build an end-to-end search query intent classifier from data labeling to serving live requests in production.

We’ll cover:

Classic ETL pipelines with Dask Dataframes,
Weak-supervision (labeling) w/ Snorkel & Dask (what a modern combo!),
Language Modeling and deployment w/ Tensorflow.

Scikit-learn in the Cloud with Coiled

2022-01-01T00:00:00+00:00

You can train scikit-learn models in parallel using the scikit-learn joblib interface. This allows scikit-learn to take full advantage of the multiple cores in your machine (or, spoiler alert, on your cluster) and speed up training.

Using the Dask joblib backend, you can maximize parallelism by scaling your scikit-learn model training out to a remote cluster.

This works for all scikit-learn models that expose the n_jobs keyword, like random forests, linear regressions, and other machine learning algorithms. You can also use the joblib library for distributed hyperparameter tuning with grid search cross-validation.

Scikit-learn Models vs Data Size

When training machine learning models, you can run into 2 types of scalability issues: your model size may get too large, or your data size may start to cause issues (or even worse: both!). You can typically recognize a scalability problem related to your model size by long training times: the computation will complete (eventually), but it’s becoming a bottleneck in your data processing pipeline. This means you’re in the top-left corner of the diagram below:

Using the scikit-learn joblib integration, you can address this problem by processing your scikit-learn models in parallel, distributed over the multiple cores in your machine.

Note that this is only a good solution if your training data and test data fit in memory comfortably. If this is not the case (and you’re in the “I give up!” quadrant of large models and large datasets), see the “Dataset Too Large?” section below to learn how Dask-ML or XGBoost-on-Dask can help you out.

Sklearn + Joblib: Run Random Forest in Parallel Locally

Let’s see this in action with an scikit-learn algorithm that is embarrassingly parallel: a random forest model. A random forest model fits a number of decision tree classifiers on sub-samples of the training data , then combines the results from the various decision trees to make a prediction. Because each decision tree is an independent process, this model can easily be trained in parallel.

First, we’ll create a new virtual environment for this work. (If you’re following along, you can also use pip or any other package manager, or your own existing environment.)

conda create -n sklearn-example -c conda-forge python=3.9 scikit-learn joblib ipython
conda activate sklearn-example

We’ll import the necessary libraries and create a synthetic classification dataset.

from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import make_classification
X, y = make_classification(n_samples=1000, n_features=4,
                          n_informative=2, n_redundant=0,
                          random_state=0, shuffle=False)

Then we’ll instantiate our Random Forest classifier

clf = RandomForestClassifier(max_depth=2, random_state=0)

We can then use a joblib context manager to train this classifier in parallel, specifying the number of cores we want to use with the n_jobs argument.

import joblib
with joblib.parallel_backend(backend='loky', n_jobs=4):
    clf.fit(X, y)

To save users from having to use context managers every time they want to train a model in parallel, scikit-learn’s developers exposed the n_jobs argument as part of the instantiating call:

clf = RandomForestClassifier(max_depth=2, random_state=0, n_jobs=-1)

The n_jobs keyword communicates to the joblib backend, so you can directly call clf.fit(X, y) without wrapping it in a context manager. This is the recommended approach for using joblib to train sklearn models in parallel locally.

Running this locally with n_jobs = -1 on a MacBook Pro with 8 cores and 16GB of RAM takes just under 3 minutes.

%%time
clf.fit(X,y)
CPU times: user 13min 21s, sys: 17.8 s, total: 13min 38s
Wall time: 2min 6s

Sklearn + Joblib: Run RF on the Cluster

The context manager syntax can still come in handy when your model size increases beyond the capabilities of your local machine, or if you want to train your model even faster. In those situations, you could consider scaling out to remote clusters on the cloud. This can be especially useful if you’re running heavy grid search cross-validation, or other forms of hyperparameter tuning.

You can use the Dask joblib backend to delegate the distributed training of your model to a Dask cluster of virtual machines in the cloud.

We first need to install some additional packages:

conda install -c conda-forge coiled dask xgboost dask-ml

Next, we’ll launch a Dask cluster of 20 machines in the cloud.

When you create the cluster, Coiled will automatically replicate your local sklearn-example environment in your cluster (see Package Sync).

import coiled
cluster = coiled.Cluster(
    n_workers=20,
)

We’ll then connect Dask to our remote cluster. This will ensure that all future Dask computations are routed to the remote cluster instead of to our local cores.

# point Dask to remote cluster
client = cluster.get_client()

Now use the joblib context manager to specify the Dask backend. This will delegate all model training to the workers in your remote cluster:

%%time
with joblib.parallel_backend("dask"):
    clf.fit(X, y)

Model training is now occurring in parallel on your remote Dask cluster.

This runs in about a minute and a half on my machine—that’s a 2x speed-up. You could make this run even faster by adding more workers to your cluster.

CPU times: user 1.93 s, sys: 601 ms, total: 2.53 s
Wall time: 1min 1s

You can also use your Dask cluster to run an extensive hyperparameter grid search that would be too heavy to run on your local machine:

from sklearn.model_selection import GridSearchCV

# Create a parameter grid
param_grid = {
    'bootstrap': [True],
    'max_depth': [80, 90, 100, 110],
    'max_features': [2, 3],
    'min_samples_leaf': [3, 4, 5],
    'min_samples_split': [8, 10, 12],
    'n_estimators': [100, 200, 300, 1000]
}

# Instantiate the grid search model
grid_search = GridSearchCV(
    estimator=clf, 
    param_grid=param_grid, 
    cv=5, 
    n_jobs=-1
)

Before we execute this compute-heavy Grid Search, let’s just scale up our cluster to 100 workers to speed up training and ensure we don’t run into memory overloads:

cluster.scale(100)
client.wait_for_workers(100)

Now let’s execute the grid search in parallel:

%%time
with joblib.parallel_backend("dask"):
    grid_search.fit(X, y)

Dataset too large?

Is your dataset becoming too large, for example because you have acquired new data? In that case, you may be experiencing two scalability issues at once: both your model and your dataset are becoming too large. You typically notice an issue with your dataset size by pandas throwing a MemoryError when trying to run a computation over the entire dataset.

Dask exists precisely to solve this problem. Using Dask, you can scale your existing Python data processing workflows to larger-than-memory datasets. You can use Dask DataFrames or Dask Arrays to process data that is too large for pandas or NumPy to handle.

Dask also scales machine learning for larger-than-memory datasets. Use dask-ml or the distributed Dask backend to XGBoost to train machine learning models on data that doesn’t fit into memory. Dask-ML is a library that contains parallel implementation of many scikit-learn models.

The code snippet below demonstrates training a Logistic Regression model on larger-than-memory data in parallel.

from dask_ml.linear_model import LogisticRegression
from dask_ml.datasets import make_classification

X, y = make_classification(chunks=50)
lr = LogisticRegression()
lr.fit(X, y)

Dask also integrates natively with XGBoost to train XGBoost models in parallel:

import xgboost as xgb

# create DMatrix
dtrain = xgb.dask.DaskDMatrix(client, X, y)

# train model
output = xgb.dask.train(
   client, params, dtrain, num_boost_round=4,
   evals=[(dtrain, 'train')]
)

Scikit-learn + Joblib Summary

We’ve seen that:

You can use the scikit-learn’s joblib integration to distribute certain scikit-learn tasks over all the cores in your machine for a faster runtime.
You can connect joblib to the Dask backend to scale out to a remote cluster for even faster processing times.
You can use XGBoost-on-Dask and/or dask-ml for distributed machine learning training on datasets that don’t fit into local memory.

Scale your data science workflows with Python and Dask

2022-01-01T00:00:00+00:00

Data Scientists are increasingly using Python and the Python ecosystem of tools for their analysis. Combined with the growing popularity of big data, this brings the challenge of scaling data science workflows. Dask is a library built for this exact purpose - making it easy to scale your Python code, and serve as a toolbox for distributed computing!

In this post, we will discuss:

Why use Python for Data Science?
Why is there a need for scalability?
How does Dask make scaling simple, and where is it used?

Python for Data Science

Python’s wide adoption is attributed to three key factors:

It has a more readable and human-friendly syntax when compared to a lot of other programming languages.
It has a free and open source foundation, which allows more people to learn, teach, and improve the language.
It has an active and growing community around it. There are numerous supporting tools and libraries that make Python one of the most powerful languages to use, including for data science.

Python has made Data Science more accessible to everyone from researchers to analysts to students:

Libraries like NumPy, pandas, scikit-learn, XGBoost, Matplotlib, etc. are free and open source alternatives to proprietary data science tools like MATLAB.
Tools like Dask also help reduce reliance on high-cost infrastructure like supercomputers for scalability.
The Jupyter ecosystem makes collaboration easy by providing executable and shareable environments.
New specialized libraries are being developed everyday, adding more layers of abstraction, and speaking the language of a specific community like Astropy for Astrophysics, SunPy for Solar physics, Biopython for computational biology, and many more.

The Scientific Python community is growing rapidly, hence we also have a growing pool of resources to support new users!

Scaling Data Science Workflows

Recent advancements in technology, algorithms, and computational power has led to a “big data revolution”. Data scientists require scalable solutions to work with these large and complex datasets, and they look towards parallel and distributed computing.

Scalable computing can be anything that extends beyond a single thread on a single core. Parallel computing is the process of computing multiple tasks simultaneously to accelerate workflows, and distributed computing is the process of leveraging computing resources of multiple machines.

Another important concept in scalable computing is cloud computing. Providers like AWS, Azure, and GCP, provide a lot of computing resources that can be accessed from anywhere. Cloud computing involves leveraging these resources to perform big computations.

Scaling made simpler with Dask

Dask is a library for parallel and distributed computing in Python. Dask makes it easy to use all resources on your local machine, set up distributed computing environments, as well as scale to the cloud. It’s familiar API, flexible design, and synergy with the Python ecosystem make Dask the tool of choice for both individuals and institutions.

Dask is used by industry leaders from Walmart and Capital One, to government agencies like NASA and the UK Met Office, among many others to scale their data science and machine learning pipelines. Dask is also used behind-the-scenes in many popular tools including Nvidia RAPIDS, Apache Airflow, and PyTorch!

Dask is a very powerful library, but scaling Dask computations to the cloud involves many DevOps challenges like networking, security, and environment configurations. Coiled Cloud is a product built by the creators and maintainers of Dask, that takes care of these challenges, so that data scientists can focus on their analysis. Coiled lets you scale to the cloud in just one-click, while also providing essential features for team and cost management.

Save Money with Spot

2022-01-01T00:00:00+00:00

The cloud is wonderful but expensive.

Spot/preemptible instances offer dramatic cost savings (2–3x), but only if you can…

Find enough in your region
Handle their sudden disappearance

Doing this well requires considerable nuance which we explore in this post.

We optimized this process while making Coiled, a cloud SaaS product around deploying Python distributed computing with Dask. We’ll use Coiled/Dask workloads as a running example, but these lessons should apply broadly.

Spot is Cheap

Both AWS and Google Cloud let you buy idle compute resources at a significant discount. For example, here are the costs for a 2vCPU 8GiB machine on each cloud:

Spot instances are about a quarter the cost… so what’s the catch?

Because spot instances are excess capacity, they can be harder to get and harder to keep. Example issues:

You might request 100 spot workers and only get 80
You might get 100, but after your cluster is up and running you might lose some

For some workloads, this is fine. For others, less so. Which raises the question:

How can we get as much spot as possible, while also delivering a consistent user experience?

After iterating on this problem with our users, we’ve come up with an approach that we’ve found to work well:

Select the best zone for getting spot instances
Use on-demand instances when spot is unavailable
Automatically replace spot instances with changes in cloud provider availability

When you spin up a cluster using Coiled, you can now enable these features with these keyword arguments:

import coiled

cluster = coiled.Cluster(
    use_best_zone=True, 
    spot_policy="spot_with_fallback",
    ...,
)

1 – Choose the right availability zone

Before we worry about spot/preemptible instances going away, first we need to find enough spot. To do this well, we want to choose the data center or availability zone with the most excess capacity.

Clouds (AWS, GCP, Azure) have lots of data centers. These data centers are grouped into “zones”, and one or more zones make up a “region”. Usually you care about what region you use, because the sorts of workloads that people run on Coiled usually involve processing lots of data, and moving data between regions costs money.

You should also care that all of your instances are in a single zone, because moving data between zones also costs money, but usually you don’t need to care which specific zone you use. Data is typically stored in a region and not a specific zone within that region.

Different availability zones have different capacity. This applies to Spot, and also other rare instance types like GPUs and TPUs. By being flexible to the right availability zone we can often get more of the thing we want, while still being local to our data.

Example

We applied this approach to internal workloads, large scale Dask benchmarks run by our in-house OSS teams. These regular jobs spin up many hundreds of workers for benchmarking Dask performance at scale. In the first two weeks of December, we spun up over 52000 m6i.large spot workers. Out of all those requests for workers, there were exactly 60 times when we failed to get the spot instance but would have been able to get the worker as an on-demand instance.

As availability shifted, you can see that these workers came from different AWS availability zones:

Prior to this feature, we didn’t even use Spot. Consistency in these benchmarks was important enough to our internal teams, and the experience was poor enough that we just paid the higher bill. Since implementing this feature (and those below) we’ve heard no complaints from engineering, and nothing but praise from finance.

2 – Spot with Fallback to On-Demand

Even when choosing the best zone, sometimes you still hit limits. This is especially common when you want many hundreds of instances, or when you want very large instances, or when you want GPUs. All of these have limited spot availability.

Sometimes this is fine. Sometimes it isn’t. Sometimes you need things to run predictably with the requested number of workers at the requested time. Our solution to this is to use spot but fall back to using on-demand if spot isn’t available. We offer three different policies

Spot: If you’re very cost conscious and never want to pay the full price for your compute, we’ll only give you spot instances.
Spot-with-fallback: good if you’re using spot instances but it’s also important to you that you get clusters of a certain size and are willing to pay full price as necessary.
On-demand: which we’ve pretty much stopped using, but which you could use if you want to be very very sure that everything is stable.

We’ve switched our internal Dask benchmarking to use spot instances with fallback. As I mentioned before, we get spot instances about 99.9% of the time—but for the remaining 0.1% we get on-demand fallback instances so that there’s no impact on our benchmarks from the ups and downs of spot availability.

On-demand is currently the default, we’re strongly considering making spot-with-fallback the default.

3 – Instance replacement

Are there any downsides to using spot with fallback compared to using only on-demand instances?

Unfortunately there are.

Spot instances aren’t just harder to get, they’re also harder to keep—the cloud provider can “interrupt” (AWS) or “preempt” (Google Cloud) your instances when they need to free up availability for other customers.

So how do we address this?

Coiled relies on Dask for scheduling work on your cluster. Dask’s architecture is to have a single scheduler (we run this on its own on-demand instance) and potentially very many workers (this is where you get scale). The scheduler needs to stay up for the entire duration of your workload, but Dask is resilient to workers coming and going—this makes it easy to scale up and down in response to changes in your workload or changes in cloud provider instance availability.

We always use on-demand for the scheduler even when there’s available spot capacity. On the one hand, that means you’re paying full price for the scheduler instance. But on the other hand, the scheduler rarely needs to be a large instance—for our benchmarking on AWS we use an m6i.large—and there’s only one scheduler, whereas there are usually many workers.

For workers we watch for advanced notification that a spot worker is going to be terminated, and we attempt to then gracefully shut it down so that computed results can be transferred to other workers. At the same time, we start spinning up another instance—which may be a spot instance, but if none are available and you’ve enabled fallback, it may be an on-demand instance. Dask is able to handle worker replacement midstream, so your workloads are able to keep running.

Smooth User Experience for Cost Savings

Coiled makes it easy to quickly and reliably spin up clusters in the cloud to run your Python code in parallel at scale. Because we remove the hassle of manually provisioning and configuring servers and infrastructure, it’s trivial for you to get a cluster with hundreds of cores and terabytes of memory that you might only need for a few minutes.

Before making the changes above, we found that while spot provided substantial cost savings in theory, it wasn’t really usable. Users kept not getting the machines they needed, or their workloads would crash unpredictably. In practice folks stopped using spot and kept on with on-demand instances. Today folks are pretty happy. We use the default configuration in our own testing suites and our internal bills have gone down. Same with many of our customers.

You can see our docs for the relevant keyword arguments when creating Coiled clusters.

Repartitioning Dask DataFrames

2022-01-01T00:00:00+00:00

This article explains how to redistribute data among partitions in a Dask DataFrame with repartitioning…

Dask DataFrames consist of partitions, each of which is a pandas DataFrame. Dask performance will suffer if there are lots of partitions that are too small or some partitions that are too big. Repartitioning a Dask DataFrame solves the issue of “partition imbalance”.

Let’s start with some simple minimal complete verifiable examples (MCVE) repartition examples to get you familiar with the repartition syntax.

We’ll then illustrate the performance gains from repartitioning a 63 million row dataset. This example demonstrates how the small file problem causes a computation to run 18 times slower and how this can be fixed with repartitioning.

Simple examples

Let’s create a Dask DataFrame with six rows of data organized in three partitions.

import pandas as pd
import dask.dataframe as dd

df = pd.DataFrame({"nums":[1, 2, 3, 4, 5, 6], "letters":["a", "b", "c", "d", "e", "f"]})
ddf = dd.from_pandas(df, npartitions=3)

Print the content of each DataFrame partition.

for i in range(ddf.npartitions):
    print(ddf.partitions[i].compute())

 nums letters
0     1       a
1     2       b
 nums letters
2     3       c
3     4       d
 nums letters
4     5       e
5     6       f

Repartition the DataFrame into two partitions.

ddf2 = ddf.repartition(2)

for i in range(ddf2.npartitions):
    print(ddf2.partitions[i].compute())

 nums letters
0     1       a
1     2       b
 nums letters
2     3       c
3     4       d
4     5       e
5     6       f

repartition(2) causes Dask to combine partition 1 and partition 2 into a single partition. Dask’s repartition algorithm is smart to coalesce existing partitions and avoid full data shuffles.

You can also increase the number of partitions with repartition. Repartition the DataFrame into 5 partitions.

ddf5 = ddf.repartition(5)

for i in range(ddf5.npartitions):
    print(ddf5.partitions[i].compute())

 nums letters
0     1       a
 nums letters
1     2       b
 nums letters
2     3       c
 nums letters
3     4       d
 nums letters
4     5       e
5     6       f

In practice, it’s easier to repartition by specifying a target size for each partition (e.g. 100 MB per partition). You want Dask to do the hard work of figuring out the optimal number of partitions for your dataset. Here’s the syntax for repartitioning into 100MB partitions.

ddf.repartition(partition_size="100MB")

Example of small file problem

Let’s illustrate the small file problem and then explain how repartitioning can eliminate an excessive number of small files.

Suppose you have two years of data, one row per second.

That’s 63 million rows of data.

Let’s run a query on a 5 node Dask cluster powered by Coiled.

When the data is stored in 104 CSV files (41 MB each), it takes 16.6 seconds to get the number of unique values in the name column.

When the data is stored in 17,496 tiny CSV files (248 KB each), it takes a whopping 4 minutes and 16 seconds to perform the same computation. Having too many files causes the computation to run 15 times slower!

When the tiny CSV files are repartitioned into 100 partitions and written out as Parquet files, then the count operation can be performed in 4 seconds. Tiny files should be eliminated via repartitioning, so Dask can perform computations more efficiently.

See this notebook for the full benchmarking analysis with computations to support the results. Here’s the conda environment you can use to run the notebook, as described in the project README.

When to repartition

Of course, repartitioning isn’t free and takes time. The cost of performing a full data shuffle can outweigh the benefits of subsequent query performance.

You shouldn’t always repartition whenever a dataset is imbalanced. Repartitioning should be approached on a case-by-case basis and only performed when the benefits outweigh the costs.

Common causes of partition imbalance

Filtering is a common cause of DataFrame partition imbalance.

Suppose you have a DataFrame with a first_name column and the following data:

Partition 0: Everyone has a first_name “Allie”
Partition 1: Everyone has first_name “Matt”
Partition 2: Everyone has first_name “Sam”

If you filter for all the rows with first_name equal to “Allie”, then Partition 1 and Partition 2 will be empty. Empty partitions cause inefficient Dask execution. It’s often wise to repartition after filtering.

Next steps

This post has shown you how to fix DataFrame partition imbalance with repartitioning.

Repartitioning can be costly because it requires data to be shuffled. The overall performance gains from repartitioning should be measured on a case-by-case basis. The performance gains from having data evenly distributed across DataFrame partitions may be outweighed by the cost of performing the shuffle.

You’re well equipped to determine when it’s best to run this important performance optimization after learning about repartitioning in this post.

Thanks for reading. If you’re interested in trying out Coiled, which provides hosted Dask clusters, docker-less managed software, and one-click deployments, you can do so for free today when you click below.

Reducing memory usage in Dask workloads by 80%

2022-01-01T00:00:00+00:00

There’s a saying in emergency response: “slow is smooth, smooth is fast”.

That saying has always bothered me, because it doesn’t make sense at first, yet it’s entirely correct.

By applying this philosophy to the scheduling algorithm in the latest release of Dask (2022.11.0), we’re seeing common workloads use up to 80% less memory than before. This means some workloads that used to be outright un-runnable are now running smoothly—an infinity-X speedup!

Cluster memory use on common workloads—blue is the latest release. The blue lines are lower by a factor of “a lot”.

To see this in action, here’s a demo by Florian from last month’s Dask Demo Day:

The second-most upvoted and commented issue of all time on the dask/distributed repo describes, “tasks early in my graph generate data faster than it can be consumed downstream, causing data to pile up, eventually overwhelming my workers”.

Dask users often struggle with workloads that run out of memory like this. Studying these situations, we realized that the Dask scheduler wasn’t following this “slow is smooth, smooth is fast” adage.

Here’s what the problem was, and how we’ve addressed it:

Root task overproduction

Historically, Dask tried hard to get each individual task done as fast as possible: if a task could run, it would. So sometimes, tasks would run even if their outputs weren’t going to be used immediately—leaving them sitting around in memory.

If you had thousands of initial tasks loading data—say, fetching Parquet from S3, CSVs from disk, or rows from a database—all those tasks would be scheduled and sent to workers up front.

The workers would churn through them, fetching chunks of data (and accumulating it in memory) as quickly as possible. A worker would tell the scheduler when each chunk was loaded, and the scheduler would respond with what to do with it next—but until that message arrived, there were more data-loading tasks runnable right now, so why not run them?

This slight gap in timing—between the worker immediately starting on a less-useful task, then only later finding out about a more-useful task it should have run instead—allowed this lower-priority data to pile up in memory. We call this “root task overproduction”.

Overall, there could be at least twice as much initial data in memory at once as necessary—and, therefore, twice as many intermediate results. (See this comment for a detailed explanation of why the 2x happens.)

When this put workers under memory pressure, this initial problem would snowball into a bigger one. Workers had to spill data to disk (slow), then read it back from disk to use it or transfer it (slow). Workers might exceed their memory limits and crash, losing progress and requiring tasks to be recomputed on a pool of workers that were already struggling.

In the end, this meant that a whole class of workloads were slow, or even un-runnable, without using hugely oversized clusters.

Queuing

There were plenty of ways to approach this problem, but we wanted to try the simplest thing first: just don’t tell the workers about more tasks than they can run at once.

We’re calling this mode of scheduling “queuing”, or “root task withholding”. The scheduler puts data-loading tasks in an internal queue, and only drips one out to a worker once it’s finished its current work and there’s nothing more useful to run instead that utilizes the work it just completed.

Slow is smooth

Queuing adds a cost in latency. Every time workers finish a task, they have to ask the scheduler what to do next and sit under-utilized, or even idle, until they get an answer. (Before, they had a backlog of things they could do before the answer came back.)

For a while, we hadn’t considered this approach, because intuitively, we assumed the latency would be too much of a slow-down.

However, by slowing down the pace of task assignment, and running only the best tasks, scheduling gets much smoother. And with that smoothness, we see that most benchmark workloads use much less memory across the board:

Percent decrease in peak memory use in the latest release. Notice the axes: up to 80% reduction.

This large a reduction in memory use is a big deal!

For many users, this might mean that workloads that weren’t runnable before can now run smoothly.

By setting worker saturation, I can reliably calculate a trend over time for the first time 🎉

– Climate scientist Julius Busecke reporting that a common geoscience task which used to always crash now works out of the box.

Some workloads will also be cheaper to run in the cloud, since they can use instances with less memory. We see some benchmarks that could, in theory, be run for 30-50% less total cost. This is not universal: others would cost more because they get slower. More on that later.

Beyond this, execution is just more predictable. Memory usage is much more consistent and less likely to spike rapidly:

Smooth is fast

In a few cases, it turns out that smooth scheduling can be even faster.

On average, one representative oceanography workload ran 20% faster. A few other workloads showed modest speedups as well. This is mostly because they no longer load too much data into memory and then have to spill it to disk, which creates significant slowdowns.

We have immediately seen 30% reduction in runtime and thus cost, 60% reduction in peak memory usage, and significantly more reasonable and consistent task assignment

– @roe246 reporting on a DataFrame workload.

Additionally, we found that the extra latency we were worried about didn’t actually slow things down in typical cases. There was no measurable change in pure task throughput on a cluster with fast networking and multi-CPU workers, like Coiled clusters or a single-machine LocalCluster. This was a good lesson in trying the simplest thing first.

Sometimes, slow is still slow (but not why you’d expect)

However, we did notice that a few benchmarks run slower with scheduler-side queuing. The typical slowdown is 5-10%, but in the extreme case, they are ~50% slower (though they also use about half the memory).

These cases use less memory now (blue), but take longer.

The problem is that implementing queuing meant giving up a scheduling feature introduced last year called co-assignment.

As described in the docs, co-assignment tries to schedule initial tasks on the same worker if their outputs will be combined later. This avoids having to transfer data from one worker to another when the downstream task runs, because all the data is already on one worker.

In a graph like this, we’d want a and b to run on the same worker. Otherwise, one of e or f would have to be transferred between workers before i can run (read graph bottom-up):

  i       j
 / \     / \
e   f   g   h
|   |   |   |
a   b   c   d

Avoiding these transfers speeds things up, because network is relatively slow. It also reduces memory usage by avoiding having to hold replicas of the same data on multiple workers.

Unfortunately, the current implementation of co-assignment isn’t compatible with queuing, and updating it is non-trivial. We plan to do this next, to get the best of both worlds.

A new default for scheduling

In the short term, we had to decide if queuing was beneficial enough to be worth enabling by default right away, despite the loss of co-assignment.

After running a number of benchmarks, and getting some initial community feedback, we think it is.

Queuing makes things possible that used to not work at all. But it doesn’t break anything that works today: everything will still work, some things just could be slower. We feel that that’s a worthwhile enough tradeoff to enable it by default—especially given how much Dask users have struggled with memory issues.

Additionally, to avoid impacting workloads that could become latency-bound, the new algorithm is still doing a little overproduction. It’s pushing a handful of extra root tasks to the worker in advance (as opposed to all of them, like before). This comes at the price of some extra memory use, but prevents painful slowdowns in high-latency clusters.

So in the latest release (2022.11.0), queuing is enabled by default. Most memory-intensive Array and DataFrame workloads should see reductions in memory use out of the box, ranging from noticeable to 🤩.

Let us know how it goes

We’ve opened a discussion on GitHub for feedback on this change. Please let us know how it helps (or doesn’t).

Keeping old behavior

For users who are sensitive to runtime and have low memory use, you can deactivate queuing and use the old scheduling mode (including co-assignment) via Dask configuration, by setting the new distributed.scheduler.worker-saturation config value to inf.

You can read more about adjusting this setting in the docs.

On Coiled, you can set it with:

import dask
import coiled

with dask.config.set({"distributed.scheduler.worker-saturation": "inf"}):
    cluster = coiled.Cluster(...)  # coiled sends current dask config automatically

‍

You can see examples of setting this configuration for various deployment systems on the discussion issue (when copy-pasting, be sure to change the 1.0 to inf!). And if you find the need to set worker-saturation back to inf, please let us know on the discussion.

Acknowledgements

Much credit goes to Tom Nicholas for his collaboration and persistence in advocating for this problem, developing minimal reproducers for it, and testing out multiple iterations of fixes on his real scientific workloads (he even gave a talk at SciPy about why this matters!). Similarly, Ryan Abernathy, @JSKenyon, Dougie Squire, and Tom Augspurger all (unwittingly, years ago) provided example code used in validating this work. This was a great example of the power of open-source!

Additionally, thanks to the many people who tested out the new setting and gave feedback on GitHub, plus the whole team at Coiled who collaborated to get this feature over the line.

Reduce memory usage with Dask dtypes

2022-01-01T00:00:00+00:00

Columns in Dask DataFrames are typed, which means they can only hold certain values (e.g. integer columns can’t hold string values). This post gives an overview of DataFrame datatypes (dtypes), explains how to set dtypes when reading data, and shows how to change column types.

Using column types that require less memory can be a great way to speed up your workflows. Properly setting dtypes when reading files is sometimes needed for your code to run without error.

Here’s what’s covered in this post:

Inspecting DataFrame column types (dtypes)
Changing column types
Inferring types when reading CSVs
Manually specifying types
Problems with type inference
Using Parquet to avoid type inference

Inspecting DataFrame types

Let’s start by creating some DataFrames and viewing the dtypes.

Create a pandas DataFrame and print the dtypes. All code snippets in this post are from this notebook.

import pandas as pd

df = pd.DataFrame({"nums":[1, 2, 3, 4], "letters":["a", "b", "c", "d"]})

print(df.dtypes)

nums       int64
letters    object
dtype: object

The nums column has the int64 type and the letters column has the object type.

Let’s use this pandas DataFrame to create a Dask DataFrame and inspect the dtypes of the Dask DataFrame.

import dask.dataframe as dd

ddf = dd.from_pandas(df, npartitions=2)

ddf.dtypes

nums       int64
letters    object
dtype: object

The Dask DataFrame has the same dtypes as the pandas DataFrame.

Changing column types

Change the nums column to int8. You can use Dask’s astype function to cast an object to a different type.

ddf['nums'] = ddf['nums'].astype('int8')

ddf.dtypes

nums        int8
letters    object
dtype: object

8 bit integers don’t require as much memory as 64 bit integers, but can only hold values between -128 and 127, so this datatype is only suitable for columns that contain small numbers.

Selecting the best dtypes for your DataFrame is discussed in more detail later in this post.

Inferring types

Let’s read a 663 million row CSV time series dataset from a cloud filestore (S3) and print the dtypes.

We’ll use the Coiled Dask platform to execute these queries on a cluster.

import coiled
import dask

cluster = coiled.Cluster(name="demo-cluster", n_workers=5)
client = dask.distributed.Client(cluster)

ddf = dd.read_csv(
    "s3://coiled-datasets/timeseries/20-years/csv/*.part", 
    storage_options={"anon": True, 'use_ssl': True}
)

print(ddf.dtypes)

timestamp    object
id           int64
name         object
x            float64
y            float64
dtype: object

Dask infers the column types by taking a sample of the data. It doesn’t analyze every row of data when inferring types because that’d be really slow on a large dataset.

Dask may incorrectly infer types because it only uses a sample. If the inferred type is wrong then subsequent computations will error out. We’ll discuss these errors and common work-arounds later in this post.

Manually specifying types

This section demonstrates how manually specifying types can reduce memory usage.

Let’s look at the memory usage of the DataFrame:

ddf.memory_usage(deep=True).compute()

Index           140160
id          5298048000
name       41289103692
timestamp  50331456000
x           5298048000
y           5298048000
dtype: int64

The id column takes 5.3GB of memory and is typed as an int64. The id column contains values between 815 and 1,193. We don’t need 64 bit integers to store values that are so small. Let’s type id as a 16 bit integer and see how much memory that saves.

ddf = dd.read_csv(
    "s3://coiled-datasets/timeseries/20-years/csv/*.part", 
    storage_options={"anon": True, 'use_ssl': True},
    dtype={
      "id": "int16"
    }
)

ddf.memory_usage(deep=True).compute()

Index           140160
id          1324512000
name       41289103692
timestamp  50331456000
x           5298048000
y           5298048000
dtype: int64

The id column only takes up 1.3 GB of memory when it’s stored as a 16 bit integer.

It’s not surprising that the memory requirements are 4 times lower because int16 is 4x smaller than int64 ;)

Manually specifying column dtypes can be tedious, especially for wide DataFrames, but it’s less error prone than inferring types.

Problems when inferring types

As mentioned earlier, Dask does not look at every value when inferring types. It only looks at a sample of the rows.

Suppose a column has the following values: 14, 11, 10, 22, and “hi”. If Dask only samples 3 values (e.g. 11, 10, and 22), then it might incorrectly assume this is an int64 column.

Let’s illustrate this issue with the New York City Parking Tickets dataset.

ddf = dd.read_csv(
    "s3://coiled-datasets/nyc-parking-tickets/csv/*.csv",
    storage_options={"anon": True, 'use_ssl': True})

len(ddf) # this works

ddf.head() # this errors out

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/var/folders/d2/116lnkgd0l7f51xr7msb2jnh0000gn/T/ipykernel_9266/95726415.py in <module>
---> 1 ddf.head()

ValueError: Mismatched dtypes found in `pd.read_csv`/`pd.read_table`.

+-----------------------+---------+----------+
| Column                | Found   | Expected |
+-----------------------+---------+----------+
| House Number          | object  | float64  |
| Issuer Command        | object  | int64    |
| Issuer Squad          | object  | int64    |
| Time First Observed   | object  | float64  |
| Unregistered Vehicle? | float64 | int64    |
| Violation Description | object  | float64  |
| Violation Legal Code  | object  | float64  |
| Violation Location    | float64 | int64    |
| Violation Post Code   | object  | float64  |
+-----------------------+---------+----------+

The following columns also raised exceptions on conversion:

- House Number
  ValueError("could not convert string to float: '67-21'")
- Issuer Command
  ValueError("invalid literal for int() with base 10: 'T730'")
- Issuer Squad
  ValueError('cannot convert float NaN to integer')
- Time First Observed
  ValueError("could not convert string to float: '1134P'")
- Violation Description
  ValueError("could not convert string to float: 'BUS LANE VIOLATION'")
- Violation Legal Code
  ValueError("could not convert string to float: 'T'")
- Violation Post Code
  ValueError("could not convert string to float: 'H -'")

Usually this is due to Dask’s dtype inference failing, and may be fixed by specifying dtypes manually by adding:

dtype={'House Number': 'object',
       'Issuer Command': 'object',
       'Issuer Squad': 'object',
       'Time First Observed': 'object',
       'Unregistered Vehicle?': 'float64',
       'Violation Description': 'object',
       'Violation Legal Code': 'object',
       'Violation Location': 'float64',
       'Violation Post Code': 'object'}

to the call to `read_csv`/`read_table`.

House Number was inferred to be a float64, but it should actually be an object because it contains values like “67-21”.

Let’s look at the full set of dtypes in this DataFrame with print(ddf.dtypes).

Summons Number                    int64
Plate ID                          object
Registration State                object
Plate Type                        object
Issue Date                        object
Violation Code                     int64
Vehicle Body Type                 object
Vehicle Make                      object
Issuing Agency                    object
Street Code1                       int64
Street Code2                       int64
Street Code3                       int64
Vehicle Expiration Date            int64
Violation Location                 int64
Violation Precinct                 int64
Issuer Precinct                    int64
Issuer Code                        int64
Issuer Command                     int64
Issuer Squad                       int64
Violation Time                    object
Time First Observed              float64
Violation County                  object
Violation In Front Of Or Opposite object
House Number                     float64
Street Name                       object
Intersecting Street               object
Date First Observed                int64
Law Section                        int64
Sub Division                      object
Violation Legal Code             float64
Days Parking In Effect            object
From Hours In Effect              object
To Hours In Effect                object
Vehicle Color                     object
Unregistered Vehicle?              int64
Vehicle Year                       int64
Meter Number                      object
Feet From Curb                     int64
Violation Post Code              float64
Violation Description            float64
No Standing or Stopping Violation float64
Hydrant Violation                float64
Double Parking Violation         float64
Latitude                         float64
Longitude                        float64
Community Board                  float64
Community Council                float64
Census Tract                     float64
BIN                              float64
BBL                              float64
NTA                              float64
dtype: object

Figuring out the right dtypes for each column is really annoying.

We can get past this initial error by specifying the dtypes, as recommended in the error message.

ddf = dd.read_csv(
    "s3://coiled-datasets/nyc-parking-tickets/csv/*.csv",
    dtype={'House Number': 'object',
       'Issuer Command': 'object',
       'Issuer Squad': 'object',
       'Time First Observed': 'object',
       'Unregistered Vehicle?': 'float64',
       'Violation Description': 'object',
       'Violation Legal Code': 'object',
       'Violation Location': 'float64',
       'Violation Post Code': 'object'}
)

ddf.head() will work now but ddf.describe().compute() will error out with this message:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)

ValueError: Mismatched dtypes found in `pd.read_csv`/`pd.read_table`.

+-------------------------+--------+----------+
| Column                  | Found  | Expected |
+-------------------------+--------+----------+
| Date First Observed     | object | int64    |
| Vehicle Expiration Date | object | int64    |
+-------------------------+--------+----------+

The following columns also raised exceptions on conversion:

- Date First Observed
  ValueError("invalid literal for int() with base 10: '01/05/0001 12:00:00 PM'")
- Vehicle Expiration Date
  ValueError("invalid literal for int() with base 10: '01/01/20160411 12:00:00 PM'")

Usually this is due to Dask’s dtype inference failing, and may be fixed by specifying dtypes manually by adding:

dtype={'Date First Observed': 'object',
       'Vehicle Expiration Date': 'object'}

to the call to `read_csv`/`read_table`.

These dtype inference problems are common when using CSV files. This is one of the many reasons to avoid the CSV file format and use files better suited for data analyses.

Avoiding type inference

You should avoid type inference whenever possible.

Certain file formats, like Parquet, let you store file metadata in the footer. The Parquet metadata includes the schema, so you don’t need to manually specify the dtypes or rely on Dask to perform schema inference.

See this blog post on the Advantages of the Parquet format for more detail.

New string type

See this video for the string[pyarrow] datatype that was added to pandas 1.3.

The object datatype is notoriously memory hungry and this new datatype will help a lot!

Conclusion

Dask DataFrame columns are typed, just like pandas DataFrames.

You’ll sometimes run into type-inference errors, especially when you’re working with CSV files.

Use file formats that don’t require type inference so you don’t need to deal with this headache.

It’s great to see the PyData ecosystem constantly evolving and the pandas team adding a new dtype for strings that’s a lot more efficient. Constant improvements are part of what make the PyData ecosystem so wonderful!

Thanks for reading!

PyArrow Strings in Dask DataFrames

2022-01-01T00:00:00+00:00

pandas 2.0 has been released! 🎉

Improved PyArrow data type support is a major part of this release, notably for PyArrow strings, which are faster and more compact in memory than Python object strings, the historic solution. This change impacts pandas users everywhere, but especially impacts Dask DataFrame users, who often run at the capacity of their hardware.

PyArrow strings often use less memory, and so allow you to process more data on less hardware. In the example below we’ll find that we can operate on the same data, faster, using a cluster of one third the size. This corresponds to about a 75% overall cost reduction.

How to use PyArrow strings in Dask

pip install pandas==2

import dask
dask.config.set({"dataframe.convert-string": True})

Note, support isn’t perfect yet. Most operations work fine, but some (any, all, …) still require upstream fixes. Thank you to pandas and PyArrow maintainers for their collaboration in this effort.

This option is available starting with Dask 2023.03.01 and it requires pandas 2.0 to be installed.

Dask PyArrow Example

Let’s consider the following example, where we load some public Uber/Lyft Parquet data onto a cluster running on the cloud.

import coiled
from dask.distributed import wait
import dask.dataframe as dd
import dask

dask.config.set({"dataframe.convert-string": True})

cluster = coiled.Cluster(
	 n_workers=15,
   backend_options={"region": "us-east-2"},
 )
  
client = cluster.get_client()
 
df = dd.read_parquet(
  "s3://coiled-datasets/uber-lyft-tlc/",
  storage_options={"anon": True},
).persist()
wait(df)

This takes about 33 seconds to load. We can then do operations to find out, for example, how many rides are tipped, broken down by Uber and Lyft rides:

df["tipped"] = df.tips != 0
df.groupby(df.hvfhs_license_num).tipped.mean().compute()

If you’d like to play more with this dataset (it’s fun!) there’s a monthly free tutorial that you can sign up here.

Turn off PyArrow Strings

We can run this exact same computation, but now without PyArrow strings:

dask.config.set({"dataframe.convert-string": False})

Unfortunately, the first step to make this work is to ask for 3x the hardware. Otherwise the dataset doesn’t fit into memory, and we’re stuck repeatedly reading from storage:

cluster.scale(45)
client.wait_for_workers(45)

Then everything works fine, and we get the same results as above, although operations like the groupby-aggregation above take twice as long (even with 3x the hardware).

Learn, Try, Report

PyArrow strings are a significant win for the entire PyData community. They are also a work in progress. While we’re excited to adopt them within our own benchmarking suites, we also know that early adopters will experience some rough edges. If you’d like to learn more, here are a few resources:

pandas: What’s new in pandas 2.0.0
Uber/Lyft dataset: Regular free tutorial playing with this dataset
Dask + PyArrow: Discussion issue where you can report experience, positive or negative

If you have the time to use PyArrow strings with (or without) Dask and report your experience we’d love to hear about what’s working well and what needs improvement.

Thank you to everyone who contributed to this, but especially Matthew Roeschke (NVIDIA), Patrick Hoefler (Coiled), and Rick Zamora (NVIDIA) who worked to make sure all of the different libraries needed here came together.

Prioritizing Pragmatic Performance for Dask

2022-01-01T00:00:00+00:00

Many people say the following to me:

You should focus on X in Dask! X makes things really fast!

Some examples of X:

Polars
High level query optimization
C/C++/Rust
POSIX shared memory
Reducing memory copies
Massive scale

In some cases I think that they’re absolutely right. In other cases I agree that the thing that they’re talking about is fast, but I don’t think that it will meaningfully improve common workloads, especially in a distributed cloud computing environment.

Instead, I tend to orient development pretty strongly around what profiling, benchmarking, and user engagement say is the biggest bottleneck, and go from there. This points me to topics like the following:

Improving S3 bandwidth
Intelligent scheduling to avoid memory buildup
Removing Python object dtypes
Removing Python for loops from user code
Better observability to help users know what’s going on

This means that Dask often loses in flat-out performance benchmarks in artificial scenarios, but that it continues to be the standard of use when it comes to scalable Python data processing.

This article will go into a few things that I think are great ideas, but are not yet major bottlenecks and why. Then it will go into a few things that I think are less-than-exciting-but-really-important improvements for common workloads.

Things people ask for

Anything that is faster than S3

First, many workloads are bound by S3/GCS/object-store read performance. This tends to be 200 MB/s per node (with variation). If the thing that you’re talking about is much faster than this then it probably doesn’t matter that much (unless you’re doing it hundreds of times per byte (in which case you’re maybe doing something wrong)).

Memory Copies

Memory copies make things slow, but only by a tiny bit

Memory bandwidth often operates at 20 GB/s on a modest machine. So while it’s true that memory copies slow things down, they only slow them down a little bit. It’s like asking a racecar to go once more around the racetrack when we’re still waiting for a donkey (s3/disk/network) to finish going around once. It just doesn’t matter.

But for the sake of argument let’s imagine for a moment that you’re at a performance point where this does matter. That means that you are processing data at 20 GB/s on a modest machine, or around 1 TB per minute. Per day you’re processing about 1 PB on this single node.

If you’re in that world, then you either …

Generate many petabytes of data per day
Don’t need parallel computing, and so don’t need Dask (Hooray! Fewer tools!)

High Level Query Optimization

Doing less work is good

In contrast, I think that high level optimization is very very good. For one, it often means that you get to reduce the amount of data that you load, which reduces the demand on the slow parts of your system (like S3). Dask Dataframes/Arrays doesn’t do this well today, but it should.

Updated August 2023: You can now try out Dask Expressions with pip install dask-expr. Learn more in our blog post.

C/C++/Rust

You’re already using these for 99% of your computation, the last 1% doesn’t matter much

If you’re doing things well then you’re already relying on mostly low level code (numpy, pandas, pytorch, scikit-learn, arrow, …) The code calling those routines accounts for a small fraction of your total runtime and is not worth optimizing.

Said another way, most data workloads are saturated by network and disk. I can, with Python, saturate network and disk and so go as fast as possible. If we’re in this regime then, unless a faster language can somehow make your disk/network hardware go faster, then language choice doesn’t matter.

Polars

Taking a look at the options above, I’m really excited about Polars mostly due to the high level query optimization. The other bits are nice, but mostly on single machines where you have really fast access to disk and don’t need to worry about network as much.

Ritchie Vink (polars core dev) and I actually started a dask-polars project several months ago. I’d love to work more on it (it’s fun!). I’m waiting for this need to come up more in benchmarks before I personally prioritize it more.

The best arguments that I’ve heard for prioritizing it early so far are …

High level query optimization (Dask dataframes should have this, and don’t)
Some folks like the API more than Pandas

The more I hear folks report #2 above the more I’ll shift efforts over to it. So far I’m seeing 50x more Pandas usage (but down from 1000x a few months ago, so the trend here is worth watching).

To be clear though, folks should use Polars. It helps to avoid the distributed computing problem significantly. Polars is great.

POSIX Shared Memory

This mostly comes from ray users who like the idea of sharing memory between processes. The conversation usually goes like this:

Them: Dask should use POSIX shared memory so that you can share data between cores!
Me: Dask does that today! We use a single normal process with lots of normal threads and normal shared memory!
Them: But then what about the GIL?!
Me: The GIL isn’t an issue! The GIL has been released in most PyData libraries ever since Dask came out seven years ago! Here is an ancient blogpost on the topic!

Threads + normal memory is less fancy, but also way simpler than posix-shared-memory!

Them: Hrm, maybe this is important if people are writing lots of for-loopy python code?
Me: That’s true! But that’s also a 100x performance hit! I encourage you to address that much much deeper issue before trying to parallelize it. Maybe consider teaching the pandas API or using Numba?

Now, this is really important if you’re doing distributed training of very large machine learning models. Almost no one I talk to is doing this though. Mostly this issue comes up because, I think, people misunderstand the GIL.

To address this misunderstanding, I think that we need to improve observability around the GIL. See dask/distributed#7290.

But let’s say that folks here are right and that we do want to reduce intra-machine communication costs with shared memory. How much of a savings is this, actually? The common alternative is just to ask for many more smaller machines. It’s true that inter-machine communication is more expensive than intra-machine communication, but most communication is going to be inter-machine anyway, so the savings here is minimal. There’s just no way that this respects Amdahl’s law.

Massive Scale

We work at really big scale. What if we want to scale to thousands of machines?

Well, Dask actually works fine at thousands of machines. You’ll probably want to change some of the default configuration; let’s chat. Really though, if you’re doing things well then you probably don’t need thousands of machines.

I work with some of the largest organizations doing some of the largest computations. It is very rare for someone to need thousands of machines. It is much more common for them to think that they need thousands of machines, get feedback from Dask about how their workflow is running, and then learn that they can accomplish more than they wanted with far fewer resources.

Most Dask clusters in large organizations operate at the 10-500 machines range with the bottom of that range being far more common. Some folks get up to a thousand machines. Even at this point the scheduler can be idling most of the time if you’re doing things efficiently.

Things I’m excited about

Most dataframe workloads that are slow are due to one of these things:

Running out of memory
Object dtypes
Iterating over rows with a for loop
Network/disk bandwidth
Loading in data we didn’t need

This covers about 95% of slow workflows. People who get rid of these issues are invariably happy. Let’s talk through a few things that are in flight today that can be done to address them.

Task Queuing / more intelligent memory scheduling

The Dask scheduler is very smart about placing and prioritizing tasks to reduce overall memory use. If you are running into a workload that you think should run in small memory but doesn’t then I encourage you to try out the 2022.11.0 release. It’s pretty awesome. We’ve written more about it in this blog post.

Here is a snippet from the last Dask Demo Day with engineer Florian Jetter talking about it and showing the impact:

PyArrow Strings

Updated August 2023: as of Dask 2023.7.1, PyArrow backed strings are the default in Dask DataFrames. This option is automatically enabled if you have at least pandas 2.0 and PyArrow 12.0 installed. See the Dask release notes for implementation details and our blog post on how this change improves performance.

Python is slow. Normally we get past this by writing performant code in C/C++/Cython and then linking to it. Pandas does this, but unfortunately they didn’t have a good way to handle strings, so they stayed with Python for that kind of data. 😞

Fortunately, Arrow handles this. They have nice strings implemented in C++. This results in lower memory use and faster speeds. Here is a demonstrative video:

For a long time PyArrow strings were exciting, but not fully supported. It seems like today they are. Pandas itself is considering switching over by default. My guess is that Dask may lead pandas here a bit and make the switch earlier. This is in active development.

Extension data types more generally

This also holds for other data types, including nullable dtypes (ints, bools, etc..), decimals, and others. There’s good movement towards just switching over generally.

Pandas 2.0 is coming out soon (rc branch already tagged) and in general I expect to see a big shift in this direction over the next year.

S3 Bandwidth

Sometimes we don’t handle things as well as we should. Recently we found a nice 2x speed boost.

dask/dask#9619

Shuffle service

When you want to sort/set-index/join dataframes you often need to do an all-to-all shuffle. Dask’s task scheduling system wasn’t designed for this and while it does it, it does it poorly in a way that often stresses the scheduler and bloats memory.

There is an experimental shuffle service (similar to what Spark uses) in Dask today that has solid performance, but doesn’t yet have rock-solid-reliability (and so is not yet default). We hope to have this in by early 2023.

High Level Query Optimization

Updated August 2023: You can now try out Dask Expressions with pip install dask-expr. Learn more in our blog post.

Probably sometime in 2023 we should focus on updating our high level graph specification, which grew organically and has a bit of cruft. After this we should be able to start development on basic query optimizations. This plus shuffling and good dtype support should get us into a point where Dask dataframes are genuinely competitive performance-wise in ways that matter.

Additionally, high level query optimization should also be applicable to other Dask collections, like Arrays/Xarray, which should be fun.

Things that aren’t performance

It’s worth pointing out that most distributed data computing users I meet would happily trade in a 3x performance hit if things would “just work”. Topics like ease of deployment, visibility, stability, and ease-of-use for less technical teammates often supersede a desire for things to go fast.

Dask has never been the fastest framework (go to MPI for that) but we have probably always been the easiest to use and understand and adapt to user needs. This results in a joint human+computer performance that can’t be beat.

Summary

Dask developers care about performance, we’ve always taken a pragmatic rather than exciting approach to the problem. We interact with users, identify user pain, and focus on that pain.

This focus on user pain has served us well in the past, but often orients us differently from the focus of the public, which is understandably excited by shiny developments. I encourage folks to profile and benchmark their workloads (Dask provides excellent tools for this) and let us know what you see. We’re excited for more input.

Perform a Spatial Join in Python

2022-01-01T00:00:00+00:00

This blog explains how to perform a spatial join in Python. Knowing how to perform a spatial join is an important asset in your data-processing toolkit: it enables you to join two datasets based on spatial predicates. For example, you can join a point-based dataset with a polygon-based dataset based on whether the points fall within the polygon.

	GeoPandas	Dask-GeoPandas
Perform Spatial Join
Support Geographic Data
Parallel Processing
Larger-than-memory Datasets

In this blog, we will use the New York City Taxi and Limousine Commission dataset to illustrate performing a spatial join in Python. We will join the taxi pickup location of each individual trip record to the NYC neighborhood it is located within. This will then allow us to bring in public health data at the neighborhood level and discover interesting patterns in the data that we could not have seen without performing a spatial join.

Spatial Join: Different Python Libraries

There are several different ways to perform a spatial join in Python. This blog will walk you through two popular open-source Python you can use to perform spatial joins:

Geopandas
Dask-GeoPandas

Dask-GeoPandas has recently undergone some impressive changes due to the hard work of maintainers Joris Van den Bossche, Martin Fleischmann and Julia Signell.

For each library, we will first demonstrate and explain the syntax, and then tell you when and why you should use it. You can follow along in this Jupyter notebook.

Let’s jump in.

Spatial Join: GeoPandas

GeoPandas extends pandas DataFrames to make it easier to work with geospatial data directly from Python, instead of using a dedicated spatial database such as PostGIS.

A GeoPandas GeoDataFrame looks and feels a lot like a regular pandas DataFrame. The main difference is that it has a geometry column that contains shapely geometries. This column is what enables us to perform a spatial join.

Sidenote: geometry is just the default name of this column, a geometry column can have any name. There can also be multiple geometry columns, but one is always considered the “primary” geometry column on which the geospatial predicates operate.

Let’s load in the NYC TLC data. We’ll start by working with first 100K trips made in January 2009:

import pandas as pd
df = pd.read_csv(
    "s3://nyc-tlc/trip data/yellow_tripdata_2009-01.csv",
    nrows=100_000
)

The dataset requires some basic cleaning up which we perform in the accompanying notebook.

Next, let’s read in our geographic data. This is a shapefile sourced from the NYC Health Department and contains the shapely geometries of 35 New York neighborhoods. We’ll load this into a GeoPandas GeoDataFrame:

import geopandas as gpd
ngbhoods = gpd.read_file(
    "../data/CHS_2009_DOHMH_2010B.shp"
)[["geometry", "FIRST_UHF_", "UHF_CODE"]]

Let’s inspect some characteristics of the nghboods object including the object type, how many rows it contains, and the first 3 rows of data:

type(ngbhoods)
geopandas.geodataframe.GeoDataFrame

len(ngbhoods)
35

ngbhoods.head(3)

ngbhoods is a GeoPandas GeoDataFrame containing the shapely polygons that define the 35 New York U neighborhoods as defined by the United Hospital Fund, as well as their names and unique ID codes.

A neat trick to know about GeoPandas is that calling .plot() on a GeoDataFrame will automatically plot the geometries of the shapes contained:

ngbhoods.plot()

When joining spatial datasets, it’s important to standardize the Coordinate Reference System (CRS). A CRS defines how coordinates (numbers like 42.01 or 180000) map to points on the surface of the Earth. Discussing why you’d pick a particular CRS is out of scope for this post, but you visit these links to learn more. What is important is that the two datasets you want to join together have the same CRS. Joining datasets with mismatched CRSs will lead to incorrect results downstream.

The NYC TLC data uses latitude and longitude, with the WGS84 datum, which is the EPSG 4326 coordinate reference system. Let’s inspect the CRS of our ngbhoods GeoDataFrame:

ngbhoods.crs
<Derived Projected CRS: EPSG:2263>

The CRSs don’t match, so we’ll have to transform one of them. It’s more efficient to transform a smaller rather than a larger dataset, so let’s project ngbhoods to EPSG 4326:

ngbhoods = ngbhoods.to_crs(epsg=4326)

Sidenote: if you are working with multiple GeoDataFrames you can set the CRS of one gdf to that of another using gdf1.to_crs(gdf2.crs) without worrying about the specific CRS numbers.

Calling head() on ngbhoods now will reveal that the values for the geometry column have changed:

Plotting the GeoDataFrame confirms that our spatial information is still intact. But if you look at the axes, you’ll notice the coordinates are now in degrees latitude and longitude instead of meters northing and easting:

Next, let’s take a look at the regular pandas DataFrame containing our taxi trip data. You’ll notice that this DataFrame also contains geographic information in the pickup_longitude/latitude and dropoff_longitude/latitude columns.

df.head(3)

Looking at the dtypes, however, reveals them to be regular numpy floats:

df.dtypes

vendor_id           object
pickup_datetime      object
dropoff_datetime     object
passenger_count       int64
trip_distance       float64
pickup_longitude    float64
pickup_latitude     float64
rate_code             int64
store_and_fwd_flag   object
dropoff_longitude   float64
dropoff_latitude    float64
payment_type         object
fare_amount         float64
surcharge           float64
mta_tax             float64
tip_amount          float64
tolls_amount        float64
total_amount        float64
dtype: object

Calling plot() on the latitude and longitude columns does not give us a meaningful spatial representation. Let’s see how latitude / longitude values are difficult to interpret without a geospatial library.

Start by plotting the latitude / longitude values:

df.plot(
    kind='scatter', x="pickup_longitude", y="pickup_latitude"
)

This result is difficult to interpret, partly because there are clearly some data points with incorrect coordinate values in the dataset. This won’t be a problem for our particular use case, but may be worth paying attention to if you’re using this data for other analyses.

Let’s turn df into a GeoPandas GeoDataFrame, so we can leverage the built-in geospatial functions that’ll make our lives a lot easier. Convert the separate pickup_longitude and pickup_latitude columns into a single geometry column containing shapely Point data:

from geopandas import points_from_xy
# turn taxi_df into geodataframe
taxi_gdf = gpd.GeoDataFrame(
    df, crs="EPSG:4326",
    geometry=points_from_xy(
        df["pickup_longitude"], df["pickup_latitude"]
    ),
)

Take a look at first 3 rows:

taxi_gdf.head(3)

We now have a GeoDataFrame with a dedicated geometry column.

Calling .explore() on a GeoDataFrame allows you to explore the dataset on a map. Run this notebook to interactively explore the map yourself:

taxi_gdf.explore(
    tiles="CartoDB positron", # use "CartoDB positron" tiles
    cmap="Set1", # use "Set1" matplotlib colormap
    style_kwds=dict(color="black") # use black outline
)

Now that both of your datasets have geometry columns, you can perform a spatial join.

Perform a left spatial join with taxi_gdf as the left_df and ngbhoods as the right_df. Setting the predicate keyword to ‘within’ will join points in the left_df to polygons from the right_df they are located within:

%%timeit
joined = gpd.sjoin(
    taxi_gdf, 
    ngbhoods, 
    how='left',
    predicate="within",
)

1.13 s ± 33.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Clean up and clarify some column names:

# drop index column
joined = joined.drop(columns=['index_right'])

# rename columns
joined.rename(
    columns={"FIRST_UHF_": "pickup_ngbhood_name", "UHF_CODE": "pickup_ngbhood_id"},
    inplace=True,
)

Now that you’ve performed the spatial join, you can analyze spatial patterns in the data.

For example, you can run a groupby that aggregates the mean trip_distance per pickup neighborhood:

res = joined.groupby('pickup_ngbhood_name').trip_distance.mean().sort_values()
res

pickup_ngbhood_name
Rockaway                              18.640000
Southwest Queens                       7.583590
Bayside - Meadows                      5.461818
East New York                          5.195714
Jamaica                                4.933902
Washington Heights - Inwood            4.746627
Pelham - Throgs Neck                   4.478235
Southeast Queens                       4.322500
Kingsbridge - Riverdale                4.140000
Downtown  - Heights - Slope            3.833111
Flushing - Clearview                   3.761765
Sunset Park                            3.612857
Ridgewood - Forest Hills               3.311461
Canarsie - Flatlands                   3.290000
Greenpoint                             3.153171
Long Island City - Astoria             3.055458
Coney Island - Sheepshead Bay          3.036667
West Queens                            3.034350
Central Harlem - Morningside Heights   2.866855
Union Square, Lower Manhattan          2.721932
Bedford Stuyvesant - Crown Heights     2.643019
East Harlem                            2.638558

Spatial Join: Dask-GeoPandas

Now that we’ve seen how to perform a spatial join with GeoPandas, let’s look at how to do the same thing with Dask-GeoPandas. Dask-GeoPandas is particularly useful when you’re working with datasets that are too large to fit into memory.

Let’s load the full NYC TLC dataset for January 2009 (about 120GB of data) into a Dask DataFrame:

import dask.dataframe as dd

ddf = dd.read_csv(
    "s3://nyc-tlc/trip data/yellow_tripdata_2009-01.csv",
    storage_options={"anon": True, "use_ssl": True},
)

Let’s perform some basic data cleaning and filtering:

# convert datetime columns
ddf["pickup_datetime"] = dd.to_datetime(ddf["pickup_datetime"])
ddf["dropoff_datetime"] = dd.to_datetime(ddf["dropoff_datetime"])

# filter out invalid coordinates
ddf = ddf[
    (ddf["pickup_longitude"] != 0) & 
    (ddf["pickup_latitude"] != 0) &
    (ddf["pickup_longitude"].between(-74.3, -73.7)) &
    (ddf["pickup_latitude"].between(40.5, 40.9))
]

Let’s also load some public health data at the neighborhood level:

health = gpd.read_file(
    "../data/CHS_2009_DOHMH_2010B.shp"
)[["geometry", "FIRST_UHF_", "UHF_CODE"]]
health = health.rename(columns={"FIRST_UHF_": "nhbd_name", "UHF_CODE": "nhbd_id"})
health = health.to_crs(epsg=4326)

Now let’s convert our Dask DataFrame into a Dask-GeoPandas GeoDataFrame:

import dask_geopandas

ddf = dask_geopandas.from_dask_dataframe(
    ddf,
    geometry=dask_geopandas.points_from_xy(ddf, "pickup_longitude", "pickup_latitude"),
)
ddf = ddf.set_crs(4326)

Before we perform the spatial join, we need to make sure all workers have access to the health data:

from dask.distributed import Client
client = Client()

# broadcast health data to all workers
health_future = client.scatter(health)

Now we can perform the spatial join:

%%time
joined = ddf.sjoin(health_future, predicate="within")

The spatial join is now distributed across all workers in the cluster. Let’s compute some statistics:

%%time
stats = joined.groupby("nhbd_name").agg({
    "trip_distance": "mean",
    "fare_amount": "mean",
    "tip_amount": "mean",
    "passenger_count": "mean"
}).compute()

This will give us average trip statistics per neighborhood, which we can use to analyze patterns in the data.

Spatial Join: the Final Verdict

Spatial joins are powerful because they allow you to join datasets based on geographic data.
GeoPandas is great for performing a spatial join on datasets that fit into memory.
Dask-GeoPandas scales GeoPandas for larger-than-memory datasets using the parallel processing power of Dask.
You can speed up a spatial join by running your computations in the cloud using a Coiled Cluster.

	GeoPandas	Dask-GeoPandas	on Coiled
Perform Spatial Join
Support Geographic Data
Parallel Processing
Larger-than-memory Datasets
Scale to the cloud

Introducing the Dask Active Memory Manager

2022-01-01T00:00:00+00:00

Historically, out-of-memory errors and excessive memory requirements have frequently been a pain point for Dask users. Two of the main causes of memory-related headaches are data duplication and imbalance between workers.

Dask release 2021.10.0 introduces the first piece of a new modular system called Active Memory Manager, which aims to alleviate memory issues. The examples in this article use Dask release 2021.11.2, which features improved API and tools.

Data duplication

Whenever a Dask task returns data, it is stored on the worker that executed the task for as long as it’s referenced by a Future or another task:

from distributed import Client
client = Client(n_workers=2)
w1, w2 = client.has_what()
x = client.submit(lambda i: i + 1, 1, workers=[w1], key="x")
x.result() # Output: 2
client.has_what() # Output: {'tcp://127.0.0.1:37465': ('x',), 'tcp://127.0.0.1:44715': ()}

When a task runs on a worker and requires as its input data that was returned by another task on a different worker, Dask will transparently transfer the data between workers, ending up with multiple copies of the same data on different workers. This is generally desirable, as it avoids re-transferring the data if it’s required again later on. However, it also causes increased overall memory usage across the cluster:

y = client.submit(lambda i: i + 1, x, workers=[w2], key="y")
y.result() # Output: 3
client.has_what() # Output: {'tcp://127.0.0.1:37465': ('x',), 'tcp://127.0.0.1:44715': ('x', 'y')}

Memory imbalances

Dask assigns tasks to workers following criteria of CPU occupancy, locality, and worker resources. This design aims to optimize the overall computation time; as a side effect, however, it can lead to some workers ending up with substantially higher memory usage than others.

The Active Memory Manager

Active Memory Manager, or AMM for short, is a background routine that periodically runs in the scheduler and that copies, moves, and removes data on the Dask cluster. It is a plugin-based design where the core component listens to and then enacts suggestions from a list of policies. In other words, the policies implement arbitrarily sophisticated decisions regarding where data should be, while the AMM core component makes sure that these decisions are put in practice without compromising data integrity.

At the moment of writing, the Active Memory Manager is shipped with a single built-in policy, ReduceReplicas, which cleans duplicated data on the workers when it’s no longer necessary.

When it comes to choosing which copies to preserve or create and which to delete, the AMM always prefers preserving/creating on the workers with the lowest memory usage and deleting on those with the highest - thus rebalancing the cluster as it runs.

Enabling the AMM

Since it is an experimental feature, the Active Memory Manager is disabled by default.

The simplest way to turn it on is through the dask config:

distributed:
  scheduler:
   active-memory-manager:
    start: true

You can do the same with an environment variable before you start the scheduler:

export DASK_DISTRIBUTED__SCHEDULER__ACTIVE_MEMORY_MANAGER__START=True
dask-scheduler

Or if you’re using Coiled:

cluster = coiled.Cluster(
    environ={
        "DASK_DISTRIBUTED__SCHEDULER__ACTIVE_MEMORY_MANAGER__START": "True"
    }
)

You can also enable/disable/test it on the fly from the dask client:

client.amm.start()
client.amm.running() # Output: True
client.amm.stop()
client.amm.running() # Output: False

With any of the above methods, the AMM will run all default policies (at the moment, just ReduceReplicas) every 2 seconds. If you’re a power user, you can further alter the Dask configuration to change the run interval, cherry-pick and configure individual policies, or write your own custom policies.

You can read the full documentation at https://distributed.dask.org/en/latest/active_memory_manager.html.

A demonstration with dot product

The dot product between two matrices is a O(n²) problem where each point of the matrices needs to be paired with each other:

import dask.array as da
A = da.random.random((40_000, 40_000), chunks=(4_000, 4_000))
b = (A @ A.T).sum()

A is a square matrix worth 12 GiB in total, split across 122 MiB chunks:

If you run b.compute() on either the local threaded scheduler or on a single distributed threaded worker, the peak memory usage will be roughly 19 GiB; this is because no copies of the data need to be moved across workers. If you run the same problem on a cluster with multiple workers, however, you will necessarily incur data duplication. While the individual workers can mount a lot less than in the local use case, you should expect the cluster-wide occupation (the sum of the occupation on all the workers) to be substantially higher.

Let’s use a free trial account on Coiled to quickly set up a cluster with 48 workers and 96 CPUs:

import coiled
from distributed import Client

cluster = coiled.Cluster(
    n_workers=48,
    environ={
        "DASK_DISTRIBUTED__SCHEDULER__ACTIVE_MEMORY_MANAGER__START": "True"
    },
)
client = Client(cluster)

Let’s test that the Active Memory Manager is running:

client.amm.running() # Output: True

We’re going to use MemorySampler (new in Dask 2021.11.2) to record memory usage on the cluster:

from distributed.diagnostics import MemorySampler
ms = MemorySampler()

Then we’re going to run our computation on the cluster and fetch the history of our memory usage:

with ms.sample("AMM on"):
    b.compute()

Now let’s do it again, but this time without the Active Memory Manager:

client.amm.stop()
with ms.sample("AMM off"):
    b.compute()

Finally, let’s plot the data to see the difference:

ms.plot(align=True, grid=True)

ms.to_pandas().max(axis=0) / 2**30

#
# Output:
#
# With AMM     135.177975
# No AMM       165.925610
# dtype: float64

The Active Memory Manager reduced the peak cluster-wide memory usage from 166 GiB to 135 GiB, with no runtime degradation!

Next steps

ReduceReplicas is just the first of a series of AMM policies that will be released in the future:

Worker retirement is being reimplemented to run on top of the AMM. The key benefit to this is that it will become possible to gracefully retire a worker on a busy cluster, while jobs are running on it. When a worker runs out of memory, it is automatically retired and restarted; this change will prevent random crashes in the jobs when that happens.
Graceful worker retirement will also let you partially downscale an adaptive cluster before it’s completely idle; this will potentially result in substantial monetary savings whenever a job features an initial burst of parallelism followed by a long “tail” of serial tasks, or when small jobs are continuously pushed to the cluster and there’s a spike in usage (more jobs, or a much larger job) at a certain time of the day.
Worker pause (which, by default, is triggered when a worker reaches 80% memory usage) is being redesigned to gracefully transition into retirement after a timeout expires.
Rebalance is being reimplemented as an AMM policy and will automatically run every two seconds, while computations are running and without any need for user intervention. The manual method Client.rebalance is going to be phased out.
Replicate is also becoming an AMM policy. Today, Client.replicate synchronously creates replicas, but does not track them after it returns; if a replica is lost later on, nothing will regenerate it unless the user manually invokes replicate() again. In the future, Client.replicate will inform the AMM of the desired number of replicas and immediately return; the AMM will start tracking them and ensure that, whenever a replica is lost, it is recreated somewhere else (as long as at least one copy of the data survives).

You’ll also be able to ask for a key to be replicated on all workers, including those that will join the cluster in the future.

Thanks for reading!

If you’d like to scale your Dask work to the cloud, check out Coiled — Coiled provides quick and on-demand Dask clusters along with tools to manage environments, teams, and costs.

How to Merge Dask DataFrames

2022-01-01T00:00:00+00:00

This post demonstrates how to merge Dask DataFrames and discusses important considerations when making large joins.

You’ll learn:

how to join a large Dask DataFrame to a small pandas DataFrame
how to join two large Dask DataFrames
how to structure your joins for optimal performance

The lessons in this post will help you execute your data pipelines faster and more reliably, enabling you to deliver value to your clients in shorter cycles.

Let’s start by diving right into the Python syntax and then build reproducible data science examples you can run on your machine.

Dask Dataframe Merge

You can join a Dask DataFrame to a small pandas DataFrame by using the dask.dataframe.merge() method, similar to the pandas api.

Below we create a Dask DataFrame with multiple partitions and execute a left join with a small pandas DataFrame:

import dask.dataframe as dd
import pandas as pd

# create sample large pandas dataframe
df_large = pd.DataFrame(
     {
        "Name": ["Azza", "Brandon", "Cedric", "Devonte", "Eli", "Fabio"], 
        "Age": [29, 30, 21, 57, 32, 19]
     }
)
# create multi-partition dask dataframe from pandas
large = dd.from_pandas(df_large, npartitions=2)

# create sample small pandas dataframe
small = pd.DataFrame(
     {
        "Name": ["Azza", "Cedric", "Fabio"], 
        "City": ["Beirut", "Dublin", "Rosario"]
     }
)

# merge dask dataframe to pandas dataframe
join = large.merge(small, how="left", on=["Name"])

# inspect results
join.compute()

To join two large Dask DataFrames, you can use the exact same Python syntax. If you are planning to run repeated joins against a large Dask DataFrame, it’s best to sort the Dask DataFrame using the .set_index() method first to improve performance.

large.set_index()

large_join = large.merge(also_large, how="left", left_index=True, right_index=True)

Dask DataFrame merge to a small pandas DataFrame

Dask DataFrames are divided into multiple partitions. Each partition is a pandas DataFrame with its own index. Merging a Dask DataFrame to a pandas DataFrame is therefore an embarrassingly parallel problem. Each partition in the Dask DataFrame can be joined against the single small pandas DataFrame without incurring overhead relative to normal pandas joins.

Let’s demonstrate with a reproducible Python code example. Import dask.dataframe and pandas and then load in the datasets from the public Coiled Datasets S3 bucket. This time the Dask DataFrame is actually large and not just a placeholder: it contains a 35GB dataset of time series data. This means the data is too large to run with pandas on almost all machines.

import dask.dataframe as dd
import pandas as pd 

# create dataframes
large = dd.read_parquet("s3://coiled-datasets/dask-merge/large.parquet")
small = pd.read_parquet("s3://coiled-datasets/dask-merge/small.parquet")

# inspect large dataframe
large.head()

# check number of partitions
>>> large.npartitions
359

# inspect small dataframe
small.head()

The large dataframe contains a large dataset of synthetic time series data with entries at a frequency of 1 second. The small dataframe contains synthetic data over the same time interval but at a frequency of one entry per day.

Let’s execute a left join on the timestamp column by calling dask.dataframe.merge(). We’ll use large as the left dataframe.

joined = large.merge(
     small, 
     how="left", 
     on=["timestamp"]
)

joined.head()

‍

As expected, the column z is filled with NaN for all entries except the first per-second entry of every day.

If you’re working with a small Dask DataFrame instead of a pandas DataFrame, you have two options. You can convert it into a pandas DataFrame using .compute(). This will load the DataFrame into memory. Alternatively, if you can’t or don’t want to load it into your single machine memory, you can turn the small Dask DataFrame into a single partition by using the .repartition() method instead. These two operations are programmatically equivalent which means there’s no meaningful difference in performance between them. Rule of thumb here is to keep your Dask partitions under 100MB each.

# turn dask dataframe into pandas dataframe
small = small.compute()

# OR turn dask dataframe into one partition
small = small.repartition(npartitions=1)

Merge two large Dask DataFrames

You can merge two large Dask DataFrames with the same .merge() API syntax.

large_joined = large.merge(
     also_large, 
     how="left", 
     on=["timestamp"]
)

However, merging two large Dask DataFrames requires careful consideration of your data structure and the final result you’re interested in. Joins are expensive operations, especially in a distributed computing context. Understanding both your data and your desired end result can help you set up your computations efficiently to optimize performance. The most important consideration is whether and how to set your DataFrame’s index before executing the join. Note that in the previous example, the timestamp column is the index for both dataframes.

Unsorted vs Sorted Joins

As explained above, Dask DataFrames are divided into partitions, where each single partition is a pandas DataFrame. Dask can track how the data is partitioned (i.e. where one partition starts and the next begins) using a DataFrame’s divisions.

If a Dask DataFrame’s divisions are known, then Dask knows the minimum value of every partition’s index and the maximum value of the last partition’s index. This enables Dask to take efficient shortcuts when looking up specific values. Instead of searching the entire dataset, it can find out which partition the value is in by looking at the divisions and then limit its search to only that specific partition. This is called a sorted join. The join stored in large_joined we executed above is an example of a sorted join since the timestamp column is the index for both of the dataframes in that join.

Let’s look at another example. The divisions of the DataFrame df below are known: it has 4 divisions. This means that if we look up the row with index 2015-02-12, Dask will only search the 2nd partition and won’t bother with the other three. In reality, Dask Dataframes often have hundreds or even thousands of partitions, which means the benefit of knowing where to look for a specific value becomes even greater.

df.known_divisions
# True

df.npartitions
# 4

df.divisions
# ['2015-01-01', '2015-02-01', '2015-03-01', '2015-04-01', '2015-04-31']

If divisions are not known, then Dask will need to move all of your data around so that rows with matching values in the joining columns end up in the same partition. This is called an unsorted join and it’s an extremely memory-intensive process, especially if your machine runs out of memory and Dask will have to read and write data to disk instead. This is a situation you want to avoid. Read more about unsorted large to large joins in the Dask documentation.

Sorted Join with set index

To perform a sorted join of two large Dask DataFrames, you will need to ensure that the DataFrame’s divisions are known by setting the DataFrame’s index. You can set a Dask DataFrame’s index and pass it the known divisions using:

# use set index to get divisions
dask_divisions = large.set_index("id").divisions
unique_divisions = list(dict.fromkeys(list(dask_divisions)))

# apply set index to both dataframes
large_sorted = large.set_index("id", divisions=unique_divisions)
also_large_sorted = also_large.set_index("id", divisions=unique_divisions)

large_join = large_sorted.merge(
     also_large_sorted, 
     how="left", 
     left_index=True, 
     right_index=True
)

Note that setting the index is itself also an expensive operation. The rule of thumb here is that if you’re going to be joining against a large Dask DataFrame more than once, it’s a good idea to set that DataFrame’s index first. Read our blog about setting a Dask DataFrame index to learn more about how and when to do this.

It’s good practice to write sorted DataFrames to the Apache Parquet file format in order to preserve the index. If you’re not familiar with Parquet, then you might want to check out our blog about the advantages of Parquet for Dask analyses.

Joining along a non-Index column

You may find yourself in the situation of wanting to perform a join between two large Dask DataFrames along a column that is not the index. This is basically the same situation as not having set the index (i.e. an unsorted join) and will require a complete data shuffle, which is an expensive operation. Ideally you’ll want to think about your computation in advance and set the index right from the start. The Open Source Engineering team at Coiled is working actively to improve shuffling. Read the Proof of Concept for better shuffling in Dask if you’d like to learn more.

Sorted Join Fails Locally (MemoryError)

Even a sorted join may fail locally if the datasets are simply too large for your local machine. For example, this join:

large = dd.read_parquet("s3://coiled-datasets/dask-merge/large.parquet")
also_large = dask.datasets.timeseries(start="1990-01-01", end="2020-01-01", freq="1s", partition_freq="1M", dtypes={"foo": int})

large_join = large.merge(
     also_large, 
     how="left", 
     left_index=True, 
     right_index=True
)

large_join.persist()

Will throw an error when run on a laptop with 32GB of RAM or less.

Note that we didn’t set the index explicitly here because the index of large is preserved in the Parquet file format and dask.datasets.timeseries() automatically sets the index when creating the synthetic data.

Run Massive Joins on Dask Cluster

When this happens, you can scale out to a Dask cluster in the cloud with Coiled and run the join there in 3 steps:

Spin up a Coiled cluster:

cluster = coiled.Cluster(
     n_workers=50,
     worker_memory='16Gib',
     spot_policy="spot_with_fallback",
)

Connect Dask to the running cluster:

client = cluster.get_client()

Run the massive join on 50 workers’ multiple cores in parallel:

large_join = large.merge(
     also_large, 
     how="left", 
     left_index=True, 
     right_index=True
)

%%time
joined = large_join.persist()
distributed.wait(joined)

# CPU times: user 385 ms, sys: 38.5 ms, total: 424 ms
# Wall time: 14.6 s

This Coiled cluster with 50 workers can run a join of two 35GB DataFrames in 14.6s.

Dask Dataframe Merge Summary

You can merge a Dask DataFrame to a small pandas DataFrame using the merge method. This is an embarrassingly parallel problem that requires little to no extra overhead compared to a regular pandas join.
You can merge two large Dask DataFrames using the same merge method. Think carefully about whether to run an unsorted join or a sorted join using set_index first to speed up the join.
For very large joins, delegate computations to a Dask cluster in the cloud with Coiled for high-performance joins using parallel computing.

How to Convert a pandas Dataframe into a Dask Dataframe

2022-01-01T00:00:00+00:00

In this post, we will cover:

How (and when) to convert a pandas DataFrame into a Dask DataFrame;
Demonstrating 2x (or more!) speedup with an example;
Discuss some best practices.

pandas is a very powerful Python library for manipulating and analyzing structured data, but it has some limitations. pandas does not leverage parallel computing capabilities of modern systems and throws a MemoryError when you try to read larger-than-memory data. Dask can help solve these limitations. Dask is a library for parallel and distributed computing in Python. It provides a collection called Dask DataFrame that helps parallelize the pandas API.

This data is stored in a public bucket and you can use Coiled for free if you’d like to run these computations yourself.

You can follow along with the code in this notebook!

Analyzing library data with pandas

Let’s start with a sample pandas workflow of reading data into a pandas DataFrame and doing a groupby operation. We read a subset of the Seattle Public Library checkouts dataset (~4GB) and find the total number of items checked out based on the ‘Usage Class’ (i.e., physical vs. digital).

df = pd.read_csv("checkouts-subset.csv") #

df.groupby("UsageClass").Checkouts.sum() # ~1.2 seconds

The computation takes about 1.2 seconds on my computer. Can we speed this up using Dask DataFrame?

Parallelizing using Dask DataFrame

To leverage single-machine parallelism for this analysis, we can convert the pandas DataFrame into a Dask DataFrame. A Dask DataFrame consists of multiple pandas DataFrames split across an index as shown in the following diagram:

![](/blog/images/webflow-images/638657ff104d4050037d09ec_DataFrame Image - Blue 2@3x.png)

We can use Dask’s from_pandas function for this conversion. This function splits the in-memory pandas DataFrame into multiple sections and creates a Dask DataFrame. We can then operate on the Dask DataFrame in parallel using its pandas-like interface. Note that we need to specify the number of partitions or the size of each chunk while converting.

ddf = dd.from_pandas(df, npartitions=10)

This conversion might take a minute, but it’s a one-time cost. We can now perform the same computation as earlier:

ddf.groupby("UsageClass").Checkouts.sum().compute() # ~679 ms
 seconds

This took ~680 milliseconds on my machine. That’s more than twice as fast as pandas!

It’s important to note here that Dask DataFrame provides parallelism, but it may not always be faster than pandas. For example, operations that require sorting take more time because data needs to be moved between different sections. We suggest sticking with pandas if your dataset fits in memory.

If you work with sizable data, it is likely stored on a remote system across multiple files, and not on your local machine. Dask allows you to read this data in parallel, which can be much faster than reading with pandas.

You can read data into a Dask DataFrame directly using Dask’s read_csv function:

import dask.dataframe as dd
ddf = dd.read_csv("s3://coiled-datasets/checkouts-subset.csv")

Both pandas and Dask also support several file-formats, including Parquet and HDF5, data formats optimized for scalable computing. You can check out the list of formats in the pandas documentation.

Using Dask Distributed Scheduler

Dask DataFrame uses the single-machine threaded scheduler by default, which works well for local workflows. Dask also provides a distributed scheduler with a lot more features and optimizations. It can be useful even for local development. For instance, you can take advantage of Dask’s diagnostic dashboards. To use the distributed scheduler, use:

from dask.distributed import Client
client = Client()

Then, perform your Dask computations as usual:

import dask.dataframe as dd
ddf = dd.read_csv("https://coiled-datasets.s3.us-east-2.amazonaws.com/seattle-library-checkouts/checkouts-subset.csv")
ddf.groupby("UsageClass").Checkouts.sum().compute()

Dask is also capable of distributed computing in the cloud and we’re building a service, Coiled, that allows you to deploy Dask clusters on the cloud effortlessly. Coiled offers a free tier with up to 100 CPU cores and 1000 CPU hours of computing time per month. Try it out today, and let us know how it goes!

How Coiled sets memory limit for Dask workers

2022-01-01T00:00:00+00:00

Problem

While running workloads to test Dask reliability, we noticed that some workers were freezing or dying when the OS stepped in and started killing processes when the system ran out of memory.

Having Dask workers die from memory overuse is common, so we thought that we’d investigate this further.

Background

Dask manages memory, so it needs to know how much memory it has available.

Setting the memory limit correctly is important to the performance of your Dask cluster. This memory limit affects when Dask spills data in memory to disk, and when it restarts workers that are out of memory.

When the memory limit is set too low, Dask won’t use as much memory as it could. That’s not the worst thing, but it means your work will potentially take longer than it needs to.

When the memory limit is set too high, Dask can fail to spill to disk when needed and fail to restart workers when they’re out of memory. This is bad because if a worker isn’t killed when it’s out of memory, then it can freeze and potentially your cluster will be stuck waiting for the work from this worker. If the worker is restarted, then Dask will retry the work.

Dask attempts to determine the correct memory limit to use. It does this by looking at the total system memory (including virtual memory) and the cgroup memory limit if set. (Dask also looks at RSS rlimit if set). You can also explicitly set the memory limit. Dask then uses the minimum of these values.

The upshot is that in many deployments—including how we had been deploying Dask with Coiled—the memory limit will be set to the total system memory. That’s not ideal because that limit is too high: we’re also running an OS and logging agents.

Solution

To address this, Coiled now sets the cgroup memory limit on the containers running dask. We set this to the available memory on the system before Dask starts running.

Here’s how we’ve implemented this.

First, our script that starts Dask on each VM now sets an environment variable with the available memory:

Copy

export AVAILABLE_MEMORY_FOR_DASK=$(grep -i memavailable /proc/meminfo | sed 's/[a-zA-Z]+:\s+//')

Then, we’ve modified the Docker compose file we use to start Dask (we run Dask in containers, and configure the containers using Docker compose). We’ve added a line to explicitly set mem_limit for the container to ${AVAILABLE_MEMORY_FOR_DASK:-0}.

We could have just used $AVAILABLE_MEMORY_FOR_DASK, but instead we’re setting a default value of 0 (i.e., no limit) just in case the environment variable isn’t set. Mostly that’s unnecessary, but it makes deployment safer because the script that sets the environment variable is part of our custom AMI while the Docker compose file is created for each cluster by our control-plane.

With a Dask client connected to your Coiled (or other Dask) cluster, you can check the value of the memory limit Dask is using by running:

Copy

client.run(lambda dask_worker: dask_worker.memory_manager.memory_limit)

The upshot is that since we’ve made this change, Dask running on Coiled now does a better job spilling to disk and restarting workers when appropriate, and we’re seeing increased reliability from our testing.

Filtering Dask DataFrames with loc

2022-01-01T00:00:00+00:00

This post explains how to filter Dask DataFrames based on the DataFrame index and on column values using loc.

Filtering Dask DataFrames can cause data to be unbalanced across partitions which isn’t desirable from a performance perspective. This post illustrates how filtering can cause the “empty partition problem” and how to eliminate empty partitions with repartitioning.

The blog starts with minimal complete verifiable examples (MVCEs) filtering examples on small datasets and then shows how to filter DataFrames with millions of rows of data on a cluster.

This data is stored in a public bucket and you can use Coiled for free if you’d like to run these computations yourself.

Here’s the link to the notebook with all the code snippets used in this post.

Dask Filter: Index filtering

Dask DataFrames consist of multiple partitions, each of which is a pandas DataFrame. Each pandas DataFrame has an index. Dask allows you to filter multiple pandas DataFrames on their index in parallel, which is quite fast.

Let’s create a Dask DataFrame with 6 rows of data organized in two partitions.

import pandas as pd
import dask.data
frame as dddf = pd.DataFrame({"nums":[1, 2, 3, 4, 5, 6], "letters":["a", "b", "c", "d", "e", "f"]})
ddf = dd.from_pandas(df, npartitions=2)

Let’s visualize the data in each partition.

for i in range(ddf.npartitions):
	print(ddf.partitions[i].compute())
	nums letters
0     1       a
1     2       b
2     3       c   
	nums letters
3     4       d
4     5       e
5     6       f

Dask automatically added an integer index column to our data.

Grab rows 2 and 5 from the DataFrame.

ddf.loc[[2, 5]].compute()
	nums letters
2     3       c
5     6       f

Grab rows 3, 4, and 5 from the DataFrame.

ddf.loc[3:5].compute()
	nums letters
3     4       d
4     5       e
5     6       f

Let’s learn more about how Dask tracks information about divisions in sorted DataFrames to perform loc filtering efficiently.

Dask Filter: Divisions

Dask is aware of the starting and ending index value for each partition in the DataFrame and stores this division’s metadata to perform quick filtering.

You can verify that Dask is aware of the divisions for this particular DataFrame by running ddf.known_divisions and seeing it returns True. Dask isn’t always aware of the DataFrame divisions.

Print all the divisions of the DataFrame.

ddf.divisions(0, 3, 5)

Take a look at the values in each partition of the DataFrame to better understand this division’s output.

for i in range(ddf.npartitions):
	print(ddf.partitions[i].compute())  
	nums letters
0     1       a
1     2       b
2     3       c   
	nums letters
3     4       d
4     5       e
5     6       f

The first division is from 0-3, and the second division is from 3-5. This means the first division contains rows 0 to 2, and the last division contains rows 3 to 5.

Dask’s division awareness in this example lets it know exactly what partitions it needs to fetch from when filtering.

Dask Filter: Column value filtering

You won’t always be able to filter based on index values. Sometimes you need to filter based on actual column values.

Fetch all rows in the DataFrame where nums is even.

ddf.loc[ddf["nums"] % 2 == 0].compute()
    nums letters
    1   2   b
    3   4   d
    5   6   f

Find all rows where nums is even, and letters contains either b or f.

ddf.loc[(ddf["nums"] % 2 == 0) & (ddf["letters"].isin(["b", "f"]))].compute()
    nums letters
    1   2   b
    5   6   f

Dask makes it easy to apply multiple logic conditions when filtering.

Dask Filter: Empty partition problem

Let’s read a dataset with 662 million rows into a Dask DataFrame and perform a filtering operation to illustrate the empty partition problem.

Read in the data and create the Dask DataFrame.

ddf = dd.read_parquet(
    "s3://coiled-datasets/timeseries/20-years/parquet",
    storage_options={"anon": True, 'use_ssl': True})

ddf.npartitions shows that the DataFrame has 1,095 partitions.

Run ddf.head() to take a look at the data.

Here’s how to filter the DataFrame to only include rows with an id greater than 1150.

res = ddf.loc[ddf["id"] > 1150]

Run len(res) to see that the DataFrame only has 1,103 rows after this filtering operation. This was a big filter, and only a small fraction of the original 662 million rows remain.

We can run res.npartitions to see that the DataFrame still has 1,095 partitions. The filtering operation didn’t change the number of partitions.

Run res.map_partitions(len).compute() to visually inspect how many rows of data are in each partition.

  0
 1
 0
 0
 2
    ..
  0
  2
  0
  0
  0
Length: 1095, dtype: int64

A lot of the partitions are empty and others only have a few rows of data.

Dask often works best with partition sizes of at least 100MB. Let’s repartition our data to two partitions and persist it in memory.

res2 = res.repartition(2).persist()

Subsequent operations on res2 will be really fast because the data is stored in memory. len(res) takes 57 seconds whereas len(res2) only takes 0.3 seconds.

The filtered dataset is so small in this example that you could even convert it to a pandas DataFrame with res3 = res.compute(). It only takes 0.000011 seconds to execute len(res3).

You don’t have to filter datasets at the computation engine level. You can also filter at the database level and only send a fraction of the data to the computation engine.

Dask Filter: Query pushdown

Query pushdown is when you perform data operations before sending the data to the Dask cluster. Part of the work is “pushed down” to the database level.

See the Advantages of Parquet blog post for information on column pruning and predicate pushdown filtering, query pushdown for the Parquet file format.

Here’s the high-level process for filtering with a Dask cluster:

Read all the data from disk into the cluster
Perform the filtering operation
Repartition the filtered DataFrame
Possibly write the result to disk (ETL style workflow) or persist in memory

Organizations often need to optimize data storage and leverage query pushdown in a manner that’s optimized for their query patterns and latency needs.

Dask Filter: Conclusion

Dask makes it easy to filter DataFrames, but you need to be cognizant of the implications of big filters.

After filtering a lot of data, you should consider repartitioning and persisting the data in memory.

You should also consider filtering at the database level and bypassing cluster filtering altogether. Lots of Dask analyses run slower than they should because a large filtering operation was performed, and the analyst is running operations on a DataFrame with tons of empty partitions.

Enterprise Dask Support

2022-01-01T00:00:00+00:00

Along with the Cloud SaaS product, Coiled sells enterprise support for Dask. Mostly people buy this for these three things:

Office hours with Dask experts
Private GitHub issue tracker
Commitment to quickly fix Dask bugs that affect critical workloads should they arise

Support contracts also come with a variety of other things, like indemnity insurance, response time SLAs, and so on, but these are usually mostly to satisfy procurement checklists rather than drivers on the business side.

Our contracts list everything out in detail, but I often find myself explaining to customers how things work and why we structure things the way we do. This article will be a more informal description of how we operate and why.

Who needs Enterprise Dask Support?

This tends to be larger companies that are moderately far along in their Dask journey. They usually have a few groups using Dask and it’s found its way into critical workloads. Often Dask has spread throughout the organization faster than their Dask expertise has grown, and so there are a handful of Dask experts who are somewhat overburdened trying to keep everyone happy and all critical workloads humming nicely.

We charge mid-six-figure annual deal sizes for our support contracts, so these are usually larger companies.

How deep do you get?

We need to strike a balance between getting deep enough with your problem to understand how to help, while staying distant enough that we focus on Dask (what we’re good at) rather than getting sucked into internal customer systems. Usually the customers who want enterprise support contracts are pretty sensitive about having us deeply in their systems, so everyone is aligned with finding good boundaries of communication.

As a result, we do not offer “hands-on-keyboard” support, where we would go into your company, look at your data, and build ML models / data engineering pipelines for you. Honestly, we’d be pretty bad at that anyway. It’s not our skillset. When folks do need this level of support we bring in consultants that we know and trust. Most often we recommend Quansight who are great at general purpose open source Python work.

Instead, we find that a combination of a few office hours a month, an issue tracker, and all-you-can-eat OSS bugfixing works well in combination. Let’s talk through these three below:

1 – Office Hours

It’s useful to talk to an expert. Many of the problems that our users face can be resolved by pointing them to a project or API of which they weren’t aware. It’s common that the right 30 minute meeting can accelerate weeks of work.

Contracts typically include a finite number of office hours per month, somewhere between two and ten. Folks can use these however they like, but a typical allocation of ten hours looks like the following:

One open office hour every week (4hrs)
A lunch and learn style lecture on a topic facing the group every month (1hr + 1hr prep)
A deeper dive on a harder problem encountered by some internal team, often including screensharing going through performance dashboards, logs, etc. (2hr)

These hours are designed to uncover issues and transfer knowledge. Often they result in issues raised for work upstream. This is where the bulk of the value comes.

2 – Dask Critical BugFixing

While office hours are finite and precious, Dask bugfixing is managed in more of an “all-you-can-eat” fashion. Often in office hours we’ll find that Dask was underperforming in some way (Dask is far from perfect) and we can raise this as an issue to be resolved upstream.

Once we’ve found work to do on the public open source project we get access to a much larger team to work on your problems. As a company we’re mostly focused around discovering high impact work, prioritizing that work, and executing it swiftly. If we can solve your problem by enhancing Dask then all of our incentives are very well aligned.

If you think in terms of hourly rates, the office hours are expensive (thousands of dollars per hour) while the Dask OSS work is very cheap (often $100 per hour or less). This reflects our perception of costs and our priorities. We spend the office hour time figuring out how to align our company and yours, and then we operate more as partners with a shared goal of improving the open source software.

Of course, even this isn’t entirely free. Our OSS team is large and skilled, but still finite. We also prioritize based on importance across users. Financially, we’re incentivized to provide enough value to customers so that they want to renew at the end of the year. The larger the contract, the more value we know we need to provide.

3 – Private Github Issue Tracker

We need a way to communicate on the ongoing technical work. We find that Github is the easiest and least distracting way to do this. Being outside of the open source tracker provides some privacy for sensitive communication, and also makes it easier for us to prioritize and track things without them getting lost.

We also jump on video calls on an as-needed basis and, very rarely, have slack channels open with customers. We’ve found that Github works best though to keep things private and also well-tracked.

What we don’t do

We’re very efficient at diagnosing Dask issues, but inefficient at diagnosing issues related to customers’ internal systems. This comes up in a few scenarios:

“Pods in our in-house Kubernetes system aren’t coming up. It says that it can’t load the image. Why?”

If we don’t manage your Dask deployment solution then we really can’t do much to help you here. Probably we don’t even have the permissions to diagnose this. We encourage you to try out Coiled’s deployment solution. It’s pretty slick.
“Our data scientist ran this giant notebook and it’s slow. Can you make it faster?”

This may have more to do with your code than ours. Your team is likely more efficient at figuring out where the issue is than we are, and they’re a lot cheaper.

However, we could help here in a few ways:

a. Getting on a call and looking at Dask diagnostics together, helping to pinpoint what might be going wrong. b. Teaching your team more about Dask diagnostics and how to use them to identify performance issues. c. Providing general Python performance advice during office hours.

But we need to be clear that custom workload performance is out-of-scope for the all-you-can-eat model. We can provide advice and training, but ultimately custom workloads are the customer’s responsibility. If the customer is able to provide a minimal reproducer of the issue or benchmark then we can more effectively take on ownership.
“We’ve made an in-house library around Dask to train models. Can you help us to maintain it?”

We would likely try to direct you to open source technologies that might do similar things. Alternatively, we could connect you with excellent consultants who can be more hands-on.
“We’ve found an issue in dask-kubernetes/dask-gateway/dask-cloudprovider. Can you fix it?”

This invariably leads us into Problem #1. We are happy to go over configuration on these tools during office hours, but don’t include them in our all-you-can-eat structure. We’ve found that they’re hard to support in deep enterprise cases. This is why we built the Coiled deployment platform.
“We have many new people to Dask, can you answer all of their questions on Slack?”

Sure, let’s use an office hour or two and give some broad training. However, we can’t be available for lots of novice Q&A. Instead, we ask that your company selects a small number of folks to route these questions through to us. We find that those folks are usually pretty Dask savvy and are able to provide basic level-1 support. Please think of us more as L2/L3 support.

Limited Seats

As the last question implies, we typically limit the number of folks that we talk to in an organization. This tends to be 1-2 people per project (companies paying several hundred thousand annually typically have a few groups using Dask). These people become deep experts and help to buffer us from lots of small problems. As Dask expands in the organization and we add more projects we also add more contact seats and typically start to ask for more money (which is often fine, because the new projects often come in with new budgets).

Sometimes, we also do new development

Sometimes companies come to us with paid feature development requests like “Hey, can you integrate Dask with Library X?”. We consider these, but reject more than we accept. We become much more excited if these are already on our roadmap, and the company wants to help us with design and testing as more of a partner.

Aligned Incentives

Over the years Dask has been funded by a wide variety of mechanisms. I really like the support model laid out above. It does a good job of delivering solid value to customers, while also aligning incentives to help make the project better. Everyone focuses on their core skillset and generally has a good time.

Dask has been developed in this manner for many years. It would not be nearly as pragmatic without the valuable support and technical insight of its support partners throughout the years.

Easily Run Python Functions in Parallel

2022-01-01T00:00:00+00:00

When you search for how to run a Python function in parallel, one of the first things that comes up is the multiprocessing module. The documentation describes parallelism in terms of processes versus threads and mentions it can side-step the infamous Python GIL (Global Interpreter Lock).

This is all great if you’re a Python developer comfortable navigating the intricacies of processes and threads. But what if you don’t want to care about any of that? What if all you want is a straightforward way to run a Python function in parallel?

Different models of concurrency, adapted from https://realpython.com/python-concurrency

Imagine you’re estimating the global burden of malaria. You don’t want to dive into the technicalities of parallelism—you just want to answer a question like, “How many people died of malaria in 1984?”

That’s where Dask comes in.

Why Choose Dask Over Multiprocessing?

Dask is designed to make parallel computing in Python as seamless as possible. With Dask, you can:

Run Python functions in parallel with minimal code changes.
Easily scale from a single machine to a distributed cluster
Focus on your data and analysis instead of debugging processes and threads.

Whether you’re analyzing global health data, processing large datasets, or running computational models, Dask lets you spend more time on your problem and less time wrestling with low-level parallelism.

Run a Python Function in Parallel

You can use Dask to run any Python function in parallel. Let’s say you have a function called costly_simulation that defines a long-running simulation over a set of parameters. With Dask, we can run this in parallel locally on all cores available on your laptop:

from dask.distributed import LocalCluster

cluster = LocalCluster()                                # Use all cores on local machine
client = cluster.get_client()

parameters = [...]
def costly_simulation(parameter):
    ...

futures = client.map(costly_simulation, parameters)     # Run simulation in parallel
results = client.gather(futures)

Or scale out to a cluster of many machines on the cloud using Coiled:

import coiled

cluster = coiled.Cluster(n_workers=50)                  # Start 50 machines on AWS        
client = cluster.get_client()

parameters = [...]
def costly_simulation(parameter):
    ...

futures = client.map(costly_simulation, parameters)     # Run simulation in parallel
results = client.gather(futures)

Coiled runs from your AWS, Google Cloud, or Azure account and you’ll get 500 free CPU hours each month. It’s free to sign up.

Parallel For Loop with Dask

Here’s an illustrative example doing a nested for loop over parameters and data files. In this example, we work sequentially, loading the data in the outer loop and scoring in the inner loop for all pairs.

def score(params: dict, data: object) -> float:
    ...

results = []
for filename in filenames:             # Nested for loop
    data = load(filename)              # Load data in outer loop
    for param in params:               # Score in inner loop over all pairs
        result = score(param, data)  
        results.append(result)

best = max(results)                    # Get best score

‍

Working sequentially can be quite slow, especially if you’re working through hundreds of files. We can use Dask to run this in parallel, using all cores available on your laptop:

from dask.distributed import LocalCluster

cluster = LocalCluster()                 # Use all cores on local machine
client = cluster.get_client()

futures = []
for filename in filenames:               # Nested for loop
    data = client.submit(load, filename)  # Load data in outer loop
    for param in params:                 # Score in inner loop over all pairs
        future = client.submit(score, param, data)      
        futures.append(future)

results = client.gather(futures)

best = max(results)                      # Get best score

‍

If your data is stored on the cloud or you need access to hardware you don’t have available locally (like more memory or GPUs) you can use Coiled to run the same workflow across a cluster of VMs:

import coiled

cluster = coiled.Cluster(n_workers=100)   # Scale out to 100 machines
client = cluster.get_client()

futures = []
for filename in s3_filenames:            # Nested for loop
    data = client.submit(load, filename)  # Load data in outer loop
    for param in params:                 # Score in inner loop over all pairs
        future = client.submit(score, param, data)      
        futures.append(future)

results = client.gather(futures)

best = max(results)                      # Get best score

Coiled handles things like:

Python environment management
Controlling cloud costs
Combining hardware and Python-specific metrics into a single UI

When to Use Multiprocessing Instead of Dask?

While Dask is powerful and flexible, there are cases where multiprocessing might be a better choice:

Standard Library: As part of Python’s standard library, multiprocessing requires no additional installation or setup.
Low Overhead: For smaller tasks, the overhead of Dask’s scheduler might outweigh its benefits, making multiprocessing a more efficient option. Dask also has a multiprocessing scheduler that is very lightweight for these situations.
Fine-Grained Control: If you need detailed control over processes, shared memory, or inter-process communication, multiprocessing provides the tools you need. Though it’s worth mentioning Dask is also flexible enough for this level of fine-grained tuning with tools like asynchronous futures, distributed locks, queues, etc.

Examples

If you’re tired of fighting with multiprocessing or just want an easier way to run Python functions in parallel, Dask could be a good option. It’s powerful, intuitive, and designed to help you focus on your work—not the underlying mechanics of parallel computing. For more examples using Dask for parallel Python you might consider the following:

Dask on GCP

2022-01-01T00:00:00+00:00

Deployment challenges
Software environment and credential management
Hardware flexibility
Historical tracking and observability
Cost management
Cost optimization

‍

How to Deploy a Dask Cluster

There are a number of tools you can use to manage Dask clusters, including:

Install Packages and Sync Credentials

Use GPUs, ARM, and Other Hardware

import coiled

cluster = coiled.Cluster(
    n_workers=100,
    worker_cpu=8,
    arm=True,
    region="us-east1",
    spot_policy="spot_with_fallback",
)

Coiled’s cluster configuration flexibility extends to AWS, GCP, and Azure allowing you to dynamically start machines of any type in any region, a powerful feature for hardware experimentation.

Getting Started with Dask on GCP

Create Dask clusters on Google Cloud by connecting Coiled to your GCP account. Coiled starts virtual machines when you need them and turns them off when you’re done. It’s easy to set up Coiled to run from your GCP account:

$ pip install coiled
$ coiled login
$ coiled setup gcp

Or you can skip the CLI and get started at cloud.coiled.io:

Observability Into Your Dask Cluster

Managing Cloud Costs

Optimizing Cloud Costs

Though preemptible instances are heavily discounted, most users are willing to pay more to avoid interruptions to their workflows as preemptible instances become unavailable. With Coiled, when you request preemptible instances, we look at your region and select the data center in that region with the best availability. Once we get the warning that these preemptible instances are being deprovisioned, we move your tasks onto new preemptible or standard instances. This ensures the continuity of your workloads without interruption. Learn more in our blog post.

Conclusion

It’s easy to get started with Dask on GCP, AWS, or Azure.

Dask on Azure

2022-01-01T00:00:00+00:00

Deployment challenges
Software environment and credential management
Hardware flexibility
Historical tracking and observability
Cost management
Cost optimization

‍

How to Deploy a Dask Cluster

There are a number of tools you can use to manage Dask clusters, including:

Install Packages and Sync Credentials

Use GPUs, ARM, and Other Hardware

import coiled

cluster = coiled.Cluster(
    n_workers=100,
    worker_cpu=8,
    arm=True,
    region="us-east-1",
    spot_policy="spot_with_fallback",
)

Coiled’s cluster configuration flexibility extends to AWS, GCP, and Azure allowing you to dynamically start machines of any type in any region, a powerful feature for hardware experimentation.

Getting Started with Dask on Azure

Create Dask clusters on Azure by connecting Coiled to your Azure account. Coiled starts virtual machines when you need them and turns them off when you’re done. It’s easy to set up Coiled to run from your Azure account:

$ pip install coiled
$ coiled login
$ coiled setup aws

Or you can skip the CLI and get started at cloud.coiled.io:

Observability Into Your Dask Cluster

Managing Cloud Costs

Optimizing Cloud Costs

Though spot virtual machines are heavily discounted, most users are willing to pay more to avoid interruptions to their workflows as spot instances become unavailable. With Coiled, when you request spot instances, we look at your region and select the data center in that region with the best spot availability. Once we get the warning that these spot instances are being deprovisioned, we move your tasks onto new spot or on-demand instances. This ensures the continuity of your workloads without interruption. Learn more in our blog post.

Conclusion

It’s easy to get started with Dask on Azure, GCP, or AWS.

Dask on AWS

2022-01-01T00:00:00+00:00

Deployment challenges
Software environment and credential management
Hardware flexibility
Historical tracking and observability
Cost management
Cost optimization

How to Deploy a Dask Cluster

There are a number of tools you can use to manage Dask clusters, including:

Install Packages and Sync Credentials

Use GPUs, ARM, and Other Hardware

import coiled

cluster = coiled.Cluster(
    n_workers=100,
    worker_cpu=8,
    arm=True,
    region="us-east-1",
    spot_policy="spot_with_fallback",
)

Coiled’s cluster configuration flexibility extends to AWS, GCP, and Azure allowing you to dynamically start machines of any type in any region, a powerful feature for hardware experimentation.

Getting Started with Dask on AWS

$ pip install coiled
$ coiled login
$ coiled setup aws

Or you can skip the CLI and get started at cloud.coiled.io:

Observability Into Your Dask Cluster

Managing Cloud Costs

Optimizing Cloud Costs

Conclusion

It’s easy to get started with Dask on AWS, GCP, or Azure.

Dask for Parallel Python

2022-01-01T00:00:00+00:00

Dask is a general purpose library for parallel computing. Dask can be used on its own to parallelize Python code, or with integrations to other popular libraries to scale out common workflows.

Dask has its own docs, but we’ll include a few typical use cases below.

Big pandas

Dask DataFrames use pandas under the hood, so your current code likely just works. It’s faster than Spark and easier too. Here’s an illustrative example of how you would load data from Parquet and perform some standard data manipulation with pandas:

import pandas as pd

df = pd.read_parquet('s3://mybucket/myfile.parquet')

df = df[df.value >= 0]
joined = df.merge(other, on="account")
result = joined.groupby("account").value.mean()

Dask does pandas in parallel. Dask is lazy; when you want an in-memory result add .compute(). Persist data in distributed memory with the .persist() method.

import dask.dataframe as dd

df = dd.read_parquet('s3://mybucket/myfile.*.parquet')

df = df[df.value >= 0]
joined = df.merge(other, on="account")
result = joined.groupby("account").value.mean()

result.compute()

Dask comfortably scales to handle datasets that range in size from 10 GiB to 100 TiB. Dask DataFrame is used across a wide variety of applications—anywhere you’re working with a large tabular dataset.

Here are a few large-scale examples using Coiled to deploy Dask on the cloud:

Parallel For Loops

With Dask, you can parallelize any Python code, no matter how complex. Dask is flexible and supports arbitrary dependencies and fine-grained task scheduling that extends Python’s concurrent.futures interface. Dask futures allow you to scale generic Python workflows across a Dask cluster with minimal code changes.

from dask.distributed import LocalCluster

cluster = LocalCluster()                # Runs locally, but you can deploy Dask anywhere
client = cluster.get_client()


def process(filename):                  # Define your own Python function
   ...

futures = []
for filename in filenames:              # Submit many tasks in a loop
    future = client.submit(process, filename)
    futures.append(future)

results = client.gather(futures)        # Wait until done and gather results

As your data grows or your computational demands increase, you might want to deploy Dask on a cluster. If you’re using the cloud, this is easy to do with Coiled:

from coiled import Cluster

cluster = coiled.Cluster(n_workers=100)    # Request 100 VMs on AWS 
client = cluster.get_client()


def process(filename):                     # Define your own Python function
   ...

futures = []
for filename in filenames:                 # Submit many tasks in a loop
    future = client.submit(process, filename)
    futures.append(future)

results = client.gather(futures)           # Wait until done and gather results

Here’s an example extracting text from PDFs stored in a public S3 bucket, where we use Dask to scrape 1 TiB of data in ~5 minutes.

Big Xarray

Dask integrates well with Xarray to process array data in parallel. This includes common file formats like HDF, NetCDF, TIFF, or Zarr. Working with Dask and Xarray feels much like using NumPy arrays, but enables handling significantly larger datasets. The integration between Dask and Xarray is seamless, so you rarely need to manage parallelism explicitly—these details are handled automatically.

In the following example, we can import Zarr data and perform an aggregation without explicitly importing Dask at all. If Dask is available, it will be used automatically to process the data in parallel.

import xarray as xr

ds = xr.open_zarr("/path/to/data.zarr")
timeseries = ds["temp"].mean(dim=["x", "y"]).compute()  # Compute result

This allows you to write scalable code that transitions from small, in-memory datasets on a single machine to large, distributed datasets on a cluster, requiring only minimal adjustments.

Here’s an example aggregating 250 TiB of raster data from the NOAA National Water Model to calculate the average water table depth for each county in the US.

Dask for Machine Learning

Dask integrates with a number of machine learning libraries to train or predict on large datasets, increasing model accuracy by using all of your data.

Here’s an illustrative example using Dask with XGBoost for distributed model training:

import xgboost as xgb
import dask.dataframe as dd

df = dd.read_parquet("s3://my-data/")
dtrain = xgb.dask.DaskDMatrix(df)

model = xgb.dask.train(
    dtrain,
    {"tree_method": "hist", ...},
    ...
)

Model training can often be both computationally- and memory-intensive. Here are a few examples of speeding up machine learning tasks by taking advantage of cloud resources with Coiled:

Deploying Dask

Many of these examples process larger-than-memory datasets on Dask clusters deployed on the cloud with Coiled. If your data is stored on the cloud or you need access to more/better hardware, you can use Coiled to run Dask on AWS, GCP, or Azure. See our documentation on all the ways you can deploy Dask for more details.

‍

Dask and the PyData Stack

2022-01-01T00:00:00+00:00

TL;DR

The PyData stack for scientific computing in Python is an interoperable collection of tools for data analysis, visualization, statistical inference, machine learning, interactive computing and more that is used across all types of industries and academic research. Dask, the open source package for scalable data science that was developed to meet the needs of modern data professionals, was born from the PyData community and is considered foundational in this computational stack. This post describes a schematic for thinking about the PyData stack, along with detailing the technical, cultural, and community forces that led to Dask becoming the go-to package for scalable analytics in Python, including its interoperability with packages such as NumPy and pandas, among many others.

The Python Data Science Stack

In 2017, Jake VanderPlas gave a keynote at PyCon called “The Unreasonable Effectiveness of Python in Science”, in which he presented a schematic of the scientific Python stack, also commonly referred to as the PyData stack. This is a broad collection of Python packages used for scientific computing and data science, in both industry and basic research:

At the bottom of the stack, you have the Python standard library. As you move out, the next layer extends Python in ways that are useful to scientists and foundational for data science in Python, such as

NumPy for numerical array computing,
Project Jupyter and IPython for interactive computing and IDEs,
Numba for just-in-time compilation, and
Dask for scaling your data science to larger datasets and larger models.

Jake stated that “built on top of this layer are tools that are slightly more specialized”, such as matplotlib and bokeh for data visualization, pandas for data ingestion and wrangling, Xarray for multidimensional labelled structures (these, as Jake made clear, depend on the level below them, and are “higher level tools”).

Moving up to even higher level tools, we discover scikit-learn for machine learning, statsmodels for statistical modeling, and PyMC for Bayesian inference and probabilistic programming, among others. On top of these, you have field-specific packages like astropy and biopython (for astronomy and biology, respectively).

Dask is Foundational in the PyData stack

If you’re asking “Why is Dask foundational in the PyData stack?”, you’re asking a great question. The answer is: because it composes well with all of the PyData tools. It’s a simple system designed to parallelize any PyData library. This means that, if you’re looking to natively scale up your Python code to larger datasets within the PyData ecosystem, you’ll be using Dask. One result of this is the elegance of the Dask API:

if you’re writing Dask code to scale your NumPy code, what you write will mimic the NumPy API, providing a smooth transition;
Similarly, if you’re writing Dask code to scale your pandas code, what you write will mimic the pandas API.

Dask arrays scale Numpy workflows, enabling multi-dimensional data analysis in earth science, satellite imagery, genomics, biomedical applications, and machine learning algorithms (from the Dask documentation).

Matt Rocklin, creator of Dask and CEO of Coiled, once said to me:

People really like that Dask Array follows the NumPy API almost exactly. Dask DataFrame follows the Pandas API almost exactly. Dask-ML follows the scikit-learn API almost exactly. There are definitely differences but those differences are necessary because there are some things you do need to think about in parallel algorithms. The Dask maintainers have written a lot of Python data science code and we’re actually the maintainers of all those other libraries too. So it’s rare to find someone who understands NumPy and won’t be able to transition on their own to Dask Array. There’s a very smooth transition process between all the different systems.

As an example, compare the following NumPy and Dask code chunks, noticing how you use NumPy syntax when working with your Dask array:

import numpy as np
x = np.random.random((1000, 1000))  # Create array of random numbers
y = x + x.T  # Add x to its transpose
z = y[::2, 500:].mean(axis=1)  # Compute mean of a slice of Y
import dask.array as da
x = da.random.random((10000, 10000), chunks=(1000, 1000))  # Create Dask array
y = x + x.T  # Add x to its transpose
z = y[::2, 5000:].mean(axis=1)  # Compute mean of a slice of Y

Not only is the end-user programming experience of a Dask Array (or DataFrame) “like” the NumPy or Pandas equivalent, it actually delegates to NumPy (or pandas) code running against real NumPy (or pandas) objects as much as possible. Why is this important? It means, for example, that pandas users’ mental model of what’s happening is close and, in some cases, identical to what actually runs, making a user more productive with less work. By contrast, some other tools – like Koalas (a wrapper over Spark) – imitate the Pandas API but delegate to vastly different machinery underneath. In the case of Spark, the real challenge is understanding enough of that machinery to work with it well, and Koalas can make a tough situation even harder, by hiding the traps and pitfalls until a user falls in.

This consistency of Dask with the PyData is both (1) by design and (2) a result of cultural and community forces. With respect to the design of Dask, there were several technical and social goals for Dask:

Technical: Harness the power of all of the cores on the laptop/workstation in parallel;
Technical: Support larger-than-memory computation, allowing datasets that fit on disk, but not in RAM;
Social: Invent nothing. We wanted to be as familiar as possible to what users already knew.
Sociotechnical: Dask was designed to be lightweight technically so that it could be integrated into other Python libraries.

Now let’s dive into the cultural and community forces that led to the simplicity of Dask’s API.

The Dask Community IS the PyData Community

In terms of cultural and community forces, Dask was developed by members of the Python data science community for Pythonista data scientists. Moreover, the Dask developer community is the PyData community:

Stefan Hoyer, the creator of Xarray, hooked up Dask and Xarray in around two weeks;
Min Ragan-Kelly, from Project Jupyter and IPython parallel, and Olivier Grisel, from Scikit-Learn & Joblib, helped to build the current incarnation of the dask-distributed scheduler;
The original idea for Dask arose at Anaconda Inc. (the company that built the popular Anaconda distribution, along with many other open source tools) from conversations and work with founders Peter Wang and Travis Oliphant, the former of whom created the PyData conference series, the latter of whom created NumPy in 2005!
Tom Augspurger is a core maintainer of both Pandas and Dask so when Tom makes features in Pandas, he also makes them in Dask.

These are a mere handful of examples and there are many more. The cultural and the technical aspects of any design system are interwoven. Due to the social and technical forces outlined here, a key differentiator of Dask in the scalable data science world is that it is Python top-to-bottom: it’s Python all the way down! For high-level API calls, we may not care what’s underneath. But when we need extensibility, flexibility, and support for custom use cases, non-Pythonic abstractions such as PySpark tend to leak heavily. Many Pythonistas have gotten burned trying to use PySpark, thinking it’s a Python library, only to discover Python is a minimal veneer over a complex Scala and Java system lacking Pythonic behaviors.

The PyData stack is more than a collection of disparate packages created and maintained by different people. It is, rather, a collection of interoperable tools and packages for data science developed by a community (a cool data science project would be a network analysis of PyData stack contributors and packages, using publicly available github data). And the way to scale your data science while staying in this ecosystem is by using Dask.

Dask DataFrames scale Pandas workflows, enabling applications in time series, business intelligence, and general data munging on big data (from the Dask documentation).

Learn More

To learn more about Dask and what we’re working on at Coiled to help meet the needs of teams and the enterprise in scaling data science, analytics, and machine learning, sign up to our email list below. We’ll also be writing more posts in the future about how Dask integrates with NumPy, pandas, popular machine learning such as scikit-learn and xgboost, and deep learning frameworks, including PyTorch and tensorflow so keep your eyes peeled for them!

Dask Read Parquet Files into DataFrames with read_parquet

2022-01-01T00:00:00+00:00

This blog post explains how to read Parquet files into Dask DataFrames. Parquet is a columnar, binary file format that has multiple advantages when compared to a row-based file format like CSV. Luckily Dask makes it easy to read Parquet files into Dask DataFrames with read_parquet.

It’s important to properly read Parquet files to take advantage of performance optimizations. Disk I/O can be a major bottleneck for distributed compute workflows on large datasets. Reading Parquet files properly allows you to send less data to the computation cluster, so your analysis can run faster.

Let’s look at some examples on small datasets to better understand the options when reading Parquet files. Then we’ll look at examples on larger datasets with thousands of Parquet files that are processed on a cluster in the cloud.

Dask read_parquet: basic usage

Let’s create a small DataFrame and write it out as Parquet files. This will give us some files to try out read_parquet. Start by creating the DataFrame.

import dask.dataframe as dd
import pandas as pd

df = pd.DataFrame(
    {"nums": [1, 2, 3, 4, 5, 6], "letters": ["a", "b", "c", "d", "e", "f"]}
)
ddf = dd.from_pandas(df, npartitions=2)

Now write the DataFrame to Parquet files with the pyarrow engine. The installation instructions for pyarrow are in the Conda environment for reading Parquet files section that follows.

ddf.to_parquet("data/something", engine="pyarrow")

Here are the files that are output to disk.

data/something/
  _common_metadata
  _metadata
  part.0.parquet
  part.1.parquet

You can read the files into a Dask DataFrame with read_parquet.

ddf = dd.read_parquet("data/something", engine="pyarrow")

Check the contents of the DataFrame to make sure all the Parquet data was properly read.

ddf.compute()

Dask read Parquet supports two Parquet engines, but most users can simply use pyarrow, as we’ve done in the previous example, without digging deep into this option.

Dask read_parquet: pyarrow vs fastparquet engines

You can read and write Parquet files to Dask DataFrames with the fastparquet and pyarrow engines. Both engines work fine most of the time. The subtle differences between the two engines doesn’t matter for the vast majority of use cases.

It’s generally best to avoid mixing and matching the Parquet engines. For example, you usually won’t want to write Parquet files with pyarrow and then try to read them with fastparquet.

This blog post will only use the pyarrow engine and won’t dive into the subtle differences between pyarrow and fastparquet. You can typically just use pyarrow and not think about the minor difference between the engines.

Dask read_parquet: lots of files in the cloud

Our previous example showed how to read two Parquet files on localhost, but you’ll often want to read thousands of Parquet files that are stored in a cloud based file system like Amazon S3.

Here’s how to read a 662 million row Parquet dataset into a Dask DataFrame with a 5 node computational cluster.

import coiled
import dask
import dask.dataframe as dd

cluster = coiled.Cluster(name="read-parquet-demo", n_workers=5)

client = dask.distributed.Client(cluster)

ddf = dd.read_parquet(
    "s3://coiled-datasets/timeseries/20-years/parquet",
    engine="pyarrow",
    storage_options={"anon": True, "use_ssl": True},
)

Take a look at the first 5 rows of this DataFrame to get a feel for the data.

ddf.head()

This dataset contains a timestamp index and four columns of data.

Let’s run a query to compute the number of unique values in the id column.

ddf["id"].nunique().compute()

This query takes 59 seconds to execute.

Notice that this query only requires the data in the id column. However, we transferred the data for all columns of the Parquet file to run this query. Spending time to transfer data that’s not used from the filesystem to the cluster is obviously inefficient.

Let’s see how Parquet allows you to only read the columns you need to speed up query times.

Dask read_parquet: column selection

Parquet is a columnar file format which allows you to selectively read certain columns when reading files. You can’t cherry pick certain columns when reading from row-based file formats like CSV. Parquet’s columnar nature is a major advantage.

Let’s refactor the query from the previous section to only read the id column to the cluster by setting the columns argument.

ddf = dd.read_parquet(
    "s3://coiled-datasets/timeseries/20-years/parquet",
    engine="pyarrow",
    storage_options={"anon": True, "use_ssl": True},
    columns=["id"],
)

Now let’s run the same query as before.

ddf["id"].nunique().compute()

This query only takes 43 seconds to execute, which is 27% faster. This performance enhancement can be much larger for different datasets / queries.

Cherry picking individual columns from files is often referred to as column pruning. The more columns you can skip, the more column pruning will help speed up your query.

Definitely make sure to leverage column pruning when you’re querying Parquet files with Dask.

Dask read_parquet: row group filters

Parquet files store data in row groups. Each row group contains metadata, including the min/max value for each column in the row group. For certain filtering queries, you can skip over entire row groups just based on the row group metadata.

For example, suppose columnA in row_group_3 has a min value of 2 and a max value of 34. If you’re looking for all rows with a columnA value greater than 95, then you know row_group_3 won’t contain any data that’s relevant for your query. You can skip over the row group entirely for that query.

Let’s run a query without any row group filters and then run the same query with row group filters to see the performance book predicate pushdown filtering can provide.

ddf = dd.read_parquet(
    "s3://coiled-datasets/timeseries/20-years/parquet",
    engine="pyarrow",
    storage_options={"anon": True, "use_ssl": True},
)

len(ddf[ddf.id > 1170])

This query takes 77 seconds to execute.

Let’s run the same query with row group filtering.

ddf = dd.read_parquet(
    "s3://coiled-datasets/timeseries/20-years/parquet",
    engine="pyarrow",
    storage_options={"anon": True, "use_ssl": True},
    filters=[[("id", ">", 1170)]],
)

len(ddf[ddf.id > 1170])

This query runs in 4.5 seconds and is significantly faster.

Row group filtering is also known as predicate pushdown filtering and can be applied in Dask read Parquet by setting the filters argument when invoking read_parquet.

Predicate pushdown filters can provide massive performance gains or none at all. It depends on how many row groups Dask will be able to skip for the specific query. The more row groups you can skip with the row group filters, the less data you’ll need to read to the cluster, and the faster your analysis will execute.

Dask read_parquet: ignore metadata file

When you write Parquet files with Dask, it’ll output a _metadata file by default. The _metadata file contains the Parquet file footer information for all files in the filesystem, so Dask doesn’t need to individually read the file footer for every file in the Parquet dataset every time the Parquet lake is read.

The _metadata file is a nice performance optimization for smaller datasets, but it has downsides.

_metadata is a single file, so it’s not scalable for huge datasets. For large data lakes, even the metadata can be “big data”, with the same scaling issues of “regular data”.

You can have Dask read Parquet ignore the metadata file by setting ignore_metadata_file=True.

ddf = dd.read_parquet(
    "s3://coiled-datasets/timeseries/20-years/parquet",
    engine="pyarrow",
    storage_options={"anon": True, "use_ssl": True},
    ignore_metadata_file=True,
)

Dask will gather and process the metadata for each Parquet file in the lake when it’s instructed to ignore the _metadata file.

Dask read_parquet: index

You may be surprised to see that Dask can intelligently infer the index when reading Parquet files. Dask is able to confirm the index from the Pandas parquet file metadata. You can manually specify the index as well.

ddf = dd.read_parquet(
    "s3://coiled-datasets/timeseries/20-years/parquet",
    engine="pyarrow",
    storage_options={"anon": True, "use_ssl": True},
    index="timestamp",
)

You can also read in all data as regular columns without specifying an index.

ddf = dd.read_parquet(
    "s3://coiled-datasets/timeseries/20-years/parquet",
    engine="pyarrow",
    storage_options={"anon": True, "use_ssl": True},
    index=False,
)

ddf.head()

Dask read_parquet: categories argument

You can read in a column as a category column by setting the categories option.

ddf = dd.read_parquet(
    "s3://coiled-datasets/timeseries/20-years/parquet",
    engine="pyarrow",
    storage_options={"anon": True, "use_ssl": True},
    categories=["id"],
)

Check the dtypes to make sure this was read in as a category.

ddf.dtypes

id      category
name      object
x        float64
y        float64
dtype: object

Conda environment for reading Parquet

Here’s an abbreviated Conda YAML file for creating an environment with the pyarrow and fastparquet dependencies:

name: standard-coiled
channels:
  - conda-forge
  - defaults
dependencies:
  - python=3.9
  - pandas
  - dask
  - pyarrow
  - fastparquet

You don’t need to include both pyarrow and fastparquet in your environment. Just add the Parquet engine you’ll be using.

You can use this environment when running Dask to keep your life simple.

Additional resources

Here are additional resources if you’d like to learn more:

5 Reasons Parquet is better than CSV for data analyses
Advantages of Parquet file format
How to speed up a Pandas query 10x with 6 Dask tricks, many of which involve the Parquet tactics outlined in this post

Conclusion

This blog post showed you how to properly read Parquet files with Dask.

There are a lot of options and they can impact the runtime of your analysis significantly, so knowing how to read Parquet files is quite important.

Column pruning is of particular importance. It’s easy to apply column pruning and it often yields a significant performance gain.

Creating Disk Partitioned Lakes with Dask using partition_on

2022-01-01T00:00:00+00:00

This post explains how to create disk-partitioned Parquet lakes using partition_on and how to read disk-partitioned lakes with read_parquet and filters. Disk partitioning can significantly improve performance when used correctly.

However, when used naively disk-partitioning may in fact reduce performance (because of expensive file listing operations or lots of small files). This article will build your intuition and skills around recognizing and implementing opportunities to improve performance using partition_on.

Let’s start by outputting a Dask DataFrame with to_parquet using disk partitioning to develop some intuition on how this design pattern can improve query performance.

Output disk partitioned lake

Create a Dask DataFrame with letter and number columns and write it out to disk, partitioning on the letter column.

df = pd.DataFrame(
    {"letter": ["a", "b", "c", "a", "a", "d"], "number": [1, 2, 3, 4, 5, 6]}
)
ddf = dd.from_pandas(df, npartitions=3)
ddf.to_parquet("tmp/partition/1", engine="pyarrow", partition_on="letter")

Here are the files that are written to disk.

tmp/partition/1/
  letter=a/
    part.0.parquet
    part.1.parquet
    part.2.parquet
  letter=b/
    part.0.parquet
  letter=c/
    part.1.parquet
  letter=d/
    part.2.parquet

Organizing the data in this directory structure lets you easily skip files for certain read operations. For example, if you only want the data where the letter equals a, then you can look in the tmp/partition/1/letter=a directory and skip the other Parquet files.

The letter column is referred to as the partition key in this example.

Let’s take a look at the syntax for reading the data from a certain partition.

Read disk partitioned lake

Here’s how to read the data in the letter=a disk partition.

ddf = dd.read_parquet(
    "tmp/partition/1", engine="pyarrow", filters=[("letter", "==", "a")]
)
print(ddf.compute())

   number letter
0       1      a
3       4      a
4       5      a

Dask is smart enough to apply the partition filtering based on the filters argument. Dask only reads the data from tmp/partition/1/letters=a into the DataFrame and skips all the files in other partitions.

Trade Offs

Disk partitioning improves performance for queries that filter on the partition key. It usually hurts query performance for queries that don’t filter on the partition key.

Example query that runs faster on disk partitioned lake

dd.read_parquet(
    "tmp/partition/1", engine="pyarrow", filters=[("letter", "==", "a")]
)

Example query that runs slower on disk partitioned lake

ddf = dd.read_parquet("tmp/partition/1", engine="pyarrow")
ddf = ddf.loc[ddf["number"] == 2]
print(ddf.compute())

The performance drag from querying a partitioned lake without filtering on the partition key depends on the underlying filesystem. Unix-like filesystems are good at performing file listing operations on nested directories. So you might not notice much of a performance drag on your local machine.

Cloud-based object stores, like AWS S3, are not Unix-like filesystems. They store data as key value pairs and are slow when listing nested files with a wildcard character (aka globbing).

You need to carefully consider your organization’s query patterns when evaluating the costs/benefits of disk partitioning. If you are always filtering on the partition key, then a partitioned lake is typically best. Disk partitioning might not be worth it if you only filter on the partitioned key sometimes.

Multiple partition keys

You can also partition on multiple keys (the examples thus far have only partitioned on a single key).

Create another Dask DataFrame and write it out to disk with multiple partition keys.

df = pd.DataFrame(
    [
        ["north america", "mexico", "carlos"],
        ["asia", "india", "ram"],
        ["asia", "china", "li"],
    ],
    columns=["continent", "country", "first_name"],
)
ddf = dd.from_pandas(df, npartitions=2)
ddf.to_parquet(
    "tmp/partition/2", engine="pyarrow", partition_on=["continent", "country"]
)

Here are the files that are outputted to the filesystem:

tmp/partition/2/
  continent=asia/
    country=china/
      part.0.parquet
    country=india/
      part.0.parquet
  continent=north america/
    country=mexico/
      part.0.parquet

Read in the files where continent=asia and country=china.

ddf = dd.read_parquet(
    "tmp/partition/2",
    engine="pyarrow",
    filters=[("continent", "==", "asia"), ("country", "==", "china")],
)
print(ddf.compute())

  first_name continent country
2         li      asia   china

Deeply partitioned lakes are even more susceptible to the partitioned lake performance pitfalls because file listing operations on cloud stores for nested folders can be slow. You’ll typically only be rewarded with performance gains when you filter on at least one of the partition keys.

Globbing gets more expensive as the directory structure gets more nested.

Additional considerations

Key Point

When creating a partitioned lake, you need to choose a partition key that doesn’t have too many distinct values.

Suppose you have a 5GB dataset and often filter on column_a. Should you partition your data lake on column_a?

That depends on the number of unique values in column_a. If column_a has 5 unique values, then it could be a great partition key. If column_a has 100,000 distinct values, then partitioning the 5GB dataset will break up the dataset into at least 100,000 files, which is way too much. That’d be 100,000 files that are only 0.05MB and it’s usually best to target files that are 100MB, see this blog post for more detail.

Partitioned data lakes are particularly susceptible to the small file problem. Dask analyses often work best with files that are around 100 MB when uncompressed. A data lake that’s incrementally updated frequently will often create files that are quite small. Incrementally updating partitioned lakes will create files that are even smaller.

Suppose you have a data pipeline that ingests 600 MB of data per hour. If you’re writing this data to an unpartitioned lake, then you can repartition into 100MB partitions and write out 6 files per hour. If you’re writing to a partitioned lake with 10,000 unique partition key values, then you may be writing up to 10,000 files per hour.

Make sure you have a small file compaction routine if you build a lake that’s partitioned on a high cardinality partition key and incrementally update frequently.

Conclusion

Disk partitioning is a powerful performance optimization that can make certain queries much faster.

Due to the key-value nature of cloud object stores, disk partitioning can also make certain queries a lot slower.

It’s a powerful optimization, but you need to understand your query patterns and limitations to make it work well with your data systems.

Thanks for reading! And if you’re interested in trying out Coiled Cloud, which provides hosted Dask clusters, docker-less managed software, and one-click deployments, you can do so for free today when you click below.

Cost Savings with Dask and Coiled

2022-01-01T00:00:00+00:00

Summary

Coiled can often save money for an organization running Dask. This article goes through the most common ways in which we see that happen.

We start and end with the big abstract points of “people are expensive” and “you can move faster” and fill in the middle with lots of technical details.

1 – Personnel cost savings

First and foremost, running Dask in the cloud takes effort by smart people. Smart people are expensive. We often see data science groups spending half of an FTE maintaining the cloud environment where Dask is run. This includes activities like the following:

Setting up and tending EC2 / Kubernetes clusters
Managing docker images for the team
Tracking resource usage and shutting down things that shouldn’t be running
Handling upgrades
Debugging issues from their teammates

Half-an-FTE often costs between $50k and $150k/year depending on their seniority. Allowing this person to go back to their day job is, often, the largest cost savings for groups just getting started. Learn more in our Build vs. Buy page.

2 – Track and shut down idle resources

People leave things on, either through negligence (whoops! I didn’t realize that I left this notebook open) through things slipping through the cracks (hrm, I definitely shut this down, but it looks like the signal never made it to the cloud manager) or through ignorance (what do you mean I had to clean up network resources as well as instances? We get charged for those separately?)

Coiled tracks everything and provides a zero-cost default experience, allowing people to rest easier. If, somehow, something sneaks through, our pagers go off, not yours.

This changes the way we work. As Coiled users we often don’t really care about being diligent shutting down clusters. We know that the automated systems will detect idleness and clean them up.

3 – Spot / Preemptible instances

Instances that can disappear are cheaper than instances that can’t, often by a factor of 2-3x. By leveraging spot instances we can drastically reduce costs. However, using Spot is hard, especially when you need something to work. Coiled allows the following policies:

Prefer spot, but use on-demand if no spot instances are available
As spot machines are reclaimed

a. Give Dask workers notice so that they can gracefully terminate, allowing smooth operation within Dask and no lost work b. Backfill these slots with new on-demand instances
As spot instances become available again (work in progress)

a. Gracefully retire on-demand instances b. Backfill with spot

Doing this dance well is hard, but it allows us to substitute an expensive resource with a cheap one with little degradation in experience. Cost savings here can be ~ 50%. Learn more in our blog post.

4 – ARM

ARM chips are fast and cheap, often resulting in 20-30% cost reductions. Coiled makes it easy to switch between instance types and try ARM on your existing workloads. Coiled also makes it easy to substitute out x86 Python packages for ARM packages without any effort in the common case.

This effort is new and exciting. We encourage users engaged with ARM chips to try it out!

5 – Transparency

Coiled tracks everything that is going on and gives aggregated views to cost owners in an organization. This allows organizations to make informed decisions and focus their energies appropriately.

Transparency takes a few forms. You can view the cost of currently running clusters, as well as the cost of historical computations both started by yourself, and by others in your team. You can also dive in to individual clusters and see what operations and even what lines of code accounted for the largest part of the cost.

You can also view aggregated cost information by user and by time to help identify opportunities for future improvement and cost assignment. See our documentation on managing costs for more details.

Account administrators can see a detailed breakdown of Coiled usage by a number of dimensions like user, GitHub reference, or tag.

6 – Limits and bounds

We also make it easy to bound the amount that you spend with us both through sensible tagging of cloud resources, and through our own limiting controls.

Cloud limits

Each resource that Coiled creates will be tagged with **owner: coiled, allowing you to see all the resources **created by your account. Alternatively, you can set custom tags for each cluster created so you have better visibility into cost breakdown. Refer to the section about tags in the documentation.

This allows you to more easily set alerts in your cloud provider account.

When you set alerts, AWS or GCP will send you an email notification when you have reached a certain threshold, so you can know how much you have spent so far. You can follow the official documentation on how to set up billing alerts:

Coiled usage limits

Additionally Coiled can help you bound how much you or your teammates spend with us. Coiled offers the following bounds:

A monthly total spending limit
Instantaneous spending limit on the full account (set by Coiled)
Instantaneous spending limits on each user (set by account admins)

These allow users to bound their overall spend, as well as enforce bounds on individual teammates. See our documentation to learn more.

7 – Risk and support

Coiled is maintained by a professional team and monitored 24x7. The risk associated with such a platform is lower than a DIY approach. This team understands cloud infrastructure, Dask, and the underlying computational libraries well enough to provide a more consistent and predictable experience for organizations that rely on this technology stack for business-critical workloads.

8 – Sensible configuration

There are many other small configuration choices that we make in our cloud architecture that each have non-trivial effects. These compound to significant savings. We’ll list a few below:

Good default instance choices, particularly those in the t3 line, which support cheap burstable compute
Avoiding NAT gateways which are common causes of surprising data transfer bills.
Multi-Zone worker pools, allowing users to get the best availability possible for Spot or rare instance types, even when supply is low
…

These choices are technical and in the weeds, but can each have a strongly positive impact on overall costs.

9 – Opportunity costs and time to delivery

You’ve opted to use the cloud because you value moving quickly.

The median Coiled user spends low thousands of dollars per month. While real money, this isn’t that big of a deal related to the value of moving faster. Coiled empowers teams to shave several months off of their timelines. When compared to the costs of monthly burn and waiting to build something internally, the costs here are often trivial in relation.

Conclusion

Coiled is free to try, with $25 of usage included. Coiled provides a cost-constrained experience that also delights data science and data engineering users, making them more productive with their time.

Convert Large JSON to Parquet with Dask

2022-01-01T00:00:00+00:00

tl;dr

You can use Coiled, the cloud-based Dask platform, to easily convert large JSON data into a tabular DataFrame stored as Parquet in a cloud object-store. Start off by iterating with Dask locally first to build and test your pipeline, then transfer the same workflow to Coiled with minimal code changes. We demonstrate a JSON to Parquet conversion for a 75GB dataset that runs without downloading the dataset to your local machine.

Why convert JSON to Parquet

Data scraped from the web in nested JSON format often needs to be converted into a tabular format for exploratory data analysis (EDA) and/or machine learning (ML). The Parquet file format is an optimal method for storing tabular data, allowing operations like column pruning and predicate pushdown filtering which greatly increases the performance of your workflows .

This post demonstrates a JSON to Parquet pipeline for a 75GB dataset, using Dask and Coiled to convert and store the data to a cloud object store. This pipeline bypasses the need for the dataset to be stored locally on your machine.

Upon completing this notebook, you will be able to:

Build and test your ETL workflow locally first, using a single test file representing 1 hour of Github activity data
Scale that same workflow out to the cloud using Coiled to process the entire dataset.

Spoiler alert – you’ll be running the exact same code in both cases, just changing the place where the computations are run.

You can find a full-code example in this notebook. To run the notebook locally, build a conda environment with the environment.yml file located in the notebook repo.

JSON to Parquet: Build your Pipeline Locally

It’s good practice to begin by building your pipeline locally first. The notebook linked above walks you through this process step-by-step. We’ll just summarize the steps here.

Extract a single file

We’ll be working with data from the Github Archive project for the year 2015. This dataset logs all public activity on Github and takes up ~75GB in uncompressed form.

Begin by extracting a single file from the Github Archive. This represents 1 hour of data and takes up ~5MB of data. There’s no need to work with any kind of parallel or cloud computing here, so you can iterate locally for now.

wget https://data.gharchive.org/2015-01-01-15.json.gz

Only scale out to the cloud if and when necessary to avoid unnecessary costs and code complexity.

Transform JSON data into DataFrame.

Great, you’ve extracted the data from its source. Now you can proceed to transform it into tabular DataFrame format.** **There are several different schemas overlapping in the data, which means you can’t simply cast this into a pandas or Dask DataFrame. Instead, you could filter out one subset, such as PushEvents, and work with that.

records.filter(lambda record: record["type"] == "PushEvent").take(1)

You can apply the process function (defined in the notebook) to flatten the nested JSON data into tabular format, with each row now representing a single Github commit.

records.filter(lambda record: record["type"] == "PushEvent").take(1)flattened = records.filter(lambda record: record["type"] == "PushEvent").map(process).flatten()

Then cast this data into a DataFrame using the .to_dataframe() method.

df = flattened.to_dataframe()

Load data to a local directory in Parquet file format

You’re now all set to write your DataFrame to a local directory as a .parquet file using the Dask DataFrame .to_parquet() method.

df.to_parquet(
    "test.parq", 
    engine="pyarrow",
    compression="snappy"
)

JSON to Parquet: Scaling out with Dask Clusters on Coiled

Great job building and testing out your workflow locally! Now let’s build a workflow that will collect the data for a full year, process it, and save it to the cloud object storage.

Spin up your Cluster

We’ll start by launching a Dask cluster in the cloud that can run your pipeline on the entire dataset. To run the code in this section, you’ll need a free Coiled account by logging into Coiled Cloud. You’ll only need to provide your Github credentials to create an account.

You will then need to make a software environment with the correct libraries so that the workers in your cluster are able to execute our computations.

import coiled

# create Coiled software environment
coiled.create_software_environment(
    name="github-parquet",
    conda=["dask", "pyarrow", "s3fs", "ujson", "requests", "lz4", "fastparquet"],
)

You can also create Coiled software environments using Docker images, environment.yml (conda), or requirements.txt (pip) files. For more information, check out the Coiled Docs.

Now, let’s spin up your Coiled cluster, specifying the cluster name, the software environment it’s running, and the number of Dask workers.

# spin up a Coiled cluster
cluster = coiled.Cluster(
    name="github-parquet",
    software="coiled-examples/github-parquet", 
    n_workers=10, 
)

Finally, point Dask to run on computations on your Coiled cluster.

# connect Dask to your Coiled cluster
from dask.distributed import Client
client = Client(cluster)
client

Run your pipeline on the Github Archive dataset

The moment we’ve all been waiting for! Your cluster is up and running, which means you’re all set to run the JSON to Parquet pipeline you built above on the entire dataset.

This requires 2 subtle changes to your code:

download all of the Github Archive files instead of just 1 test file
point df.to_parquet() to write to your s3 bucket instead of locally

Note that the code below uses a filenames list that contains all of the files for the year 2015 and the process function mentioned above. Refer to the notebook for the definitions of these two objects.

%%time
# read in json data
records = db.read_text(filenames).map(ujson.loads)

# filter out PushEvents
push = records.filter(lambda record: record["type"] == "PushEvent")

# process into tabular format, each row is a single commit
processed = push.map(process)

# flatten and cast to dataframe
df = processed.flatten().to_dataframe()

# write to parquet
df.to_parquet(
    's3://coiled-datasets/etl/test.parq',
    engine='pyarrow',
    compression='snappy'
)


CPU times: user 15.1 s, sys: 1.74 s, total: 16.8 s
Wall time: 19min 17s

Excellent, that works. But let’s see if we can speed it up a little…

Let’s scale our cluster up to boost performance. We’ll use the cluster.scale() command to double the number of workers in our cluster. We’ll also include a call to client.wait_for_workers() which will block activity until all of the workers are online. This way we can be sure that we’re throwing all the muscle we have at our computation.

# double n_workers
cluster.scale(20)

# this blocks activity until the specified number of workers have joined the cluster
client.wait_for_workers(20)

Let’s now re-run the same ETL pipeline on our scaled cluster.

%%time
# re-run etl pipeline
records = db.read_text(filenames).map(ujson.loads)
push = records.filter(lambda record: record["type"] == "PushEvent")
processed = push.map(process)
df = processed.flatten().to_dataframe()
df.to_parquet(
    's3://coiled-datasets/etl/test.parq',
    engine='pyarrow',
    compression='snappy'
)


CPU times: user 11.4 s, sys: 1.1 s, total: 12.5 s
Wall time: 9min 53s

We’ve cut the runtime in half, great job!

Converting Large JSON to Parquet Summary

In this notebook, we showed how to convert JSON to Parquet by converting a raw JSON data into a flattened DataFrame and stored it in the efficient Parquet file format on a cloud object-store. We performed this workflow first on a single test file locally. We then scaled the same workflow out to run on the cloud using Dask clusters on Coiled in order to process the entire 75GB dataset.

Main takeaways:

Coiled allows you to scale common ETL workflows to larger-than-memory datasets.
Only scale to the cloud if and when you need to. Cloud computing comes with its own set of challenges and overhead. So be strategic about deciding if and when to import Coiled and spin up a cluster.
Scale up your cluster for increased performance. We cut the runtime of the ETL function in half by scaling our cluster from 10 to 20 workers.

If you have any questions or suggestions for future material, feel free to drop us a line at support@coiled.io or in our Coiled Community Slack channel. We’d love to hear from you!

Coiled, one year in

2022-01-01T00:00:00+00:00

Coiled, a Dask company, is about one year old. We’ll have a more official celebration in mid-February (official date of incorporation), but I wanted to take this opportunity to talk a little bit about the journey over the last year, where that has placed us today, and what I think comes next.

History: building a company from an existing community

January last year I (matt) left NVIDIA to start Coiled. This was an interesting and challenging process in which I received a lot of help. This section chronicles that journey in three phases, and should be a good peek into what starting a company looks like for anyone interested in the process.

January-February — Company size: 1

I spent most of this period meeting with many different kinds of people

Other entrepreneurs who knew what to expect
Lawyers to help incorporate the company
Potential anchor clients to help us bootstrap
Potential first-hire employees to both help build out the company, and then run parts of it
Investors to explore if we wanted to take venture or angel funding

I found that I had to make lots of high impact decisions on very little information. This was an uncomfortable process, but in the end, things mostly worked out. After you talk to a few people in any category (like lawyers, potential employees, clients, investors) you get a sense for what you’re looking for, and can usually make a sensible decision.

This period ended with the Dask Developer Summit, a great reminder that I wasn’t alone in this process, and then the beginning of COVID-19, a less-than-great reminder that business success relies on adapting to changing circumstances.

Entrepreneurs

Special thanks to folks like Travis Oliphant/Peter Wang (Anaconda), Wes McKinney (Ursa), Andrew Montalenti (Parse.ly), Jeremiah Lowin (Prefect), Rodrigo and Felipe Aramburu (BlazingSQL) and many others who helped out both logistically and emotionally.

Employees

Most of Coiled’s current leadership team was assembled during this period (Rami, Hugo, James) while others were picked up later on in the process (Scott).

There were also some mis-steps. Very early on there was an engineer who built out a product prototype on his own time in order to see if we would work well together. This ended up not working out, which caused mutual stress. There were also early hires that ended up not being a good fit. We found and used contractors for various functions, some of which worked well and some which didn’t. In general, I learned that expecting and dealing with failure is probably more important than learning to avoid it.

Investors

My original plan was to bootstrap. I personally know many large companies that use Dask today that were happy to be early clients.

However, it quickly became clear that VCs valued this nascent company at levels that, to me, seemed silly to pass up at the time. We ended up taking $5,000,000 in seed funding, giving up 20% of the company while retaining all control (the only thing I couldn’t do was pay myself an unreasonable salary). I ended up taking this deal, which in hindsight was great. It allowed us to hire people with a wide range of risk tolerances, and allowed for a sensible work-life balance.

March-August — Company Size: 6

So we hired some engineers, a marketing guru, and went to work. This was a very small team relative to the funding that we had, but it worked well. An engineering team of four can be incredibly efficient. We had a prototype up in June, and launched in the beginning of September.

At the same time there was a lot of client negotiation and brand building.

Client negotiation

We had good momentum in February with a few large financial clients in February. In March all momentum was put on hold due to the COVID-19 outbreak. Eventually COVID became a boon to technology companies, but back in March everyone was concerned about a financial crisis, and so spending in many organizations froze.

By the summer things had thawed, and I was spending a good deal of time in business and legal negotiations. I have some experience working on the business side of deals, but almost no experience navigating the legal and internal corporate politics of enterprise sales. At this point we found Scott, who joined us as a sales advisor. This made a world of difference.

Renewed focus on smaller teams

At the same time we were also talking to some much smaller groups using Dask and wanting to use Coiled. These conversations were much more similar, and required an order of magnitude less work. In general the person that we found ourselves talking to was …

A team leader of a data science / machine learning team of 2-10
6-12 months in to using Dask
Had successfully used it on a big machine
Had then stood up some cloud infrastructure, usually with Kubernetes to scale out
Realized quickly that they definitely didn’t want to do this. This was taking 25-50% of an FTE, and wasn’t even that good of a solution.
Coiled’s lightweight platform did everything that they wanted, and was exactly the user experience that they were looking for
They’re happy to double their current Dask spending, typically around $3000-5000 / month in order to remove the DevOps pain, and hopefully become more efficient cost wise as well

These deals were not nearly as lucrative as the larger enterprise deals, but they happened quickly, and resulted in us improving our product much more quickly. Coiled’s engineering team transitioned from “doing what Matt thought was a good idea” to “doing what several customers were all telling us in unison”.

Spreading the Word

During this time, we also wanted to make sure that people could find us. To this end, we built out the MVP of our Marketing and Evangelism engines. This included our website, mailing list, a consistent blog with a consistent cadence, and a variety of social channels (Twitter, LinkedIn, YouTube, Medium). We also collaborated with broader PyData activities by appearing on podcasts, speaking at conferences, co-marketing with partner organizations, and presenting at Meetups. These activities then lead to tight feedback loops between Product, Marketing, and Sales, which helped us when we launched in late September.

September-December — Company Size: 15

Things were clearly working at this stage. We had signed a couple of the large enterprise contracts, had a bit less than $1M coming in the door, and had about 1000 signups on the public cloud service (only about 5-10% of which do meaningful work).

We were also slammed.

Everyone was busy, and so we decided to grow. Hugo/Marketing hired a couple of people to take the day-to-day burden off of him. Scott/Sales hired our first inside sales rep. Engineering went from four people to ten.

A lot of this period was about going from a small nimble team where everyone knew what was going on to building out a machine. Today if you sign up for and spin up a cluster on Coiled there are now automatic signals that go to Marketing and Sales. You’ll get a ping from Oren (our inside sales rep) pretty soon after asking if you’d like to talk. Some fraction of these turn into meetings with Scott/Sales, and some fraction of those turn into a commercial partnership going forward.

It’s surprising how consistent/predictable this flow is. As a Dask maintainer it has also been thrilling to see how many people out there use this software, and how many of them we can help get over some of the common early humps of devops, and Dask support.

At the beginning of this period we were very busy as individuals. Now that we have a machine flowing we’re very busy as a company. The amount of pressure is just as high, but we’re more able to scale to adapt to it.

Where Coiled is today

Today Coiled is focused around cloud.coiled.io, a hosted Dask-as-a-service product that makes it easy for anyone to be successful scaling their Python data science / machine learning workloads.

Most of our feedback comes from individual users
Most of our customers are small teams
Most of our revenue comes from a few large enterprise contracts.

This is great, we’ve developed a product that large companies want to pay for that benefits individuals as well. We get to be a successful business and help the world at the same time.

We’ve also found that many companies benefit from a little bit of advice. Dask is designed to be easy to use, but running anything at scale generally benefits from access to expertise. Coiled generally doesn’t take on services contracts (although we partner with and highly recommend Quansight for this work). However, we have found that we’re a good fit for both moderate annual contracts ($250k+) and small Q&A only monthly contracts ($2-5k) in order to make Dask perform optimally for companies. It turns out that Dask is much more powerful if you have access to some of the core maintainers.

The future of Coiled

We started the company conservatively, hiring only a few engineers and a skeleton support staff. This has allowed us to safely test a few hypotheses and learn a few lessons:

A lot of small teams have started using Dask in the last 6-12 months
Companies are quite happy to pay for managed Dask-as-a-service
We knew a lot more about what we needed to build than we realized

Because of this, we have high confidence that we’re on a good path, and so we’re more comfortable spending from our war chest (thank you investors). We’re hiring, and planning to expand significantly over the next few months in a few different directions.

Product

Our Coiled-on-AWS product is, today, the easiest and fastest way to get started with Dask on the Cloud. Universally users and customers say that this is exactly the user experience that they want.

However, they also say that they want many more things, including Azure/GCP support, integration with other products, a smoother GPU experience, and more. Coiled also falls over from time to time when under heavy load. These are all great features, and we’re working on them, but it’s also a lot for the current team to handle on our own. We’re stretched thin and need help so come join us!

Enterprise

Now that we’re within some larger companies there are a whole new set of constraints to deal with. We have this covered with current staff, but the deals in the sales pipeline loom large over us, and we’re likely to be quickly overwhelmed.

Partnerships

Coiled, like Dask, is designed to solve a specific piece of a larger problem. Dask adds scalability, but relies on other projects like Pandas/Scikit-Learn for computation and UI. Coiled adds scalability+DevOps but relies on other products like Prefect/Dash/Sagemaker to provide the other components of a data platform.

Many product companies at about our size see the potential here, and are looking to add a scale button to their product. We’re excited about these opportunities in order to see how effectively we can weave together a collection of similarly aligned companies.

Dask

Originally Coiled didn’t devote substantial resources to Dask itself other than having existing maintainer/employees organize and maintain on the side, and as part of paid work with some of the larger clients (thank you NVIDIA and DE Shaw).

However, there is still a lot going on in Dask right now. We need to step up on maintenance, and also drive forward some of the long running design changes that we’ve wanted for years. We’re happy to say that we have about one maintainer joining the company per month for the next few months. We’re excited about the added bandwidth that this will give the project.

Final Thoughts

We’ve built a machine that both generates revenue and serves the open-source community. Commercial open source businesses are hard to get right, and while we’re not there yet, I think that we’re headed in a good direction.

In a recent Talk Python to Me podcast with other startup founders, we talked a little bit about a sailing analogy to commercial open source companies.

Working on your own in OSS is like rowing a boat. You have full control on where you go, but you can only go so fast, and you can’t stop rowing
Building a business around OSS introduces wind. You go much faster, don’t have to row all the time, but you lose control about exactly where you end up
Our goal is to build directional technology like rudders and jibs, which allow us to harness wind (commercial interest) and direct it to where we want to go (saving the world)

Coiled today looks like a little skiff, with a jury-rigged rudder and a plastic sheet as a front sail. It works! But it could also use some love. We’re in calm waters and winds are strong, and so it feels like we could have a large impact.

Coiled today has grown a lot from me sitting alone in an AirBnB in Oakland, CA. We’re about twenty people strong today. However, we’re also still very young and need scrappy individuals to jump on board and help us build stuff ranging from product, to Dask itself, to messaging, to sales infrastructure and operations. If you’re looking for a new challenge with the potential for high impact then please take a look at our current jobs postings, or if nothing there looks right, but you think that you yourself would be a good fit then e-mail us at jobs.coiled.io. At our stage companies are more about high-impact people than specific roles.

‍

Coiled Cloud Architecture

2022-01-01T00:00:00+00:00

Running Dask in the cloud is easy.

Running Dask in the cloud securely is hard.

Running Dask for an enterprise in the cloud is really hard.

Over the last couple of years, Coiled has made a cloud-SaaS application that runs Dask for folks smoothly and securely in the cloud.

We thought you would like to hear a bit about the choices we made and why.

Design objectives

Coiled had three main objectives when making our design:

Easy for users
Trustworthy for IT
Easy for us to maintain

This led us down a path of a cloud SaaS application with:

A centralized control plane to make our development faster
It’s ok for Coiled to see users’ metadata
Deployment of Dask clusters within user cloud accounts
It’s not ok for Coiled to see users’ data
As few opinions as possible, following a “best of breed” philosophy

We felt that this was the best choice to enable scalable Python for the most people in a way that we could execute on with a small team. We’re pretty happy with the result.

Paths we didn’t choose

There are lots of excellent paths here that make sense, but that we chose not to pursue. For example:

Host computation on our servers, like Snowflake does

We felt that most customers wouldn’t be comfortable giving a small company like us their data directly.

Deploy entirely on customer accounts, like Cloudera does

We felt that this would make larger customers more comfortable with us, but would slow down our development cycle too much.

Build an end-to-end data science/engineering work environment, like Domino/Saturn Cloud/IBM/Databricks do

We had built these kinds of systems pre-Coiled and found that they’re just really really hard to do well. Data science environments were, in our opinion, too varied to give people a good experience.

These are all great choices, and target different kinds of users.

Coiled architecture

Coiled is a web application that manages cloud resources within user accounts on their behalf. You give us highly restricted access to your cloud. Members of your team then ask for Dask clusters, and we set up those Dask clusters for them on the fly.

We’ve found that this approach provides a good balance of giving users the data security that they need while letting us iterate very quickly on the product and so provide a better user experience. It’s a little nuanced though, let’s dig in a bit.

Admin users give Coiled permissions to manage cloud resources

First, let’s go through how we coordinate cloud access safely to set up resources for users.

As a user you need to to give Coiled enough permissions to do things like:

Create and destroy instances
Create and destroy network infrastructure
Write and read logs into cloud services
Write and read images into Docker registries

This is done once and then the Coiled web app has these permissions until you revoke them. Notably, you don’t give any permission to read and write from sensitive data sources like cloud object stores, Snowflake, and so on. All Coiled can really do is run up your cloud bill.

Users ask for Dask clusters

Then when users ask for a Dask cluster

import coiled
cluster = coiled.Cluster(...)

Coiled looks up your credentials, creates the necessary infrastructure, creates trusted security tokens, and gives those tokens to either side so that the user can then connect directly from their system to that cluster.

from dask.distributed import Client
client = Client(cluster)

Coiled sets up infrastructure, brokers a secure connection between the user and that infrastructure, and then gets out of the way as quickly as possible.

Users send data access credentials to the cluster

The user then has access to a Dask cluster without Coiled acting in the middle. This is great. The user then sends data access credentials from their machine up to the Dask cluster. This allows the cluster to access sensitive data resources within the company without Coiled itself ever seeing the data or the credentials.

Results

The result is that anyone on your team can quickly get access to a fully authenticated and secure Dask cluster that can access all of the data that they have access to in a minute or two.

Let’s talk a bit about how we actually set up and deploy resources.

Set up Dask workers in the cloud

To actually create the Dask scheduler and worker processes Coiled uses raw instances. This is in contrast to other options like using Kubernetes or container-as-a-service options like AWS Fargate.

Kubernetes is great. We love Kubernetes here at Coiled, but it also introduces complexity that we actually don’t need to accomplish our goals. Also, removing the intermediate layer of Kubernetes allowed us to develop a product that was more flexible and raw-metal.

Instead, we use technologies like EC2 fleets so that we could easily adapt to our customers’ many varied workloads with often very different hardware requirements. Once we have a bunch of VMs at our disposal we need to install just the right versions of Python packages on them. This can be tricky.

Manage python libraries in the cloud

In order for Dask to work well, the versions of software libraries need to match exactly between where Coiled/Dask is invoked, and the workers that are run in the cloud. We do this in two ways (learn more in our documentation):

We enable users to ask for specific conda/pip packages, and build docker images to their specification.

import coiled
coiled.create_software_environment(
    "prefect-flow",
    conda={"channels": ["conda-forge"],
           "dependencies": ["coiled-runtime", "prefect"]
)

This works well, but can get a bit finicky especially when the user’s environment changes rapidly. Adding a conda solve and docker build step into development cycles causes a lot of slowdown and pain.

We dynamically query the user’s local environment and install those packages on-the-fly This is far more magical, but works pretty well:

cluster = coiled.Cluster(package_sync=True)

Handle Dask logs in the cloud

When something breaks users often want to go and look at logs. Where should these logs go?

As always, Coiled tries hard not to store anything from the user, and instead rely on cloud-native solutions. This means that for AWS we store user logs in CloudWatch, and for Google Cloud we store user logs in Google Cloud Logging. We provide easy methods to query these logs, but they are also easily accessible from however the user already handles logs within their organization.

As always, Coiled tries to be convenient, but not intrusive.

24x7 tracking of cluster status and Dask usage

Finally, Coiled also acts as a watchdog over your Dask processes. If a cluster has been idle for too long we’ll kill it. If a user is launching GPU resources and not using them, we make folks aware of it. Gathering usage information is key.

To accomplish this, Coiled installs a SchedulerPlugin in the Dask schedulers that tracks various usage metrics like idleness, which operations are running, etc. This information is really useful for debugging and cost optimization, but can also be somewhat invasive. We allow users to turn various bits of this data tracking on and off with configuration (see the Coiled documentation on analytics).

Billing

Coiled is free for the first 500 CPU hours a month (although you still have to pay your cloud provider). Beyond that we charge $0.05 per managed CPU hour, with rates that decline based on usage.

We’ve found that this is enough for individuals to operate entirely for free, and for folks at larger organizations to “try before they buy” pretty effectively. For companies who do spend thousands of dollars or more per month, we look comparable to services like Databricks or SageMaker. Especially when compared against the personnel cost of building and maintaining a system like Coiled, this very obviously becomes worth the cost.

Results

Anyone with a cloud account can connect their accounts to Coiled in a couple of minutes, and can be scaling happily afterwards, knowing that they’re secure. Our development teams have enough information to see what’s going on to rapidly iterate on a better user experience, while users can rest assured that both their data and budgets are safe and secure. For a more details, see Why Coiled?.

If you’re interested, Coiled is easy to try out. Just run the following in any Python-enabled environment after you’ve created a Coiled account:

pip install coiled
coiled login
coiled setup aws
ipython

import coiled
cluster = coiled.Cluster(n_workers=100, package_sync=True)

from dask.distributed import Client
client = Client(cluster)

import dask.array as da
da.random.random((10000, 10000)).sum().compute()

Code Formatting Jupyter Notebooks with Black

2022-01-01T00:00:00+00:00

Black is an amazing Python code formatting tool that automatically corrects your code.

The jupyterlab_code_formatter extension makes it easy for you to run black in a Jupyter Notebook with the click of a button.

This post explains how to use the plugin and why consistent code formatting is so important.

Installation

You can add the necessary libraries to your conda environment with this command:

conda install -c conda-forge jupyterlab_code_formatter black isort

You can also add the three required dependencies to a YAML file if you’re using a multi-environment conda workflow.

Usage

Restart your Jupyter Lab session, and you’ll now see this formatting button:

Click the button to format all the cells in your notebook with Black’s code formatting rules.

This notebook is running locally, but you can use Coiled to run Jupyter Notebooks on the cloud.

Example code formatting

Let’s take a look at a chunk of code that’s not formatted and then see how it looks after it’s been formatted. Here’s the unformatted code.

ddf = dd.read_parquet(
   "s3://coiled-datasets/timeseries/20-years/parquet",
   storage_options={"anon": True, 'use_ssl': True}
)

This code uses inconsistent apostrophe symbols, and the last argument to read_parquet doesn’t have a trailing comma.

Click the code formatting button, and this cell is automatically reformatted as follows:

ddf = dd.read_parquet(
   "s3://coiled-datasets/timeseries/20-years/parquet",
   storage_options={"anon": True, "use_ssl": True},
)

You might feel like code formatting is minor, but it’s actually really important.

Why code formatting is important

Before automatic code formatting was popular, teams would create style guides and manually enforce coding rules.

This was tedious and sometimes contentious. Some programmers are more passionate about their whitespace preferences than you might think!

Spending time creating style guides, commenting about formatting in pull requests, and arguing about double or single quotes is not a good use of developer time. It’s better to follow a community-accepted, fully automated solution.

Other code formatting options

jupyterlab_code_formatter also supports YAPF and Autopep8 code formatting. We recommend using Black because it’s popular and used in several popular open-source libraries.

Conclusion

JupyterLab has an extension that makes it easy to automatically format the Python code in your notebooks.

Automatic code formatters make it easy for teams to consistently style their code without extra work.

Cloud ETL automation with Github Actions and Coiled

2022-01-01T00:00:00+00:00

Github Actions let you launch automated jobs from your Github repository. Coiled lets you run your Python code in the cloud. Combining the two gives you lightweight workflow orchestration for heavy ETL (extract-transform-load) jobs that can run in the cloud without any complicated infrastructure provisioning or DevOps.

This Github Action runs an ETL job on 15 GB of data and writes the results to an Amazon S3 bucket. Running this with a Github Action alone would not be possible as the default virtual machines on which Github Actions are run only have 7GB of memory. Running the action on a Coiled cluster lets you crunch the entire dataset. Since the Coiled cluster runs on AWS, the data can easily be written to an S3 bucket.

This post will demonstrate how running Github Actions on Coiled can be a useful way to schedule automated data-processing jobs. We’ll start with some background on Github Actions and Coiled and then walk through an example script in detail.

Github Actions, S3 and Coiled: Workflow Orchestration

Many data analysts have tasks that need to be run on a fixed schedule. For example, you might want to fetch new data every day at 8AM, preprocess it, and use it to train or update a machine learning model. Workflow orchestration tools allow you to schedule jobs like this without human intervention. Popular orchestration tools include Apache Airflow, Prefect, Argo and others.

Managing these workflow orchestration tools can be more work than you want to do. Github Actions is a lightweight and cheap (or even free) alternative that lets you run automated jobs based on predefined trigger events. Since many developers and data analysts already have their code in Github repositories, this can be an ideal way to run automated jobs without having to learn another tool or manage a full-fledged orchestration deployment.

Github Actions, S3 and Coiled: Scale to the Cloud

Github Actions are run on virtual machines that are limited in the hardware resources they let you access. This makes sense because Github Actions are free (entirely for public repositories, up to a limit for private ones) and Github probably doesn’t want you running massive workloads on their free service. By default, Github Actions are hosted on ‘runners’ with 7GB of RAM.

So what do you do if you want to run a scheduled ETL job with Github Actions that requires more memory resources than the Github Action default?

If you want to use Github Actions for more compute-intensive or memory-bound tasks, you’ll need to deploy the task to another virtual machine or cluster of machines. There are ways to do this yourself with Amazon ECS, Azure or Google Kubernetes. But deploying Github Actions to the cloud yourself is complex and time-intensive. Github also offers an option to use self-hosted runners for Github Actions. These also require an extensive installation, setup and maintenance protocol.

Coiled makes it easy to deploy Python code to the cloud. Coiled handles all the infrastructure setup for you, including software environments, security and cost-management. Coiled Clusters are launched on-demand and can be stopped immediately once the computation is completed, so there will be no unexpected costs. Even better: Coiled currently provides a generous free tier, so you’ll only be paying your own cloud account fees.

Github Actions, S3 and Coiled: In Action

We’ve put together a Github Action that you can customize to run your own automated ETL jobs. Let’s walk through the script in detail. If you’re totally new to Github Actions, we recommend following the official Github tutorial first.

Let’s break down the workflow that our Github Action performs step by step.

Defining your Github Action

We’ll set this ETL job to run on a schedule: it will be triggered at 12pm (noon) every day.

on:
 schedule:
   # run at noon every day
   - cron: "0 12 * * *"

Github Actions uses the cron format to specify dates and time. You can also have your Github Actions triggered by other events, such as push events to specific repositories.

We’ll then define our quickstart job:

jobs:
 quickstart:
   runs-on: "ubuntu-latest"
   strategy:
     fail-fast: false

And define the steps this job will consist of:

steps:
  # Needed to run on PR or branches
  - name: Checkout source
    uses: actions/checkout@v2
        
  - name: Set up local environment
    uses: conda-incubator/setup-miniconda@v2
    with:
      miniforge-variant: Mambaforge
      use-mamba: true
      environment-file: ci/environment-3.9.yaml
  
  - name: Run quickstart
    env:
      DASK_COILED__TOKEN: ${{ secrets.COILED_TOKEN }}
      AWS_ACCESS_KEY_ID: ${{ secrets.BLOG_AWS_ACCESS_KEY_ID }} 
      AWS_SECRET_ACCESS_KEY: ${{ secrets.BLOG_AWS_SECRET_ACCESS_KEY }}
    run: python scripts/quickstart.py

‍

The first and second steps are standard setup protocols. The third step is the one that will actually perform our ETL job: it will run the Python script quickstart.py located in the scripts directory.

Authenticate with S3

Notice that we passed four environment variables to this script:

DASK_COILED__TOKEN contains the token your script will need to authenticate with Coiled and launch cloud computing resources.
AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY contain your AWS credentials that let you authenticate with S3.

Secrets are pulled from your Github account settings. Follow these steps if you do not know how to store secrets in your Github repository.

Running your ETL Job

Now let’s take a closer look at the quickstart.py script that our Github Action will run.

After importing dependencies, the script starts by launching a Coiled cluster:

GITHUB_RUN_ID = os.environ["GITHUB_RUN_ID"]


cluster = coiled.Cluster(
    name=f"github-actions-{GITHUB_RUN_ID}",
    n_workers=10,
    worker_memory="8Gib",
)

This spins up a cluster of 10 virtual machines with 8GB of RAM each. Coiled automatically replicates the local Python environment in your cluster. You can customize your cluster to specify different instance types and hardware specs. Each cluster that is launched is given a unique name using the Github Run ID.

We then connect the Python client running the script (in this case the virtual machine that’s running the Github Action) to our Coiled cluster:

client = cluster.get_client()

Now we’re all set to load in our large dataset. We’ll use Dask to do this. Dask is a parallel computing library for Python that makes it easy to scale your pandas and NumPy workloads to larger-than-memory datasets.

The code below loads the Parquet data and performs a groupby aggregation:

ddf = dd.read_parquet(
    "s3://coiled-datasets/nyc-tlc/2019",
    columns=["passenger_count", "tip_amount"],
    storage_options={"anon": True},
).persist()

# perform groupby aggregation
result = ddf.groupby("passenger_count").tip_amount.mean()

Write to S3

Chances are that once you have this result, you’d like to store it somewhere. Let’s write the result to an Amazon S3 bucket. We’ll write to a dedicated S3 bucket and create a unique directory for each Github Action run:

# write result to s3
bucket_path = (
    f"s3://coiled-github-actions-blog/github-action-{GITHUB_RUN_ID}/quickstart.parquet"
)
result.to_parquet(
    bucket_path,
)
print(f"The result was successfully written to {bucket_path}")

The cluster will shutdown automatically after 20 minutes, but we can turn it off with:

cluster.shutdown()

All set! Once pushed to the repository, the Github Action will run this script at noon every day. You can of course adjust the cron schedule to run this at a different interval. After a few runs, your S3 bucket UI will look something like this, with a unique directory for each completed run:

Why Coiled?

While there are other ways to run your Github Actions on cloud-computing resources (or with self-hosted runners), these all require complicated setup processes and careful maintenance to implement security protocols and avoid unnecessary costs. Coiled takes care of all of this DevOps for you so you can focus on writing and running your Python code.

Coiled is also very scriptable and is not picky about where you are running your Python code. Unlike some of the other cloud-computing services, Coiled does not limit you to launching clusters from hosted notebooks. This means you can launch Coiled clusters from anywhere you can run Python code, including Github Actions (see our integration guide for more details).

Github Actions, S3 and Coiled: Your turn

If you’re ready to build your own Github Actions that leverage cloud-computing resources, follow these steps:

Fork the template repository
Sign up
Edit the quickstart.py and quickstart.yaml files to suit your specific use case
Include your Coiled token and AWS S3 credentials in your Github settings
Launch your Github Action!

Blog

Coiled Stability Update

Database CPU Saturation

Application Memory Saturation

Silver Lining

Closing

MLflow on Coiled with Sidecars and Filestores

Running MLflow as a Sidecar

Persisting Data with Filestores

Training Code and Launching the Cluster

MLflow Is for Sharing

What Else Can You Do With Sidecars and Filestores?

How to run pandas on GPUs with Coiled

Why GPU-accelerated pandas delivers faster performance

How Coiled simplifies GPU access

Example: How often do New Yorker’s tip?

Sign up

Install Coiled

Create a GPU-enabled software environment

Launch a GPU-enabled Jupyter Notebook

Running GPU-accelerated cudf.pandas

Run a Python script on a cloud GPU

When to use GPU acceleration

FAQs

Building a scientific Python stack with Coiled

The challenge: Scaling beyond a workstation

Coiled: A platform for scalable Python computing on the cloud

Integration with the scientific Python stack

Example: Mineral exploration at scale

Observability and optimization

Cost control and cloud efficiency

Focus on science, not infrastructure

Example: Mapping tropical forests at planetary scale

Business and scientific impact

Conclusion

FAQs

What is scientific Python?

How is Python used in science?

Does NASA use Python?

Marimo on Coiled

Running in Docker on Your Local Computer

Running in the Cloud

Sync Files

Install Packages Automatically

Run marimo on a VM With a GPU

Wrapping Up

Is AWS Batch right for data professionals?

Navigating infrastructure setup and configuration

Understanding networking requirements

Managing complex IAM permissions

Debugging across multiple AWS services

Scaling and GPU considerations

Container-first approach vs. script-based workflows

DevOps integration requirements

When AWS Batch works well

Alternatives for data-focused teams

Batch jobs with Coiled

FAQs

What is AWS Batch used for?

Why might AWS Batch feel complex for data professionals?

Does AWS Batch support GPUs?

How do I submit scripts to AWS Batch?

What are common AWS Batch considerations for data teams?

What are alternatives to AWS Batch?

Choosing an AWS Batch alternative

What makes AWS Batch challenging?

AWS Batch with Fargate

AWS Batch with EC2 for GPU access

What to look for in an AWS Batch alternative?

Coiled as an AWS Batch alternative: An easy way to scale scripts on the cloud

Coiled Batch for single jobs to large-scale parallelism

Next steps

Long-running Lambda workloads: Challenges and alternatives

The challenge of long-running Lambda tasks

Better alternatives to AWS Lambda for long-running tasks

Fargate for serverless containers

Step Functions for workflow orchestration

AWS Batch for large-scale compute jobs

Using Coiled instead of AWS services for long-running Lambda workarounds

Coiled serverless Python functions

Running GPU-accelerated `cudf.pandas`