Docker 環境の追加と Google Drive upload/download スクリプトの追加

.gitignore

@@ -0,0 +1,2 @@
+data_and_model/**
+secrets/**

Dockerfile

@@ -0,0 +1,26 @@
+FROM nvcr.io/nvidia/tensorrt:25.01-py3
+
+ENV PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1 \
+    PIP_DISABLE_PIP_VERSION_CHECK=1 \
+    PIP_NO_CACHE_DIR=1 \
+    DEBIAN_FRONTEND=noninteractive \
+    LC_ALL="C.UTF-8" \
+    LANG="C.UTF-8" \
+    TZ=Asia/Tokyo
+
+WORKDIR /workspace
+
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+    ca-certificates \
+    tzdata \
+    curl \
+    build-essential && \
+    rm -rf /var/lib/apt/lists/*
+
+COPY requirements.txt /app/requirements.txt
+RUN python3 -m pip install --upgrade pip && \
+    python3 -m pip install -r /app/requirements.txt
+
+CMD ["/bin/bash"]

README.md

@@ -1,2 +1,48 @@
 # psv-utils
 Utilities related to PackedSfenValue format
+
+# how to use with Docker
+
+Build docker image
+
+```bash
+./docker/docker_build.sh
+```
+
+Run container
+
+```bash
+./docker/docker_run.sh --name rescore --gpus all
+```
+
+Run rescore code in the container
+
+```bash
+python rescore_with_dlshogi.py <PATH_TO_INPUT> <PATH_TO_OUTPUT> --model-path <MODEL_PATH> --batch-size 4096 --score-scaling 1200.0 --blend-ratio 1.0 --enable-cuda --enable-tensorrt
+```
+
+# Download & Upload files between Google Drive
+
+After building a docker image, run container with specifying port
+
+```bash
+./docker/docker_run.sh --name download --port 8765
+```
+
+Download:
+
+```bash
+python3 google_drive_utils/download_from_drive.py --port 8765 --out-dir <OUTPUT_DIRECTORY> --file-id <GOOGLE_DRIVEs_FILE_ID>
+```
+
+Then, follow the instructions given in your terminal.
+(e.g. auth, login, etc.)
+
+Upload:
+
+```bash
+python3 google_drive_utils/upload_into_drive.py --port 8765 --file <YOUR_FILE_PATH> --folder-id <GOOGLE_DRIVE_FOLDER_ID>
+```
+
+Then, follow the instructions given in your terminal.
+(e.g. auth, login, etc.)

docker/docker_build.sh

@@ -0,0 +1 @@
+docker build -t psv_utils .

docker/docker_run.sh

@@ -0,0 +1,56 @@
+NAME_ARG=""
+PORT_ARG=""
+GPU_ARG=""
+
+while [[ $# -gt 0 ]]; do
+  case "$1" in 
+    --name|-n)
+      if [[ -z "$2" ]]; then
+        echo "Error: --name requires a value." >&2
+        exit 1
+      fi
+      NAME_ARG="--name $2"
+      shift 2
+      ;;
+    --port|-p)
+      if [[ -z "$2" ]]; then
+        echo "Error: --port requires a value." >&2
+        exit 1
+      fi
+      PORT_ARG="-p $2:$2"
+      shift 2
+      ;;
+    --gpus|-g)
+      if [[ -z "$2" ]]; then
+        echo "Error: --gpu requires a value." >&2
+        exit 1
+      fi
+      if [[ "$2" == "all" ]]; then
+        GPU_ARG="--gpus all"
+        shift 2
+        continue
+      fi
+      GPU_ARG="--gpus \"device=$2\""
+      shift 2
+      ;;
+    *)
+      echo "Error: unknown option: $1" >&2
+      exit 1
+      ;;
+  esac
+done
+
+echo "NAME_ARG: $NAME_ARG"
+echo "PORT_ARG: $PORT_ARG"
+echo "GPU_ARG: $GPU_ARG"
+
+docker run -it \
+  --shm-size 24G \
+  -e TZ=Asia/Tokyo \
+  -w /workspace \
+  -v "$(pwd)":/workspace \
+  $NAME_ARG \
+  $PORT_ARG \
+  $GPU_ARG \
+  psv_utils \
+  /bin/bash

google_drive_utils/download_from_drive.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python3
+import argparse
+import os
+import sys
+import time
+from pathlib import Path
+
+import requests
+from tqdm import tqdm
+
+from google_auth import get_drive_credentials, request_with_refresh
+
+
+def main():
+    CREDENTIALS = "secrets/credentials.json"
+    TOKEN = "secrets/token.json"
+
+    DRIVE_API = "https://www.googleapis.com/drive/v3/files"
+    CHUNK = 128 * 1024 * 1024          # 128MB
+    STREAM_CHUNK = 8 * 1024 * 1024     # 8MB
+    TIMEOUT = (10, 60)                 # (connect, read)
+    RETRY_SLEEP = 5
+
+    p = argparse.ArgumentParser()
+    p.add_argument("--file-id", required=True)
+    p.add_argument("--port", default=8100, type=int)
+    p.add_argument("--out-dir", default=".")
+    args = p.parse_args()
+
+    # token.json があれば非対話で利用（refresh も自動）
+    try:
+        creds = get_drive_credentials(
+            port=args.port, credentials_path=CREDENTIALS, token_path=TOKEN
+        )
+    except FileNotFoundError:
+        print(f"[ERROR] {CREDENTIALS} が見つかりません。")
+        sys.exit(1)
+
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    # メタデータ（name, size）
+    meta = request_with_refresh(
+        lambda token: requests.get(
+            f"{DRIVE_API}/{args.file_id}",
+            params={"fields": "name,size"},
+            headers={"Authorization": f"Bearer {token}"},
+            timeout=TIMEOUT,
+        ),
+        creds,
+        TOKEN,
+    )
+    meta.raise_for_status()
+    meta = meta.json()
+    name = meta["name"]
+    size = int(meta["size"])
+
+    part = out_dir / (name + ".part")
+    pos = part.stat().st_size if part.exists() else 0
+
+    # 本体（alt=media + Range）
+    url = f"{DRIVE_API}/{args.file_id}"
+    params = {"alt": "media"}
+    pbar = tqdm(total=size, initial=pos, unit="B", unit_scale=True, desc=name)
+
+    with open(part, "ab") as f:
+        while pos < size:
+            end = min(pos + CHUNK - 1, size - 1)
+
+            try:
+                r = request_with_refresh(
+                    lambda token: requests.get(
+                        url,
+                        params=params,
+                        headers={
+                            "Authorization": f"Bearer {token}",
+                            "Range": f"bytes={pos}-{end}",
+                        },
+                        stream=True,
+                        timeout=TIMEOUT,
+                    ),
+                    creds,
+                    TOKEN,
+                )
+
+                r.raise_for_status()
+
+                for chunk in r.iter_content(STREAM_CHUNK):
+                    if chunk:
+                        f.write(chunk)
+                        pos += len(chunk)
+                        pbar.update(len(chunk))
+
+                f.flush()
+                os.fsync(f.fileno())
+
+            except Exception:
+                time.sleep(RETRY_SLEEP)
+
+    pbar.close()
+    part.rename(out_dir / name)
+
+
+if __name__ == "__main__":
+    main()

google_drive_utils/google_auth.py

@@ -0,0 +1,68 @@
+#!/usr/bin/env python3
+import os
+from pathlib import Path
+
+from google_auth_oauthlib.flow import InstalledAppFlow
+from google.auth.transport.requests import Request
+from google.oauth2.credentials import Credentials
+
+# Scopes that allow both download and upload operations
+DEFAULT_DRIVE_SCOPES = [
+    "https://www.googleapis.com/auth/drive.readonly",
+    "https://www.googleapis.com/auth/drive.file",
+]
+
+
+def get_drive_credentials(
+    *,
+    port: int = 8100,
+    scopes=None,
+    credentials_path: str = "secrets/credentials.json",
+    token_path: str = "secrets/token.json",
+):
+    if scopes is None:
+        scopes = DEFAULT_DRIVE_SCOPES
+
+    if not os.path.exists(credentials_path):
+        raise FileNotFoundError(f"Credentials not found: {credentials_path}")
+
+    creds = None
+    if os.path.exists(token_path):
+        creds = Credentials.from_authorized_user_file(token_path, scopes)
+        if not creds.has_scopes(scopes):
+            os.remove(token_path)
+            creds = None
+
+    if creds and creds.expired and creds.refresh_token:
+        if not _refresh_token(creds, token_path):
+            creds = None
+
+    if not creds or not creds.valid:
+        flow = InstalledAppFlow.from_client_secrets_file(credentials_path, scopes)
+        creds = flow.run_local_server(
+            open_browser=False, host="localhost", bind_addr="0.0.0.0", port=int(port)
+        )
+        Path(token_path).write_text(creds.to_json())
+
+    return creds
+
+
+def request_with_refresh(request_fn, creds, token_path: str):
+    response = request_fn(creds.token)
+
+    if response.status_code == 401 and creds.refresh_token:
+        if _refresh_token(creds, token_path):
+            response = request_fn(creds.token)
+
+    return response
+
+
+def _refresh_token(creds, token_path: str) -> bool:
+    try:
+        creds.refresh(Request())
+        Path(token_path).write_text(creds.to_json())
+        return True
+    except Exception:
+        if os.path.exists(token_path):
+            os.remove(token_path)
+        return False