dr-environment 0.1.0__py3-none-any.whl

dr_environment/__init__.py

@@ -0,0 +1,14 @@
+#
+# Copyright 2026 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

dr_environment/cli.py

@@ -0,0 +1,91 @@
+#
+# Copyright 2026 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""CLI entrypoint for dr-environment plugin."""
+
+from __future__ import annotations
+
+import json
+import sys
+from importlib.metadata import PackageNotFoundError, version
+from pathlib import Path
+
+import click
+
+from dr_environment.recipe import build
+
+
+def _package_version() -> str:
+    try:
+        return version("dr-environment")
+    except PackageNotFoundError:
+        return "0.0.0"
+
+
+def print_manifest() -> None:
+    manifest = {
+        "name": "environment",
+        "version": _package_version(),
+        "description": "Build execution environment Docker contexts with offline dependency caches",
+        "authentication": False,
+    }
+    print(json.dumps(manifest, indent=2))
+
+
+@click.group()
+@click.version_option(_package_version(), "--version")
+def cli() -> None:
+    """Build DataRobot execution environment Docker contexts."""
+
+
+@cli.command("recipe")
+@click.option(
+    "--recipe-path",
+    default=".",
+    type=click.Path(exists=True, file_okay=False, path_type=Path),
+    help="Recipe root directory, defaults to the current working directory",
+)
+@click.option(
+    "--target",
+    default="docker_context",
+    help="Output directory relative to the current working directory, defaults to docker_context",
+)
+@click.option("--no-tarball", is_flag=True, help="Skip docker_context.tar.gz creation")
+def recipe_cmd(recipe_path: Path, target: str, no_tarball: bool) -> None:
+    """Build the recipe Codespace execution environment docker context."""
+    docker_context = build(
+        recipe_path.resolve(),
+        Path(target),
+        tarball=not no_tarball,
+    )
+    click.echo(f"Built docker context: {docker_context}")
+    archive = docker_context.parent / "docker_context.tar.gz"
+    if archive.is_file():
+        click.echo(f"Archive: {archive}")
+
+
+def main() -> None:
+    if "--dr-plugin-manifest" in sys.argv:
+        print_manifest()
+        return
+
+    try:
+        cli(prog_name="dr-environment")
+    except Exception as exc:
+        click.echo(f"Error: {exc}", err=True)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

dr_environment/recipe/__init__.py

@@ -0,0 +1,19 @@
+#
+# Copyright 2026 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Recipe docker context builder."""
+
+from dr_environment.recipe.build import build
+
+__all__ = ["build"]

dr_environment/recipe/build.py

@@ -0,0 +1,99 @@
+#
+# Copyright 2026 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Orchestrate docker context build."""
+
+from __future__ import annotations
+
+import shutil
+import tarfile
+from pathlib import Path
+
+import yaml
+
+from dr_environment.recipe.cache.stages import write_component_cache_fragments
+from dr_environment.recipe.discover import discover_components
+from dr_environment.recipe.hooks import run_environment_hook
+from dr_environment.recipe.layout import layout_components
+from dr_environment.recipe.models import ComponentStrategy
+from dr_environment.recipe.render import (
+    assemble_dockerfile,
+    copy_fragment_assets,
+    render_base_fragment,
+    render_build_deps_fragment,
+    render_kernel_setup_fragment,
+    render_offline_fragment,
+    render_user_fragment,
+    render_versions_fragment,
+)
+from dr_environment.recipe.validate import inspect_component, validate_all
+
+
+def load_versions(versions_file: Path) -> dict:
+    if not versions_file.is_file():
+        return {}
+    return yaml.safe_load(versions_file.read_text(encoding="utf-8")) or {}
+
+
+def build(
+    recipe_path: Path,
+    target: Path,
+    *,
+    tarball: bool = True,
+) -> Path:
+    recipe_path = recipe_path.resolve()
+    docker_context = target.resolve() if target.is_absolute() else (Path.cwd() / target).resolve()
+    versions_file = recipe_path / ".datarobot/cli/versions.yaml"
+
+    components = discover_components(recipe_path)
+    validate_all(components)
+    for component in components:
+        inspect_component(component)
+
+    if docker_context.exists():
+        shutil.rmtree(docker_context)
+    docker_context.mkdir(parents=True)
+
+    versions = load_versions(versions_file)
+    copy_fragment_assets(docker_context)
+    render_base_fragment(docker_context)
+    render_user_fragment(docker_context)
+    render_versions_fragment(docker_context, versions)
+    render_build_deps_fragment(docker_context)
+    render_kernel_setup_fragment(docker_context)
+
+    for component in components:
+        if component.strategy == ComponentStrategy.HOOK:
+            run_environment_hook(component, docker_context)
+        elif component.strategy == ComponentStrategy.DEFAULT:
+            inspect_component(component)
+            layout_components([component], docker_context)
+
+    active = [c for c in components if c.strategy == ComponentStrategy.DEFAULT and c.manifests]
+    cache_stage = write_component_cache_fragments(active, docker_context)
+
+    render_offline_fragment(docker_context, cache_stage=cache_stage)
+    assemble_dockerfile(docker_context)
+
+    if tarball:
+        create_tarball(docker_context)
+
+    return docker_context
+
+
+def create_tarball(docker_context: Path) -> Path:
+    archive = docker_context.parent / "docker_context.tar.gz"
+    with tarfile.open(archive, "w:gz") as tar:
+        tar.add(docker_context, arcname=".")
+    return archive

dr_environment/recipe/cache/__init__.py

@@ -0,0 +1,15 @@
+#
+# Copyright 2026 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Cache stage fragment generators."""

dr_environment/recipe/cache/stages.py

@@ -0,0 +1,151 @@
+#
+# Copyright 2026 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Dockerfile cache stage fragment generators."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from dr_environment.recipe.layout import LOCAL_SHARED_PACKAGE
+from dr_environment.recipe.models import Component, Ecosystem
+
+PREVIOUS_STAGE = "kernel"
+RUNTIME_USER = "notebooks"
+UV_CACHE = "/opt/cache/uv"
+NPM_CACHE = "/opt/cache/npm"
+GO_MOD_CACHE = "/opt/cache/go/pkg/mod"
+GO_BUILD_CACHE = "/opt/cache/go/build"
+GO_CACHE_ROOT = "/opt/cache/go"
+# Flat find-links directory of the exact published Python wheels/sdists. Kept separate from
+# the uv cache: DataRobot builds the FastAPI custom application FROM this image with
+# `uv pip install --system --no-cache`, and `--no-cache` makes uv ignore UV_CACHE. uv still
+# honours find-links, so the offline stage points UV_FIND_LINKS here (see _python_cache_lines).
+WHEELHOUSE = "/opt/wheelhouse"
+
+# Paths copied from the final cache stage into the offline runtime image.
+CACHE_COPY_PATHS = (
+    UV_CACHE,
+    NPM_CACHE,
+    GO_CACHE_ROOT,
+    WHEELHOUSE,
+)
+
+
+def write_component_cache_fragment(
+    component: Component,
+    docker_context: Path,
+    *,
+    previous_stage: str,
+) -> str:
+    """Write cache fragment for component; return new stage name."""
+    stage_name = f"cache-{component.name}"
+    lines: list[str] = [f"FROM {previous_stage} AS {stage_name}"]
+
+    for info in component.manifests:
+        if info.ecosystem == Ecosystem.PYTHON:
+            lines.extend(_python_cache_lines(component.name))
+        elif info.ecosystem == Ecosystem.NPM:
+            lines.extend(_npm_cache_lines(component.name))
+        elif info.ecosystem == Ecosystem.GO:
+            lines.extend(_go_cache_lines(component.name))
+
+    fragment_path = (
+        docker_context
+        / "dockerfile.d"
+        / f"{component.fragment_order:02d}-cache-{component.name}.fragment"
+    )
+    fragment_path.parent.mkdir(parents=True, exist_ok=True)
+    fragment_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
+    return stage_name
+
+
+def write_component_cache_fragments(
+    components: list[Component],
+    docker_context: Path,
+) -> str | None:
+    """Write all default component cache fragments; return final cache stage name."""
+    previous = PREVIOUS_STAGE
+    wrote = False
+    for component in components:
+        if not component.manifests:
+            continue
+        wrote = True
+        previous = write_component_cache_fragment(
+            component, docker_context, previous_stage=previous
+        )
+    return previous if wrote else None
+
+
+def _copy_component_tree(component_name: str) -> str:
+    return (
+        f"COPY --chown={RUNTIME_USER}:{RUNTIME_USER} "
+        f"components/{component_name}/ /tmp/cache-work/{component_name}/"
+    )
+
+
+def _python_cache_lines(component_name: str) -> list[str]:
+    sync_flags = "--frozen --no-install-project --all-extras --all-groups"
+    if component_name != LOCAL_SHARED_PACKAGE:
+        sync_flags += f" --no-install-package {LOCAL_SHARED_PACKAGE}"
+
+    warm_venv = f"/tmp/uv-cache-warm-{component_name}"
+    wheelhouse_req = f"/tmp/wheelhouse-req-{component_name}.txt"
+    return [
+        f"ENV UV_CACHE_DIR={UV_CACHE}",
+        _copy_component_tree(component_name),
+        f"WORKDIR /tmp/cache-work/{component_name}",
+        # Warm the uv cache; agent/MCP custom models install from it at runtime with
+        # `uv sync` (see the start_server.sh block in the offline stage). Mirror the same
+        # dependencies into the shared wheelhouse in the same RUN (see WHEELHOUSE above for
+        # why DataRobot's `--no-cache` app build cannot read the uv cache): `pip download`
+        # fetches the exact published artifacts, so their hashes match the ones `uv export`
+        # writes into the requirements file at deploy time; rebuilt wheels would not. Export
+        # exactly what DataRobot installs: production dependencies, project and workspace
+        # members excluded. warm_venv must be removed in this same RUN — splitting the sync
+        # and the rm across RUNs would bake the multi-GB venv into the earlier layer for good.
+        # No --only-binary=:all: here (unlike the copier/build-deps wheelhouse downloads):
+        # components pull in arbitrary packages we don't control, and plenty of legitimate
+        # pure-Python packages have stopped publishing wheels without being any less safe to
+        # build offline (e.g. rouge-score >=0.0.7 is sdist-only but pure Python, needing only
+        # the setuptools we already bake in) — rejecting all sdists here would fail builds
+        # over packages that build fine, not just the genuinely risky ones.
+        f"RUN UV_PROJECT_ENVIRONMENT={warm_venv} \\",
+        f"    uv sync {sync_flags} \\",
+        "        --python ${VENV_PATH}/bin/python \\",
+        f"    && uv pip install --no-cache --python {warm_venv}/bin/python pip \\",
+        "    && uv export --frozen --no-dev --no-emit-local --no-emit-project \\",
+        f"        -o {wheelhouse_req} \\",
+        f"    && {warm_venv}/bin/python -m pip download --no-deps --no-cache-dir \\",
+        f"        -r {wheelhouse_req} --dest {WHEELHOUSE} \\",
+        f"    && rm -rf {warm_venv} {wheelhouse_req}",
+    ]
+
+
+def _npm_cache_lines(component_name: str) -> list[str]:
+    return [
+        f"ENV NPM_CONFIG_CACHE={NPM_CACHE}",
+        _copy_component_tree(component_name),
+        f"WORKDIR /tmp/cache-work/{component_name}",
+        f"RUN npm ci --include=dev --cache {NPM_CACHE}",
+    ]
+
+
+def _go_cache_lines(component_name: str) -> list[str]:
+    return [
+        f"ENV GOMODCACHE={GO_MOD_CACHE} GOCACHE={GO_BUILD_CACHE}",
+        _copy_component_tree(component_name),
+        f"WORKDIR /tmp/cache-work/{component_name}",
+        "RUN go mod download all",
+    ]

dr_environment/recipe/discover.py

@@ -0,0 +1,88 @@
+#
+# Copyright 2026 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Discover recipe components from the root Taskfile."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import yaml
+
+from dr_environment.recipe.manifests import has_any_manifest
+from dr_environment.recipe.models import Component, ComponentStrategy
+
+TASKFILE_NAMES = ("Taskfile.yml", "Taskfile.yaml")
+
+
+def find_taskfile(directory: Path) -> Path | None:
+    for name in TASKFILE_NAMES:
+        candidate = directory / name
+        if candidate.is_file():
+            return candidate
+    return None
+
+
+def has_environment_task(taskfile: Path) -> bool:
+    data = yaml.safe_load(taskfile.read_text(encoding="utf-8")) or {}
+    tasks = data.get("tasks") or {}
+    return "environment" in tasks
+
+
+def discover_components(
+    recipe_path: Path,
+    *,
+    component_filter: set[str] | None = None,
+    skip_hooks: bool = False,
+) -> list[Component]:
+    """Return components from root Taskfile includes."""
+    root_taskfile = find_taskfile(recipe_path)
+    if root_taskfile is None:
+        raise FileNotFoundError(f"No Taskfile.yml found in {recipe_path}")
+
+    data = yaml.safe_load(root_taskfile.read_text(encoding="utf-8")) or {}
+    includes = data.get("includes") or {}
+
+    components: list[Component] = []
+    order = 10
+    for name, spec in includes.items():
+        if not isinstance(spec, dict):
+            continue
+        if spec.get("internal"):
+            continue
+        if component_filter and name not in component_filter:
+            continue
+
+        component_dir = Path(spec.get("dir", name))
+        if not component_dir.is_absolute():
+            component_dir = (recipe_path / component_dir).resolve()
+
+        taskfile = find_taskfile(component_dir)
+        strategy = ComponentStrategy.DEFAULT
+        if not skip_hooks and taskfile and has_environment_task(taskfile):
+            strategy = ComponentStrategy.HOOK
+        elif not has_any_manifest(component_dir):
+            strategy = ComponentStrategy.SKIP
+
+        components.append(
+            Component(
+                name=name,
+                source_dir=component_dir,
+                strategy=strategy,
+                fragment_order=order,
+            )
+        )
+        order += 1
+
+    return components

dr_environment/recipe/hooks.py

@@ -0,0 +1,57 @@
+#
+# Copyright 2026 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Invoke component Taskfile environment hooks."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+from pathlib import Path
+
+from dr_environment.recipe.discover import find_taskfile
+from dr_environment.recipe.models import Component
+
+
+def run_environment_hook(component: Component, docker_context: Path) -> None:
+    """Run `task environment` with the documented env contract."""
+    taskfile = find_taskfile(component.source_dir)
+    if taskfile is None:
+        raise FileNotFoundError(f"No Taskfile for component {component.name}")
+
+    fragment_path = (
+        docker_context
+        / "dockerfile.d"
+        / f"{component.fragment_order:02d}-cache-{component.name}.fragment"
+    )
+    fragment_path.parent.mkdir(parents=True, exist_ok=True)
+    if not fragment_path.exists():
+        fragment_path.write_text("", encoding="utf-8")
+
+    env = os.environ.copy()
+    env.update(
+        {
+            "DOCKER_CONTEXT": str(docker_context.resolve()),
+            "COMPONENT_DIR": str(component.source_dir.resolve()),
+            "COMPONENT_NAME": component.name,
+            "DOCKERFILE_FRAGMENT": str(fragment_path.resolve()),
+            "COMPONENT_DEST": str((docker_context / "components" / component.name).resolve()),
+        }
+    )
+
+    subprocess.run(
+        ["task", "-d", str(component.source_dir), "environment"],
+        env=env,
+        check=True,
+    )

dr_environment/recipe/layout.py

@@ -0,0 +1,109 @@
+#
+# Copyright 2026 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Copy component manifests into docker context."""
+
+from __future__ import annotations
+
+import re
+import shutil
+from pathlib import Path
+
+from dr_environment.recipe.models import Component, ComponentStrategy, Ecosystem
+
+# Recipe convention: sibling `core/` is a symlinked local shared Python package.
+LOCAL_SHARED_PACKAGE = "core"
+
+COPY_MAP: dict[Ecosystem, tuple[str, ...]] = {
+    Ecosystem.PYTHON: ("pyproject.toml", "uv.lock"),
+    Ecosystem.NPM: ("package.json", "package-lock.json"),
+    Ecosystem.GO: ("go.mod", "go.sum"),
+}
+
+
+def strip_local_shared_python_package(pyproject_text: str) -> str:
+    """Remove the symlinked local `core` package from a copied pyproject.toml."""
+    text = re.sub(
+        rf'^{LOCAL_SHARED_PACKAGE} = \{{ path = "{LOCAL_SHARED_PACKAGE}"[^\n]*\n',
+        "",
+        pyproject_text,
+        flags=re.MULTILINE,
+    )
+
+    def _clean_inline_deps(match: re.Match[str]) -> str:
+        inner = match.group(1)
+        parts = [part.strip().strip('"') for part in inner.split(",") if part.strip()]
+        kept = [part for part in parts if part != LOCAL_SHARED_PACKAGE]
+        quoted = ", ".join(f'"{part}"' for part in kept)
+        return f"dependencies = [{quoted}]"
+
+    text = re.sub(
+        r"^dependencies = \[(.*?)\]\s*$",
+        _clean_inline_deps,
+        text,
+        flags=re.MULTILINE,
+    )
+
+    lines = text.splitlines()
+    out: list[str] = []
+    in_dependencies = False
+
+    for line in lines:
+        stripped = line.strip()
+        if stripped == "dependencies = [":
+            in_dependencies = True
+        elif in_dependencies and stripped == "]":
+            in_dependencies = False
+        elif in_dependencies and re.fullmatch(r'"core",?', stripped):
+            continue
+        out.append(line)
+
+    trailing_newline = "\n" if pyproject_text.endswith("\n") else ""
+    return "\n".join(out) + trailing_newline
+
+
+def copy_component(component: Component, docker_context: Path) -> Path:
+    """Copy manifest files for a component into components/<name>/."""
+    dest = docker_context / "components" / component.dest_name
+    dest.mkdir(parents=True, exist_ok=True)
+
+    for info in component.manifests:
+        manifest_dir = info.manifest.parent
+        for filename in COPY_MAP[info.ecosystem]:
+            src = manifest_dir / filename
+            if not src.is_file():
+                continue
+            dest_file = dest / filename
+            if (
+                info.ecosystem == Ecosystem.PYTHON
+                and filename == "pyproject.toml"
+                and component.name != LOCAL_SHARED_PACKAGE
+            ):
+                dest_file.write_text(
+                    strip_local_shared_python_package(src.read_text(encoding="utf-8")),
+                    encoding="utf-8",
+                )
+            else:
+                shutil.copy2(src, dest_file)
+
+    return dest
+
+
+def layout_components(components: list[Component], docker_context: Path) -> None:
+    for component in components:
+        if component.strategy != ComponentStrategy.DEFAULT:
+            continue
+        if not component.manifests:
+            continue
+        copy_component(component, docker_context)