1+ #! /bin/bash
2+
3+ # Runs transformers benchmark in olmocr-bench
4+ set -e
5+
6+
7+ # Check for uncommitted changes
8+ if ! git diff-index --quiet HEAD --; then
9+ echo " Error: There are uncommitted changes in the repository."
10+ echo " Please commit or stash your changes before running the benchmark."
11+ echo " "
12+ echo " Uncommitted changes:"
13+ git status --short
14+ exit 1
15+ fi
16+
17+ # Use conda environment Python if available, otherwise use system Python
18+ if [ -n " $CONDA_PREFIX " ; then
19+ PYTHON=" $CONDA_PREFIX /bin/python"
20+ echo " Using conda Python from: $CONDA_PREFIX "
21+ else
22+ PYTHON=" python"
23+ echo " Warning: No conda environment detected, using system Python"
24+ fi
25+
26+ # Get version from version.py
27+ VERSION=$( $PYTHON -c ' import olmocr.version; print(olmocr.version.VERSION)' )
28+ echo " OlmOCR version: $VERSION "
29+
30+ # Get first 10 characters of git hash
31+ GIT_HASH=$( git rev-parse HEAD | cut -c1-10)
32+ echo " Git hash: $GIT_HASH "
33+
34+ # Get current git branch name
35+ GIT_BRANCH=$( git rev-parse --abbrev-ref HEAD)
36+ echo " Git branch: $GIT_BRANCH "
37+
38+ # Create full image tag
39+ IMAGE_TAG=" olmocr-benchmark-${VERSION} -${GIT_HASH} "
40+ echo " Building Docker image with tag: $IMAGE_TAG "
41+
42+ # Build the Docker image
43+ echo " Building Docker image..."
44+ docker build --platform linux/amd64 -f ./Dockerfile -t $IMAGE_TAG .
45+
46+ # Get Beaker username
47+ BEAKER_USER=$( beaker account whoami --format json | jq -r ' .[0].name' )
48+ echo " Beaker user: $BEAKER_USER "
49+
50+ # Push image to beaker
51+ echo " Trying to push image to Beaker..."
52+ if ! beaker image create --workspace ai2/oe-data-pdf --name $IMAGE_TAG $IMAGE_TAG 2> /dev/null; then
53+ echo " Warning: Beaker image with tag $IMAGE_TAG already exists. Using existing image."
54+ fi
55+
56+ # Create Python script to run beaker experiment
57+ cat << 'EOF ' > /tmp/run_benchmark_experiment.py
58+ import sys
59+ from beaker import Beaker, ExperimentSpec, TaskSpec, TaskContext, ResultSpec, TaskResources, ImageSource, Priority, Constraints, EnvVar
60+
61+ # Get image tag, beaker user, git branch, git hash, version from command line
62+ image_tag = sys.argv[1]
63+ beaker_user = sys.argv[2]
64+ git_branch = sys.argv[3]
65+ git_hash = sys.argv[4]
66+
67+
68+ # Initialize Beaker client
69+ b = Beaker.from_env(default_workspace="ai2/olmocr")
70+
71+
72+ # Check if AWS credentials secret exists
73+ aws_creds_secret = f"{beaker_user}-AWS_CREDENTIALS_FILE"
74+ try:
75+ # Try to get the secret to see if it exists
76+ b.secret.get(aws_creds_secret, workspace="ai2/olmocr")
77+ has_aws_creds = True
78+ print(f"Found AWS credentials secret: {aws_creds_secret}")
79+ except:
80+ has_aws_creds = False
81+ print(f"AWS credentials secret not found: {aws_creds_secret}")
82+
83+ # First experiment: Original benchmark job
84+ commands = []
85+ if has_aws_creds:
86+ commands.extend([
87+ "mkdir -p ~/.aws",
88+ 'echo "$AWS_CREDENTIALS_FILE" > ~/.aws/credentials'
89+ ])
90+ commands.extend([
91+ "git clone https://huggingface.co/datasets/allenai/olmOCR-bench",
92+ "cd olmOCR-bench && git lfs pull && cd ..",
93+ "python -m olmocr.bench.convert transformers:target_longest_image_dim=1288:prompt_template=yaml:response_template=yaml: --dir ./olmOCR-bench/bench_data",
94+ "python -m olmocr.bench.benchmark --dir ./olmOCR-bench/bench_data"
95+ ])
96+
97+ # Build task spec with optional env vars
98+ task_spec_args = {
99+ "name": "transformers-benchmark",
100+ "image": ImageSource(beaker=f"{beaker_user}/{image_tag}"),
101+ "command": [
102+ "bash", "-c",
103+ " && ".join(commands)
104+ ],
105+ "context": TaskContext(
106+ priority=Priority.normal,
107+ preemptible=True,
108+ ),
109+ "resources": TaskResources(gpu_count=1),
110+ "constraints": Constraints(cluster=["ai2/ceres-cirrascale", "ai2/jupiter-cirrascale-2"]),
111+ "result": ResultSpec(path="/noop-results"),
112+ }
113+
114+ # Add env vars if AWS credentials exist
115+ if has_aws_creds:
116+ task_spec_args["env_vars"] = [
117+ EnvVar(name="AWS_CREDENTIALS_FILE", secret=aws_creds_secret)
118+ ]
119+
120+ # Create first experiment spec
121+ experiment_spec = ExperimentSpec(
122+ description=f"Transformers Benchmark Run - Branch: {git_branch}, Commit: {git_hash}",
123+ budget="ai2/oe-data",
124+ tasks=[TaskSpec(**task_spec_args)],
125+ )
126+
127+ # Create the first experiment
128+ experiment = b.experiment.create(spec=experiment_spec, workspace="ai2/olmocr")
129+ print(f"Created benchmark experiment: {experiment.id}")
130+ print(f"View at: https://beaker.org/ex/{experiment.id}")
131+ print("-------")
132+ print("")
133+
134+
135+ EOF
136+
137+ # Run the Python script to create the experiments
138+ echo " Creating Beaker experiments..."
139+ $PYTHON /tmp/run_benchmark_experiment.py $IMAGE_TAG $BEAKER_USER $GIT_BRANCH $GIT_HASH
140+
141+ # Clean up temporary file
142+ rm /tmp/run_benchmark_experiment.py
143+
144+ echo " Benchmark experiments submitted successfully!"
0 commit comments