Skip to content

Commit ed3fc6d

Browse files
atharvasatharvas
committed
break up Dockerfile into another building stage. Saves ~200GB of docker image space
1 parent 3d63f9d commit ed3fc6d

8 files changed

Lines changed: 451 additions & 1035 deletions

‎scratch/merged_context_registry_2025-09-04T08:32:08.486247.json‎

Lines changed: 280 additions & 280 deletions
Large diffs are not rendered by default.

‎scratch/scripts/initialize_context_registry.py‎

Lines changed: 111 additions & 112 deletions
Original file line numberDiff line numberDiff line change
@@ -7,127 +7,126 @@
77

88
CONTEXT_REGISTRY = ContextRegistry(default_context=DockerContext())
99

10-
# CONTEXT_REGISTRY.register(
11-
# "astropy/astropy:pkg",
12-
# DockerContext(
13-
# building_data="""#!/usr/bin/env bash
14-
# # Purpose: Build/install the repo (editable) in one or more ASV micromamba envs, then run health checks.
15-
# set -euo pipefail
16-
17-
# ###### SETUP CODE (NOT TO BE MODIFIED) ######
18-
# # Loads micromamba, common helpers, and persisted variables from the env stage.
19-
# source /etc/profile.d/asv_utils.sh || true
20-
# source /etc/profile.d/asv_build_vars.sh || true
21-
# eval "$(micromamba shell hook --shell=bash)"
10+
CONTEXT_REGISTRY.register(
11+
"astropy/astropy:pkg",
12+
DockerContext(
13+
building_data="""#!/usr/bin/env bash
14+
# Purpose: Build/install the repo (editable) in one or more ASV micromamba envs, then run health checks.
15+
set -euo pipefail
2216
23-
# ROOT_PATH=${ROOT_PATH:-$PWD} # Usually /workspace/repo
24-
# REPO_ROOT="$ROOT_PATH"
25-
# TARGET_VERSIONS="${PY_VERSION:-${ASV_PY_VERSIONS:-}}"
26-
# EXTRAS="${ALL_EXTRAS:+[$ALL_EXTRAS]}"
27-
# if [[ -z "${TARGET_VERSIONS}" ]]; then
28-
# echo "Error: No PY_VERSION set and ASV_PY_VERSIONS not found." >&2
29-
# exit 1
30-
# fi
31-
# ###### END SETUP CODE ######
17+
###### SETUP CODE (NOT TO BE MODIFIED) ######
18+
# Loads micromamba, common helpers, and persisted variables from the env stage.
19+
source /etc/profile.d/asv_utils.sh || true
20+
source /etc/profile.d/asv_build_vars.sh || true
3221
33-
# # -----------------------------
34-
# # Agent guidance (read-first)
35-
# # -----------------------------
36-
# # GOAL: For each Python version below, install the project in EDITABLE mode into env asv_{version},
37-
# # with NO build isolation, then run health checks.
38-
# #
39-
# # Below this comment, you should do whatever is necessary to build the project without errors. Including (but not limited to):
40-
# # - Add extra conda/pip dependencies needed to build this project.
41-
# # - Run repo-specific pre-steps (e.g., submodules, generating Cython, env vars).
42-
# # - Run arbitrary micromamba/pip commands in the target env.
43-
# # - Set CFLAGS/CXXFLAGS/LDFLAGS if needed for this repo.
44-
# # - Change files in the repo if needed (e.g., fix a missing #include).
45-
# # - Anything else needed to get a successful editable install.
46-
# #
47-
# # MUST:
48-
# # - Keep this script idempotent.
49-
# # - Use: `pip install --no-build-isolation -v -e .` or `pip install -e .` or equivalent.
50-
# # - Do not modify the SETUP CODE or helper functions below.
51-
# #
52-
# # DO NOT:
53-
# # - Change env names or Python versions outside MODEL EDIT AREA.
54-
# # - Use build isolation unless absolutely necessary.
22+
ROOT_PATH=${ROOT_PATH:-$PWD} # Usually /workspace/repo
23+
REPO_ROOT="$ROOT_PATH"
24+
TARGET_VERSIONS="${PY_VERSION:-${ASV_PY_VERSIONS:-}}"
25+
EXTRAS="${ALL_EXTRAS:+[$ALL_EXTRAS]}"
26+
if [[ -z "${TARGET_VERSIONS}" ]]; then
27+
echo "Error: No PY_VERSION set and ASV_PY_VERSIONS not found." >&2
28+
exit 1
29+
fi
30+
###### END SETUP CODE ######
5531
56-
# # -----------------------------
57-
# # Helpers (do not modify)
58-
# # -----------------------------
59-
# log() { printf "\033[1;34m[build]\033[0m %s\n" "$*"; }
60-
# warn() { printf "\033[1;33m[warn]\033[0m %s\n" "$*" >&2; }
61-
# die() { printf "\033[1;31m[fail]\033[0m %s\n" "$*" >&2; exit 1; }
32+
# -----------------------------
33+
# Agent guidance (read-first)
34+
# -----------------------------
35+
# GOAL: For each Python version below, install the project in EDITABLE mode into env asv_{version},
36+
# with NO build isolation, then run health checks.
37+
#
38+
# Below this comment, you should do whatever is necessary to build the project without errors. Including (but not limited to):
39+
# - Add extra conda/pip dependencies needed to build this project.
40+
# - Run repo-specific pre-steps (e.g., submodules, generating Cython, env vars).
41+
# - Run arbitrary micromamba/pip commands in the target env.
42+
# - Set CFLAGS/CXXFLAGS/LDFLAGS if needed for this repo.
43+
# - Change files in the repo if needed (e.g., fix a missing #include).
44+
# - Anything else needed to get a successful editable install.
45+
#
46+
# MUST:
47+
# - Keep this script idempotent.
48+
# - Use: `pip install --no-build-isolation -v -e .` or `pip install -e .` or equivalent.
49+
# - Do not modify the SETUP CODE or helper functions below.
50+
#
51+
# DO NOT:
52+
# - Change env names or Python versions outside MODEL EDIT AREA.
53+
# - Use build isolation unless absolutely necessary.
6254
63-
# # Conservative default parallelism (override if the repo benefits)
64-
# export CMAKE_BUILD_PARALLEL_LEVEL="${CMAKE_BUILD_PARALLEL_LEVEL:-2}"
65-
# export NPY_NUM_BUILD_JOBS="${NPY_NUM_BUILD_JOBS:-2}"
55+
# -----------------------------
56+
# Helpers (do not modify)
57+
# -----------------------------
58+
log() { printf "\033[1;34m[build]\033[0m %s\n" "$*"; }
59+
warn() { printf "\033[1;33m[warn]\033[0m %s\n" "$*" >&2; }
60+
die() { printf "\033[1;31m[fail]\033[0m %s\n" "$*" >&2; exit 1; }
6661
67-
# # -----------------------------
68-
# # Build & test across envs
69-
# # -----------------------------
70-
# for version in $TARGET_VERSIONS; do
71-
# ENV_NAME="asv_${version}"
72-
# log "==> Building in env: $ENV_NAME (python=$version)"
62+
# Conservative default parallelism (override if the repo benefits)
63+
export CMAKE_BUILD_PARALLEL_LEVEL="${CMAKE_BUILD_PARALLEL_LEVEL:-2}"
64+
export NPY_NUM_BUILD_JOBS="${NPY_NUM_BUILD_JOBS:-2}"
7365
74-
# if ! micromamba env list | awk '{print $1}' | grep -qx "$ENV_NAME"; then
75-
# die "Env $ENV_NAME not found. Did docker_build_env.sh run?"
76-
# fi
66+
# -----------------------------
67+
# Build & test across envs
68+
# -----------------------------
69+
for version in $TARGET_VERSIONS; do
70+
ENV_NAME="asv_${version}"
71+
log "==> Building in env: $ENV_NAME (python=$version)"
7772
78-
# # Import name resolution (kept simple for the agent)
79-
# IMP="${IMPORT_NAME:-}"
80-
# if [[ -z "$IMP" ]]; then
81-
# if ! IMP="$(asv_detect_import_name --repo-root "$REPO_ROOT" 2>/dev/null)"; then
82-
# die "Could not determine import name. Set IMPORT_NAME in /etc/profile.d/asv_build_vars.sh"
83-
# fi
84-
# fi
85-
# log "Using import name: $IMP"
73+
if ! micromamba env list | awk '{print $1}' | grep -qx "$ENV_NAME"; then
74+
die "Env $ENV_NAME not found. Did docker_build_env.sh run?"
75+
fi
8676
87-
# # -----------------------------
88-
# # MODEL EDIT AREA: repo-specific tweaks (optional)
89-
# # -----------------------------
90-
# # Examples (uncomment if needed for this repo):
91-
# #
92-
# # log "Updating submodules"
93-
# # git -C "$REPO_ROOT" submodule update --init --recursive
94-
# #
95-
# # log "Installing extra system libs via conda-forge"
96-
# # micromamba install -y -n "$ENV_NAME" -c conda-forge 'openblas' 'blas=*=openblas' 'libopenmp'
97-
# #
98-
# # log "Pre-generating Cython sources"
99-
# # micromamba run -n "$ENV_NAME" python -m cython --version
100-
# #
101-
# # export CFLAGS="${CFLAGS:-}"
102-
# # export CXXFLAGS="${CXXFLAGS:-}"
103-
# # export LDFLAGS="${LDFLAGS:-}"
104-
# # -----------------------------
77+
# Import name resolution (kept simple for the agent)
78+
IMP="${IMPORT_NAME:-}"
79+
if [[ -z "$IMP" ]]; then
80+
if ! IMP="$(detect_import_name --repo-root "$REPO_ROOT" 2>/dev/null)"; then
81+
die "Could not determine import name. Set IMPORT_NAME in /etc/profile.d/asv_build_vars.sh"
82+
fi
83+
fi
84+
log "Using import name: $IMP"
10585
106-
# # Install some basic micromamba packages.
107-
# micromamba install -y -n "$ENV_NAME" -c conda-forge git conda mamba "libmambapy<=1.9.9"
86+
# -----------------------------
87+
# MODEL EDIT AREA: repo-specific tweaks (optional)
88+
# -----------------------------
89+
# Examples (uncomment if needed for this repo):
90+
#
91+
# log "Updating submodules"
92+
# git -C "$REPO_ROOT" submodule update --init --recursive
93+
#
94+
# log "Installing extra system libs via conda-forge"
95+
# micromamba install -y -n "$ENV_NAME" -c conda-forge 'openblas' 'blas=*=openblas' 'libopenmp'
96+
#
97+
# log "Pre-generating Cython sources"
98+
# micromamba run -n "$ENV_NAME" python -m cython --version
99+
#
100+
# export CFLAGS="${CFLAGS:-}"
101+
# export CXXFLAGS="${CXXFLAGS:-}"
102+
# export LDFLAGS="${LDFLAGS:-}"
103+
# -----------------------------
108104
109-
# export CFLAGS="${CFLAGS:-} -Wno-error=incompatible-pointer-types"
110-
# micromamba run -n "$ENV_NAME" pip install -e . scipy matplotlib
105+
# Install some generic packages needed for building/testing.
106+
# Libmambapy must be < 2 avoid issues running airspeed-velocity.
107+
micromamba install -y -n "$ENV_NAME" -c conda-forge pip git conda mamba "libmambapy<=1.9.9" \
108+
"numpy<2" scipy "cython<2" joblib fakeredis threadpoolctl pytest \
109+
compilers meson-python cmake ninja pkg-config tomli extension-helpers
110+
export CFLAGS="${CFLAGS:-} -Wno-error=incompatible-pointer-types"
111111
112-
# # Editable install (no build isolation preferrably). Toolchain lives in the env already.
113-
# # $EXTRAS is an optional argument to install all discovered extra dependencies.
114-
# # It will be empty if pyproject.toml does not exist or has no [project.optional-dependencies].
115-
# # In case setup.py is used, no need to append $EXTRAS.
116-
# log "Editable install with --no-build-isolation"
117-
# PIP_NO_BUILD_ISOLATION=1 micromamba run -n "$ENV_NAME" python -m pip install --no-build-isolation -v -e "$REPO_ROOT"$EXTRAS
112+
# Editable install (no build isolation preferrably). Toolchain lives in the env already.
113+
# $EXTRAS is an optional argument to install all discovered extra dependencies.
114+
# It will be empty if pyproject.toml does not exist or has no [project.optional-dependencies].
115+
# In case setup.py is used, no need to append $EXTRAS.
116+
log "Editable install with --no-build-isolation"
117+
PIP_NO_BUILD_ISOLATION=1 micromamba run -n "$ENV_NAME" python setup.py develop
118118
119-
# # Health checks (import + compiled extension probe; optional pytest smoke with RUN_PYTEST_SMOKE=1)
120-
# log "Running smoke checks"
121-
# micromamba run -n "$ENV_NAME" asv_smokecheck.py --import-name "$IMP" --repo-root "$REPO_ROOT" ${RUN_PYTEST_SMOKE:+--pytest-smoke}
119+
# Health checks (import + compiled extension probe; optional pytest smoke with RUN_PYTEST_SMOKE=1)
120+
log "Running smoke checks"
121+
micromamba run -n "$ENV_NAME" asv_smokecheck.py --import-name "$IMP" --repo-root "$REPO_ROOT" ${RUN_PYTEST_SMOKE:+--pytest-smoke}
122122
123-
# # Machine-readable markers (useful in logs)
124-
# echo "::import_name=${IMP}::env=${ENV_NAME}"
125-
# done
123+
echo "::import_name=${IMP}::env=${ENV_NAME}"
124+
done
126125
127-
# log "All builds complete ✅"
128-
# """.strip(),
129-
# ),
130-
# )
126+
log "All builds complete ✅"
127+
""".strip(),
128+
),
129+
)
131130

132131
CONTEXT_REGISTRY.register(
133132
"scikit-learn/scikit-learn:pkg",
@@ -200,7 +199,7 @@
200199
# Import name resolution (kept simple for the agent)
201200
IMP="${IMPORT_NAME:-}"
202201
if [[ -z "$IMP" ]]; then
203-
if ! IMP="$(asv_detect_import_name --repo-root "$REPO_ROOT" 2>/dev/null)"; then
202+
if ! IMP="$(detect_import_name --repo-root "$REPO_ROOT" 2>/dev/null)"; then
204203
die "Could not determine import name. Set IMPORT_NAME in /etc/profile.d/asv_build_vars.sh"
205204
fi
206205
fi
@@ -321,7 +320,7 @@
321320
# Import name resolution (kept simple for the agent)
322321
IMP="${IMPORT_NAME:-}"
323322
if [[ -z "$IMP" ]]; then
324-
if ! IMP="$(asv_detect_import_name --repo-root "$REPO_ROOT" 2>/dev/null)"; then
323+
if ! IMP="$(detect_import_name --repo-root "$REPO_ROOT" 2>/dev/null)"; then
325324
die "Could not determine import name. Set IMPORT_NAME in /etc/profile.d/asv_build_vars.sh"
326325
fi
327326
fi
@@ -449,7 +448,7 @@
449448
# Import name resolution (kept simple for the agent)
450449
IMP="${IMPORT_NAME:-}"
451450
if [[ -z "$IMP" ]]; then
452-
if ! IMP="$(asv_detect_import_name --repo-root "$REPO_ROOT" 2>/dev/null)"; then
451+
if ! IMP="$(detect_import_name --repo-root "$REPO_ROOT" 2>/dev/null)"; then
453452
die "Could not determine import name. Set IMPORT_NAME in /etc/profile.d/asv_build_vars.sh"
454453
fi
455454
fi
@@ -574,7 +573,7 @@
574573
# Import name resolution (kept simple for the agent)
575574
IMP="${IMPORT_NAME:-}"
576575
if [[ -z "$IMP" ]]; then
577-
if ! IMP="$(asv_detect_import_name --repo-root "$REPO_ROOT" 2>/dev/null)"; then
576+
if ! IMP="$(detect_import_name --repo-root "$REPO_ROOT" 2>/dev/null)"; then
578577
die "Could not determine import name. Set IMPORT_NAME in /etc/profile.d/asv_build_vars.sh"
579578
fi
580579
fi
@@ -699,7 +698,7 @@
699698
# # Import name resolution (kept simple for the agent)
700699
# IMP="${IMPORT_NAME:-}"
701700
# if [[ -z "$IMP" ]]; then
702-
# if ! IMP="$(asv_detect_import_name --repo-root "$REPO_ROOT" 2>/dev/null)"; then
701+
# if ! IMP="$(detect_import_name --repo-root "$REPO_ROOT" 2>/dev/null)"; then
703702
# die "Could not determine import name. Set IMPORT_NAME in /etc/profile.d/asv_build_vars.sh"
704703
# fi
705704
# fi

‎scratch/scripts/update_context_registry.py‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,13 +11,22 @@ def update_cr_entrypoint(cr: ContextRegistry):
1111
new_v.entrypoint_data = DockerContext().entrypoint_data
1212
new_v.dockerfile_data = DockerContext().dockerfile_data
1313
new_v.env_building_data = DockerContext().env_building_data
14+
new_v.base_building_data = DockerContext().base_building_data
1415
new_reg[k] = new_v
1516

1617
cr.registry = new_reg
1718
return cr
1819

1920

2021
if __name__ == "__main__":
21-
cr = ContextRegistry.load_from_file(Path("scratch/artifacts/pipeflush/tiny/context_registry.json"))
22-
new_cr = update_cr_entrypoint(cr)
23-
new_cr.save_to_file(Path("scratch/artifacts/pipeflush/tiny/context_registry.json"))
22+
# cr = ContextRegistry.load_from_file(Path("scratch/merged_context_registry_2025-09-04T08:32:08.486247.json"))
23+
# new_cr = update_cr_entrypoint(cr)
24+
# new_cr.save_to_file(Path("scratch/merged_context_registry_2025-09-04T08:32:08.486247.json"))
25+
26+
# Detect all context_registry.json files in scratch/ and update them. The files might be named like merged_context_registry*.json or context_registry*.json
27+
for p in Path("scratch").glob("**/*context_registry*.json"):
28+
print(f"Updating {p}...")
29+
cr = ContextRegistry.load_from_file(p)
30+
new_cr = update_cr_entrypoint(cr)
31+
new_cr.save_to_file(p)
32+
print(f"Updated {p}.")

‎src/datasmith/agents/context_synthesis.py‎

Lines changed: 18 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -193,9 +193,9 @@ class BuildScriptAgentStep(dspy.Signature):
193193
desc="Describes where the failure occured. E.g. 'N/A', 'build failed', 'asv run failed'."
194194
)
195195
last_docker_build_script = dspy.InputField(desc="Previous docker_build.sh script.")
196-
initial_template = dspy.InputField(
197-
desc="Initial template of the docker_build.sh script with important instructions."
198-
)
196+
# initial_template = dspy.InputField(
197+
# desc="Initial template of the docker_build.sh script with important instructions."
198+
# )
199199
repo_facts_json = dspy.InputField(desc="Some inferred repo facts (A JSON object with paths, candidates, versions).")
200200
toolbelt = dspy.InputField(desc="Human-readable summary of available tools.")
201201
messages_log = dspy.InputField(desc="Transcript of prior tool actions & observations.")
@@ -237,7 +237,7 @@ def forward(
237237
stdout_logs: str,
238238
failure_more: str,
239239
last_docker_build_script: str,
240-
initial_template: str,
240+
# initial_template: str,
241241
repo_facts_json: str,
242242
tool_executor: ContainerToolExecutor,
243243
max_steps: int = 4,
@@ -263,7 +263,7 @@ def forward(
263263
stdout_logs=stdout_logs or "",
264264
failure_more=failure_more or "N/A",
265265
last_docker_build_script=last_docker_build_script or "",
266-
initial_template=initial_template,
266+
# initial_template=initial_template,
267267
repo_facts_json=repo_facts_json or "{}",
268268
toolbelt=toolbelt,
269269
messages_log=messages_log,
@@ -285,6 +285,10 @@ def forward(
285285

286286
messages_log += f"\n\n# Step [{step_idx + 1}/{max_steps}]\n# Action: {action}\n# Input: {action_input}\n# Observation:\n{observation[:4000]}"
287287

288+
if action in ("none", "finish"):
289+
# Model is done but didn't provide a script. Stop.
290+
break
291+
288292
# Don't prefer build_script until model is completely done with it.
289293
# # If model already emitted a script, prefer it
290294
# if (out.docker_build_script or "").strip(): # pyright: ignore[reportAttributeAccessIssue]
@@ -305,7 +309,10 @@ def forward(
305309
# script = out.docker_build_script.strip() # pyright: ignore[reportAttributeAccessIssue]
306310
script = (iter_script or "").strip()
307311
logger.debug("DSPy: candidate script preview: %s", _preview(script, 240))
308-
must_haves = ["###### SETUP CODE (NOT TO BE MODIFIED) ######"]
312+
# source /etc/profile.d/asv_utils.sh || true
313+
# source /etc/profile.d/asv_build_vars.sh || true
314+
315+
must_haves = ["/etc/profile.d/asv_utils.sh", "/etc/profile.d/asv_build_vars.sh"]
309316
ok_template = all(m in script for m in must_haves)
310317
must_not_haves = ["```bash", "```", "import IPython", "from IPython"]
311318
no_bad = all(m not in script for m in must_not_haves)
@@ -343,7 +350,7 @@ def synthesize_script(
343350
last_script: str,
344351
stderr_tail: str,
345352
stdout_tail: str,
346-
building_template: str,
353+
# building_template: str,
347354
failure_more: str,
348355
tool_exec: ContainerToolExecutor,
349356
max_steps: int = 4,
@@ -367,7 +374,7 @@ def synthesize_script(
367374
stdout_logs=stdout_tail or "",
368375
failure_more=failure_more or "N/A",
369376
last_docker_build_script=last_script or "",
370-
initial_template=building_template,
377+
# initial_template=building_template,
371378
repo_facts_json=tool_exec.facts_json(),
372379
tool_executor=tool_exec,
373380
max_steps=max_steps,
@@ -444,11 +451,11 @@ def agent_build_and_validate( # noqa: C901
444451
assert task.sha is not None, "task.sha must be set" # noqa: S101
445452
default_building_template = context_registry.get_default(tag="env")[1].building_data
446453
if len(similar_contexts := context_registry.get_similar(task.with_tag("env"))) > 0:
447-
_, context = similar_contexts[0]
454+
t, context = similar_contexts[0]
448455
logger.info(
449456
"build_once_with_context: found %d similar contexts; using most similar with key=%s",
450457
len(similar_contexts),
451-
str(context),
458+
str(t),
452459
)
453460
first_guess = context.building_data
454461
else:
@@ -540,7 +547,7 @@ def agent_build_and_validate( # noqa: C901
540547
attempts[-1].building_data,
541548
stderr_tail=stderr_tail,
542549
stdout_tail=stdout_tail,
543-
building_template=default_building_template,
550+
# building_template=default_building_template,
544551
failure_more=failure_more,
545552
tool_exec=tool_exec,
546553
max_steps=args.max_steps,

0 commit comments

Comments
 (0)