Running PySpark Jobs in Parallel within Dagster on Dockerized Setup

[MammadTavakoli]

I am implementing a Docker setup for running Dagster and PySpark together. My docker-compose.yml file looks like this:

dagster:
    container_name: dagster
    hostname: dagster   
    build:
      context: .
      dockerfile: ./docker/dagster/Dockerfile    
    restart: always
    environment:       
      DAGSTER_HOME: ./app/mediation   
    volumes:        
      - /var/run/docker.sock:/var/run/docker.sock
      - ./app/mediation:/app/mediation      
      - ./shared:/shared:rw     
    ports:        
        - 3000:3000     
    command: dagster-webserver -h 0.0.0.0   
    depends_on:
      - minio 
    networks:
      - dl-network

  dagster-daemon:
    container_name: dagster-daemon
    hostname: dagster-daemon  
    build:
      context: .
      dockerfile: ./docker/dagster/Dockerfile
    restart: always
    environment:
      DAGSTER_HOME: ./app/mediation   
    volumes:
      - /var/run/docker.sock:/var/run/docker.sock
      - ./app/mediation:/app/mediation
      - ./shared:/shared:rw      
    command: dagster-daemon run
    depends_on:
      - dagster
    networks:
      - dl-network

spark-master:
    hostname: spark-master
    container_name: spark-master
    image: bitnami/spark:3.5.4
    environment:
      SPARK_MODE: master     
    ports:
      - "4040:4040"
      - "8080:8080"
      - "7077:7077"      
   
    networks:
      - dl-network
# Spark Worker Services
  spark-worker-1:
    hostname: spark-worker-1
    container_name: spark-worker-1
    image: bitnami/spark:5.3.4
    depends_on:
      - spark-master
    environment:
      SPARK_MODE: worker
      SPARK_MASTER_URL: spark://spark-master:7077 
    ports:
      - "8081:8081"  
    networks:
      - dl-network

  spark-worker-2:
    hostname: spark-worker-2
    container_name: spark-worker-2
    image: bitnami/spark:5.3.4
    depends_on:
      - spark-master
    environment:
      SPARK_MODE: worker
      SPARK_MASTER_URL: spark://spark-master:7077
    ports:
      - 8082:8081"  
    networks:
      - dl-network

The Dagster Dockerfile is as follows:

FROM python:3.12-slim 

# Set the working directory
WORKDIR /app

# Copy only the requirements file first to leverage Docker cache
COPY /docker/dagster/requirements.txt .


# Update and install minimal dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \ 
    openjdk-17-jdk \
    build-essential \
    python3-dev \
    curl \    
    software-properties-common \
    git \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*

# Upgrade pip and install dependencies from requirements.txt
RUN pip install --upgrade pip --no-cache-dir && \
pip install -r requirements.txt --no-cache-dir
    
# Set JAVA_HOME environment variable
ENV JAVA_HOME=/usr/lib/jvm/java-17-openjdk-amd64
ENV PATH=$JAVA_HOME/bin:$PATH

# Copy the application source code
COPY /app/mediation /app/mediation

# Set the working directory
WORKDIR /app/mediation

# Install development dependencies
RUN  pip install -e ".[dev]"   --no-cache-dir -i https://pypi.tuna.tsinghua.edu.cn/simple

Problem Description

I want to process data for five countries in parallel. For each country:

1. Data is fetched from a database, partitioned by city.
2. Each country's cities are processed sequentially (but the countries should run in
   parallel).

Here is the Dagster code I wrote for this:

from dagster import asset, op, job, graph, AssetExecutionContext, EnvVar, graph_asset, AssetIn, Nothing, In
from pyspark.sql.functions import current_date, col, coalesce
from datetime import datetime, timedelta
import time
from concurrent.futures import ThreadPoolExecutor
from dagster.core.definitions import DynamicOutput, DynamicOutputDefinition

class shopProcessor:
    def __init__(self, context: AssetExecutionContext, check_hour: int=3):
        self.context = context
        self.spark = context.resources.spark.spark_session
        self.sql_server = context.resources.sql_server
        self.check_hour = check_hour
    
    def execute_query(self, query: str, params: dict = None):
        try:                        
            return self.spark.sql(query, params) if params else self.spark.sql(query)
        except Exception as e:
            self.context.log.error(f"Query execution failed: {query}, Error: {e}")
            raise

    def get_citys(self, country_id=None):
        where_clause = f"WHERE country_id = {country_id}" if country_id else "WHERE 1=1"
        query = f"""
            SELECT                 
                city_ID, 
                city_Code, 
                city_Name, 
                country_id
            FROM 
                Deltalake.shop_city
            {where_clause}          
        """
      
        
        params = (country_id,) if country_id is not None else (None,)             
        result_df = self.execute_query(query, params)        
        return result_df.collect()
    
    def etl_shop(self, city_id, city_code):

        # Define the dynamic SQL query with OPENQUERY to handle remote SQL Server connection
        query = f"""
                SELECT  CAST({int(city_id)} as int) AS city_ID, * 
                FROM shop{city_code} d WITH(NOLOCK)                      
        """
     
        result_df = self.spark.read \
            .format("jdbc") \
            .option("url", self.sql_server["url"]) \
            .option("driver", self.sql_server["properties"]["driver"]) \
            .option("dbtable", f"({query}) AS tmp") \
            .option("user", self.sql_server["properties"]["user"]) \
            .option("password", self.sql_server["properties"]["password"]) \
            .load()

        if result_df.isEmpty():            
            row_count = 0
        else:            
            result_df.write \
                .format("delta") \
                .mode("append") \
                .option("userMetadata", str(datetime.now())) \
                .saveAsTable("Deltalake.shop")
            row_count = result_df.count()

        self.context.log.info(f"Data written to Delta Lake for {city_id} at: Deltalake.shop ")

        return row_count

    def process_city_shop(self, city: dict):
        ui_url = self.spark.sparkContext.uiWebUrl        

        city_id   = city["city_ID"]
        city_code = city["city_Code"]

        row_count = self.etl_shop(city_id, city_code) 

        self.context.log.info(f"Finished processing for city: {city_name}")

    def process_citys(self, country_id: int):       
        citys = self.get_citys(country_id)   
        for city in citys:                            
            self.process_city_shop(city)

# Define shop_processing graph
def create_shop_asset(country_id: int):
    @asset(
        kinds={"deltalake", 'sqlserver', 'minio'},
        required_resource_keys={"spark", "sql_server"},
        group_name="shop_extract",       
        name=f"shop_raw_{country_id}" 
        )
    def _shop_processing(context: AssetExecutionContext):
        shop = shopProcessor(context)
        shop.process_citys(country_id)
    return _shop_processing

shop_raw_assets = [create_shop_asset(i) for i in range(1, 6)]

The shop_raw_assets array generates five assets (one per country) and I expected these assets to run in parallel. However:

1. Only one asset runs at a time, while others wait.
2. In Spark Master UI (accessible on 8080), I see five applications created, but only one application utilizes cores while others have 0 cores.

Environment

• Dockerized environment with Dagster, PySpark, and Spark Master/Workers.
• Dagster orchestrates the ETL jobs.
• PySpark processes the data.

Question

How can I modify my setup or code to ensure that the assets in shop_raw_assets run in parallel, utilizing the available Spark cores effectively?