social-data-attribution-w2/scripts/modal/config.py

"""Constants and path helpers for SOC-91 and SOC-127 Modal workflows."""

from __future__ import annotations

import hashlib

from dolma.constants import DOLMA_6T_MIX_DATASET_ID, DOLMA_POOL_DATASET_ID

R2_BUCKET_NAME = "soc127-dedup"
R2_BUCKET = R2_BUCKET_NAME
R2_ENDPOINT_URL = "https://0934ab8e84ac8f4e81decaf3eb121337.r2.cloudflarestorage.com"
R2_SECRET_NAME = "r2-credentials"
R2_PREFIX = "soc127"

R2_INPUT_PREFIXES = [
    "soc127/phase1_pool_shared",
    "soc127/phase2_nonpool_final",
]
R2_OUTPUT_PREFIX = "soc91-labels"
R2_STATS_PREFIX = "soc91-stats"

TOPIC_URL_MODEL = "WebOrganizer/TopicClassifier"
TOPIC_NOURL_MODEL = "WebOrganizer/TopicClassifier-NoURL"
FORMAT_URL_MODEL = "WebOrganizer/FormatClassifier"
FORMAT_NOURL_MODEL = "WebOrganizer/FormatClassifier-NoURL"

ALL_MODELS = [
    TOPIC_URL_MODEL,
    TOPIC_NOURL_MODEL,
    FORMAT_URL_MODEL,
    FORMAT_NOURL_MODEL,
]

GPU_CONFIGS = {
    "T4": {"gpu": "T4", "dtype": "fp16", "memory": 16},
    "L4": {"gpu": "L4", "dtype": "bf16", "memory": 24},
    "A10G": {"gpu": "A10G", "dtype": "bf16", "memory": 24},
}

DEFAULT_GPU = "L4"
DEFAULT_BATCH_SIZE = 64
DEFAULT_MAX_LENGTH = 1024
DEFAULT_TIMEOUT = 1800
DEFAULT_CONCURRENCY = 1000
MODEL_VOLUME_NAME = "soc91-model-cache"

BLOOM_IMAGE_PATH = "/bloom/unique_6t_ids.bloom"
WORKER_CPU = 2
WORKER_MEMORY = 16384
WORKER_EPHEMERAL_DISK = 524288
WORKER_TIMEOUT = 7200
WORKER_RETRIES = 3
BUCKET_COUNT = 256
HF_SUBFOLDER_MAX = 5000

__all__ = [
    "ALL_MODELS",
    "BLOOM_IMAGE_PATH",
    "BUCKET_COUNT",
    "DEFAULT_BATCH_SIZE",
    "DEFAULT_CONCURRENCY",
    "DEFAULT_GPU",
    "DEFAULT_MAX_LENGTH",
    "DEFAULT_TIMEOUT",
    "DOLMA_6T_MIX_DATASET_ID",
    "DOLMA_POOL_DATASET_ID",
    "FORMAT_NOURL_MODEL",
    "FORMAT_URL_MODEL",
    "GPU_CONFIGS",
    "HF_SUBFOLDER_MAX",
    "MODEL_VOLUME_NAME",
    "R2_BUCKET",
    "R2_BUCKET_NAME",
    "R2_ENDPOINT_URL",
    "R2_INPUT_PREFIXES",
    "R2_OUTPUT_PREFIX",
    "R2_PREFIX",
    "R2_SECRET_NAME",
    "R2_STATS_PREFIX",
    "TOPIC_NOURL_MODEL",
    "TOPIC_URL_MODEL",
    "WORKER_CPU",
    "WORKER_EPHEMERAL_DISK",
    "WORKER_MEMORY",
    "WORKER_RETRIES",
    "WORKER_TIMEOUT",
    "r2_done_path",
    "r2_shard_output_path",
    "r2_stats_path",
    "shard_hash",
    "subfolder_index",
]


def shard_hash(shard_path: str) -> str:
    return hashlib.blake2b(shard_path.encode("utf-8"), digest_size=8).hexdigest()


def subfolder_index(shard_index: int, max_per_folder: int) -> str:
    folder_num = shard_index // max_per_folder
    return f"part_{folder_num:03d}"


def r2_shard_output_path(phase: str, shard_path: str) -> str:
    filename = shard_path.replace("/", "__")
    return f"{R2_PREFIX}/{phase}/{filename}"


def r2_stats_path(phase: str, shard_path: str) -> str:
    return f"{r2_shard_output_path(phase, shard_path)}.stats.json"


def r2_done_path(phase: str, shard_path: str) -> str:
    return f"{r2_shard_output_path(phase, shard_path)}.done"