diff --git a/README.md b/README.md
index a959185..fa865e1 100644
--- a/README.md
+++ b/README.md
@@ -90,14 +90,19 @@ Edit compose.yml file if you use repository folked from this repository.
 services:
   runner:
     privileged: true
+    restart: always
+    env_file:
+      - .env
+    environment:
+      REPOSITORY_URL: https://github.com/[OWNER]/[REPOSITORY]
+      RUNNER_NAME: self-hosted-runner
+      RUNNER_GROUP: default
+      TARGET_ARCH: x64
     build:
       context: .
       dockerfile: Dockerfile
       args:
-        REPOSITORY_URL: [Target Repository URL]
         GH_RUNNER_VERSION: 2.329.0
-        RUNNER_NAME: self-hosted-runner
-        RUNNER_GROUP: default
         TARGET_ARCH: x64
 ```
 
diff --git a/self-hosted-runner/Dockerfile b/self-hosted-runner/Dockerfile
index 1e69b5a..f46da23 100644
--- a/self-hosted-runner/Dockerfile
+++ b/self-hosted-runner/Dockerfile
@@ -1,13 +1,7 @@
 FROM ubuntu:24.04
 
-ARG REPOSITORY_URL="https://github.com/user/repo"
 ARG GH_RUNNER_VERSION="2.329.0"
-ARG RUNNER_NAME="self-hosted-github-actions-runner"
-ARG RUNNER_GROUP="default"
-ARG RUNNER_LABELS="self-hosted,Linux"
 ARG TARGET_ARCH="x64"
-ARG PERSONAL_ACCESS_TOKEN=""
-ARG REGISTRY_TOKEN=""
 
 ENV DEBIAN_FRONTEND=noninteractive \
     AGENT_TOOLSDIRECTORY=/opt/hostedtoolcache
@@ -65,27 +59,4 @@ WORKDIR /home/ubuntu
 RUN curl -fsSL -o actions-runner.tar.gz -L "https://github.com/actions/runner/releases/download/v${GH_RUNNER_VERSION}/actions-runner-linux-${TARGET_ARCH}-${GH_RUNNER_VERSION}.tar.gz" \
     && tar xf actions-runner.tar.gz \
     && rm actions-runner.tar.gz \
-    && sudo ./bin/installdependencies.sh \
-    && if [ -n "$PERSONAL_ACCESS_TOKEN" ]; then \
-    ./config.sh \
-    --unattended \
-    --url $REPOSITORY_URL \
-    --pat "$PERSONAL_ACCESS_TOKEN" \
-    --name $RUNNER_NAME \
-    --runnergroup $RUNNER_GROUP \
-    --labels "${RUNNER_LABELS},${TARGET_ARCH}" \
-    --work /home/ubuntu/actions-runner \
-    --replace; \
-    else \
-    ./config.sh \
-    --unattended \
-    --url $REPOSITORY_URL \
-    --token "$REGISTRY_TOKEN" \
-    --name $RUNNER_NAME \
-    --runnergroup $RUNNER_GROUP \
-    --labels "${RUNNER_LABELS},${TARGET_ARCH}" \
-    --work /home/ubuntu/actions-runner \
-    --replace; \
-    fi
-
-CMD ["./run.sh"]
+    && sudo ./bin/installdependencies.sh
diff --git a/self-hosted-runner/compose.yml b/self-hosted-runner/compose.yml
index bebcd1e..0e03fd4 100644
--- a/self-hosted-runner/compose.yml
+++ b/self-hosted-runner/compose.yml
@@ -4,12 +4,14 @@ services:
     restart: always
     env_file:
       - .env
+    environment:
+      REPOSITORY_URL: https://github.com/mjun0812/flash-attention-prebuild-wheels
+      RUNNER_NAME: self-hosted-runner
+      RUNNER_GROUP: default
+      TARGET_ARCH: x64
     build:
       context: .
       dockerfile: Dockerfile
       args:
-        REPOSITORY_URL: https://github.com/mjun0812/flash-attention-prebuild-wheels
         GH_RUNNER_VERSION: 2.329.0
-        RUNNER_NAME: self-hosted-runner
-        RUNNER_GROUP: default
         TARGET_ARCH: x64
diff --git a/self-hosted-runner/entrypoint.sh b/self-hosted-runner/entrypoint.sh
index d56b933..1a04ffe 100644
--- a/self-hosted-runner/entrypoint.sh
+++ b/self-hosted-runner/entrypoint.sh
@@ -1,7 +1,30 @@
 #!/bin/bash
-id
 
 # Start docker daemon
 sudo service docker start
 
-exec "$@"
\ No newline at end of file
+if [ -n "$PERSONAL_ACCESS_TOKEN" ]; then
+    echo "Using personal access token";
+    ./config.sh \
+        --unattended \
+        --url $REPOSITORY_URL \
+        --pat "$PERSONAL_ACCESS_TOKEN" \
+        --name $RUNNER_NAME \
+        --runnergroup $RUNNER_GROUP \
+        --labels "${RUNNER_LABELS},${TARGET_ARCH}" \
+        --work /home/ubuntu/actions-runner \
+        --replace;
+else
+    echo "Using registry token";
+    ./config.sh \
+        --unattended \
+        --url $REPOSITORY_URL \
+        --token "$REGISTRY_TOKEN" \
+        --name $RUNNER_NAME \
+        --runnergroup $RUNNER_GROUP \
+        --labels "${RUNNER_LABELS},${TARGET_ARCH}" \
+        --work /home/ubuntu/actions-runner \
+        --replace;
+fi
+
+exec "./run.sh"