diff --git a/src/load_generator/Dockerfile b/src/load_generator/Dockerfile
index 2e5427a34ab5f1c932d1aee378bc5c46ca70ef36..8f59bb4dbc47d6180d96187eea20cc6e8cbce60c 100644
--- a/src/load_generator/Dockerfile
+++ b/src/load_generator/Dockerfile
@@ -62,10 +62,16 @@ RUN python3 -m pip install -r requirements.txt
 
 # Add component files into working directory
 WORKDIR /var/teraflow
-COPY src/context/. context/
-COPY src/dlt/. dlt/
-COPY src/service/. service/
-COPY src/slice/. slice/
+COPY src/context/__init__.py context/__init__.py
+COPY src/context/client/. context/client/
+COPY src/dlt/__init__.py dlt/__init__.py
+COPY src/dlt/connector/__init__.py dlt/connector/__init__.py
+COPY src/dlt/connector/. dlt/connector/
+COPY src/load_generator/. load_generator/
+COPY src/service/__init__.py service/__init__.py
+COPY src/service/client/. service/client/
+COPY src/slice/__init__.py slice/__init__.py
+COPY src/slice/client/. slice/client/
 
 # Start the service
 ENTRYPOINT ["python", "-m", "load_generator.service"]