diff --git a/Dockerfile b/Dockerfile
index 879c0fe..e0747cb 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,5 +1,6 @@
 FROM mageai/mageai:latest
 
+ARG PROJECT_NAME
 ARG USER_CODE_PATH=/home/src/${PROJECT_NAME}
 
 # Note: this overwrites the requirements.txt file in your new project on first run. 
diff --git a/docker-compose.yml b/docker-compose.yml
index 832c208..a3135c5 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -1,4 +1,3 @@
-version: '3'
 services:
   magic:
     image: mageai/mageai:latest
@@ -8,6 +7,8 @@ services:
     build:
       context: .
       dockerfile: Dockerfile
+      args:
+        PROJECT_NAME: ${PROJECT_NAME}
     environment:
       USER_CODE_PATH: /home/src/${PROJECT_NAME}
       ENV: ${ENV}
@@ -15,4 +16,4 @@ services:
       - 6789:6789
     volumes:
       - .:/home/src/
-    restart: on-failure:5
+    restart: unless-stopped