diff --git a/docs/code_reference/run_config.md b/docs/code_reference/run_config.md index 2b68f266c..ae358d5e0 100644 --- a/docs/code_reference/run_config.md +++ b/docs/code_reference/run_config.md @@ -1,7 +1,14 @@ # Run Config The `run_config` module defines runtime settings that control dataset generation behavior, -including early shutdown thresholds, batch sizing, and non-inference worker concurrency. +including early shutdown thresholds, batch sizing, non-inference worker concurrency, +and the Jinja rendering engine used by the runtime. + +`JinjaRenderingEngine.SECURE` is the default. Set `JinjaRenderingEngine.NATIVE` +when you want Jinja2's broader built-in sandbox behavior instead of Data Designer's +hardened renderer. + +For guidance on when to use each mode, see [Security](../concepts/security.md). ## Usage @@ -13,6 +20,7 @@ data_designer = DataDesigner() data_designer.set_run_config(dd.RunConfig( buffer_size=500, max_conversation_restarts=3, + jinja_rendering_engine=dd.JinjaRenderingEngine.NATIVE, )) ``` diff --git a/docs/concepts/deployment-options.md b/docs/concepts/deployment-options.md index a77af1070..ca7278ffa 100644 --- a/docs/concepts/deployment-options.md +++ b/docs/concepts/deployment-options.md @@ -141,6 +141,8 @@ If you need to provide synthetic data generation as a shared service: - **Job management**: Queue, monitor, and manage generation jobs centrally - **Resource sharing**: Shared infrastructure for SDG workloads +When users can submit configs containing Jinja templates to a shared engine, template rendering becomes a remote code execution concern and part of your security boundary. See [Security](security.md) for guidance on when to keep the default `JinjaRenderingEngine.SECURE` mode. + --- ## 🧭 Decision Flowchart @@ -181,3 +183,4 @@ If you need to provide synthetic data generation as a shared service: - **Library**: Continue with this documentation - **Microservice**: See the [NeMo Data Designer Microservice documentation](https://docs.nvidia.com/nemo/microservices/latest/design-synthetic-data-from-scratch-or-seeds/index.html){target="_blank"} +- **Security model**: See [Security](security.md) diff --git a/docs/concepts/security.md b/docs/concepts/security.md new file mode 100644 index 000000000..1b98bd1a9 --- /dev/null +++ b/docs/concepts/security.md @@ -0,0 +1,203 @@ +# Security + +Data Designer can run in two very different trust models: + +- **Trusted / monolithic**: The same user or team writes the config and runs the engine. +- **Untrusted / shared execution**: One user submits a config and a different process, service, or team executes it. + +That distinction matters for features that evaluate user-supplied configuration at runtime, such as Jinja template rendering. In a trusted local workflow, broader template flexibility may be acceptable. In a shared-service deployment, user-supplied Jinja becomes part of the engine's remote code execution surface. A template sandbox escape would execute inside the process running Data Designer. + +See [Deployment Options](deployment-options.md) for the architectures where that trust boundary changes. + +## Jinja Rendering Modes + +Data Designer exposes the renderer choice through `RunConfig`: + +```python +import data_designer.config as dd + +run_config = dd.RunConfig( + jinja_rendering_engine=dd.JinjaRenderingEngine.SECURE, +) +``` + +`SECURE` is the default. Opt into `NATIVE` only when you are comfortable treating the config author and the engine operator as the same trust domain. + +| Mode | What it uses | Best fit | +|------|---------------|----------| +| `SECURE` | Data Designer's hardened renderer built on top of Jinja2's sandbox | Shared services, microservices, internal platforms, or any deployment where config submission is separated from execution | +| `NATIVE` | Jinja2's built-in sandbox with Data Designer's variable whitelist | Local library usage and other trusted, monolithic workflows that want broader Jinja behavior | + +!!! warning "Treat untrusted Jinja as a security boundary" + If many users can submit configs to one engine, or if configs are accepted over an API and executed elsewhere, keep `JinjaRenderingEngine.SECURE`. In that model, Jinja templates are no longer just prompt-formatting helpers. They are untrusted user programs being evaluated by your engine. + +## Compatibility Matrix + +`NATIVE` is not an unrestricted Python template engine. The matrix below shows what each mode permits, restricts, or adds on top of Jinja2's standard sandbox behavior. + +| Capability | `NATIVE` | `SECURE` | +|------|------|----------| +| Jinja2 `ImmutableSandboxedEnvironment` baseline | Yes | Yes | +| References to explicitly provided dataset variables only | Yes | Yes | +| Standard Jinja built-in filter set | Yes | Subset only | +| Data Designer `jsonpath` filter | Yes | Yes | +| `import`, `macro`, `set`, `extends`, `block` support | Yes | No | +| Nested or recursive `for` loops | Yes | No | +| Unbounded AST complexity | Yes | No | +| Template context sanitized to JSON-compatible types before render | No | Yes | +| Empty, oversized, or built-in-like rendered output is permitted | Yes | No | + +## What `SECURE` Adds on Top of Standard Jinja Sandbox + +The `SECURE` renderer uses a hardened environment implemented in the [renderer source file on GitHub](https://github.com/NVIDIA-NeMo/DataDesigner/blob/v0.5.6/packages/data-designer-engine/src/data_designer/engine/processing/ginja/environment.py). Compared with the standard Jinja sandbox, it adds several additional controls. + +### Record Sanitization Before Render + +Before rendering, `SECURE` forces template context through a JSON-compatible serialization step. That means remote templates operate on plain data, not arbitrary Python objects. + +```python +# Intended shape for remote template context +record = { + "user": { + "name": "alice", + "roles": ["admin", "reviewer"], + } +} +``` + +```python +# Not the kind of server-side object SECURE wants to expose directly +record = { + "user": SomePythonObject(...), +} +``` + +In a remote execution setting, exposing rich Python objects increases the risk of attribute- and method-based sandbox escapes. Jinja's [sandbox security considerations](https://jinja.palletsprojects.com/en/stable/sandbox/) note that the sandbox is not a complete security boundary, and past escapes have included [`str.format` (CVE-2016-10745)](https://nvd.nist.gov/vuln/detail/CVE-2016-10745), [`str.format_map` (CVE-2019-10906)](https://github.com/advisories/GHSA-462w-v97r-4m45), [indirect `str.format` references (CVE-2024-56326)](https://nvd.nist.gov/vuln/detail/CVE-2024-56326), and [`|attr`-based access to `format` (CVE-2025-27516)](https://nvd.nist.gov/vuln/detail/CVE-2025-27516); PortSwigger's [server-side template injection research](https://portswigger.net/research/server-side-template-injection) covers the broader object-traversal pattern. + +### Filter Allowlist + +`SECURE` keeps only a small approved subset of Jinja filters plus the Data Designer `jsonpath` filter. If a filter is not on that allowlist, the template is rejected. Common excluded filters are: + +| Disallowed filters | Why they are excluded in `SECURE` | +| --- | --- | +| `attr`, `xmlattr` | These add dynamic attribute lookup or attribute-name construction, which widens the object-traversal surface in untrusted templates. | +| `map`, `select`, `reject`, `selectattr`, `rejectattr`, `groupby`, `batch`, `slice`, `sum` | These make templates behave more like a data-processing language and can multiply compute across large inputs. | +| `join`, `format`, `indent`, `wordwrap`, `center`, `filesizeformat` | These expand presentation and composition logic inside the template. `SECURE` keeps formatting logic narrow so templates stay close to interpolation. | +| `default`, `d`, `dictsort`, `count`, `wordcount`, `pprint`, `tojson` | These encourage fallback logic, secondary data shaping, or debug-style output inside the template rather than in the engine or config layer. | +| `safe`, `striptags`, `urlize` | These are primarily HTML-oriented output transforms and are unnecessary for server-side dataset rendering. | + +Some omitted convenience filters, such as the `e` alias for `escape`, are excluded because `SECURE` uses a small explicit allowlist. The current implementation does not assign each omitted filter its own separate security rationale. + +Use `NATIVE` when full Jinja filter compatibility matters more than the additional restrictions used for untrusted template execution. + +### Template Features Removed + +`SECURE` rejects `import`, `macro`, `set`, `extends`, and `block`. + +```jinja +{% macro render_name(name) %}{{ name }}{% endmacro %} +{{ render_name(customer_name) }} +``` + +```jinja +{% set temp = user_id %} +{{ temp }} +``` + +Those features are useful in trusted authoring environments, but they also make user templates more expressive and stateful. In a remote execution model, `SECURE` intentionally narrows the language so templates stay closer to data interpolation than to a reusable programming layer. + +### Loop Restrictions + +`SECURE` rejects recursive loops and nested `for` loops. + +```jinja +{% for row in rows %} + {% for item in row %} + {{ item }} + {% endfor %} +{% endfor %} +``` + +Nested and recursive loops are especially risky in shared execution because they can amplify compute cost and output size in ways that are hard to reason about from the outside. + +### AST Complexity Limits + +`SECURE` statically analyzes the parsed Jinja AST and rejects templates that exceed the current limits of 600 nodes or depth 10. + +```jinja +{% if a %} + {% if b %} + {% if c %} + {{ value }} + {% endif %} + {% endif %} +{% endif %} +``` + +This is not about any one feature being unsafe by itself. It is about limiting how much control flow and composition untrusted templates can pack into a single server-side render operation, which helps prevent compute bombs in shared execution. + +### `self` References Blocked + +`SECURE` rejects references to `self`. + +```jinja +{{ self }} +``` + +The point is to avoid exposing template internals back to the submitter. In a remote setting, even accidental access to those internals is unnecessary surface area. + +### Rendered Output Guards + +`SECURE` validates rendered output after template execution. It rejects empty output, very large output, and strings that look like Python built-in or function representations. + +```jinja +{{ "" }} +``` + +```text + + +``` + +These checks matter because not all bad outcomes come from parse-time behavior. Some templates are syntactically valid but still produce output that is clearly broken, oversized, or revealing internal implementation details. + +### Sanitized User-Facing Errors + +At the engine boundary, `SECURE` normalizes most template failures into a generic invalid-template message. + +```text +User provided prompt generation template is invalid. +``` + +That matters in remote execution because exception details can leak information about server-side implementation, supported objects, or internal execution paths that untrusted users do not need to see. + +These controls exist because the standard sandbox is a good baseline, but shared-service deployments need a narrower and more defensive execution model. + +## Why This Matters in Multi-User Deployments + +The security posture changes as soon as config submission and execution are separated. + +Examples: + +- A centralized Data Designer service accepts configs from many users. +- An internal platform lets users upload or edit configs that are executed by a background worker. +- A REST API accepts Jinja-containing configs and runs them on server-side infrastructure. + +In those environments, templates are no longer just local convenience syntax. They are untrusted input being evaluated by infrastructure the submitter does not control. In practice, that makes Jinja rendering a remote code execution concern, which is why `SECURE` exists and why it remains the default. + +If you are deciding between local library usage and a shared service model, read [Deployment Options](deployment-options.md). The library patterns are often still "trusted" deployments. The shared microservice pattern is not. + +## When To Use `NATIVE` + +Use `NATIVE` when all of the following are true: + +- The person submitting the config is also the person running the engine, or they are in the same trusted operational boundary. +- You want broader standard Jinja behavior than `SECURE` allows. +- You understand that this is a flexibility tradeoff, not the safer default. + +For example, this is often reasonable in a notebook, local script, or other single-user library workflow. + +## Related Reading + +- [Deployment Options](deployment-options.md) +- [Run Config Reference](../code_reference/run_config.md) diff --git a/mkdocs.yml b/mkdocs.yml index cb20ed4e7..464491d4f 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -31,6 +31,7 @@ nav: - Safety & Limits: concepts/mcp/safety-and-limits.md - Architecture & Performance: concepts/architecture-and-performance.md - Deployment Options: concepts/deployment-options.md + - Security: concepts/security.md - Tutorials: - Overview: notebooks/README.md - The Basics: notebooks/1-the-basics.ipynb diff --git a/packages/data-designer-config/src/data_designer/config/__init__.py b/packages/data-designer-config/src/data_designer/config/__init__.py index ed3336259..eb385e15a 100644 --- a/packages/data-designer-config/src/data_designer/config/__init__.py +++ b/packages/data-designer-config/src/data_designer/config/__init__.py @@ -58,7 +58,7 @@ ProcessorType, SchemaTransformProcessorConfig, ) - from data_designer.config.run_config import RunConfig, ThrottleConfig # noqa: F401 + from data_designer.config.run_config import JinjaRenderingEngine, RunConfig, ThrottleConfig # noqa: F401 from data_designer.config.sampler_constraints import ( # noqa: F401 ColumnInequalityConstraint, ConstraintType, @@ -175,6 +175,7 @@ "ProcessorType": (_MOD_PROCESSORS, "ProcessorType"), "SchemaTransformProcessorConfig": (_MOD_PROCESSORS, "SchemaTransformProcessorConfig"), # run_config + "JinjaRenderingEngine": (f"{_MOD_BASE}.run_config", "JinjaRenderingEngine"), "RunConfig": (f"{_MOD_BASE}.run_config", "RunConfig"), "ThrottleConfig": (f"{_MOD_BASE}.run_config", "ThrottleConfig"), # sampler_constraints diff --git a/packages/data-designer-config/src/data_designer/config/run_config.py b/packages/data-designer-config/src/data_designer/config/run_config.py index c2e4269ec..d5f10c9e5 100644 --- a/packages/data-designer-config/src/data_designer/config/run_config.py +++ b/packages/data-designer-config/src/data_designer/config/run_config.py @@ -9,6 +9,14 @@ from typing_extensions import Self from data_designer.config.base import ConfigBase +from data_designer.config.utils.type_helpers import StrEnum + + +class JinjaRenderingEngine(StrEnum): + """Template renderer used by the engine for user-supplied Jinja templates.""" + + NATIVE = "native" + SECURE = "secure" class ThrottleConfig(ConfigBase): @@ -99,6 +107,11 @@ class RunConfig(ConfigBase): Default is False. progress_interval: How often (in seconds) the async progress reporter emits a consolidated log block. Must be > 0. Default is 5.0. + jinja_rendering_engine: Template renderer used for engine-side Jinja evaluation. + ``native`` uses Jinja2's built-in sandbox with the standard filter set and + fewer Data Designer-specific restrictions. ``secure`` uses Data Designer's + hardened sandbox with additional AST, filter, and output guards. + Default is ``secure``. throttle: AIMD throttle tuning parameters. See ``ThrottleConfig`` for details. """ @@ -112,6 +125,13 @@ class RunConfig(ConfigBase): async_trace: bool = False progress_bar: bool = False progress_interval: float = Field(default=5.0, gt=0.0) + jinja_rendering_engine: JinjaRenderingEngine = Field( + default=JinjaRenderingEngine.SECURE, + description=( + "Template renderer used for engine-side Jinja evaluation. " + "`native` uses Jinja2's built-in sandbox; `secure` uses Data Designer's hardened sandbox." + ), + ) throttle: ThrottleConfig = Field(default_factory=ThrottleConfig) @model_validator(mode="after") diff --git a/packages/data-designer-config/tests/config/test_run_config.py b/packages/data-designer-config/tests/config/test_run_config.py new file mode 100644 index 000000000..98c819b38 --- /dev/null +++ b/packages/data-designer-config/tests/config/test_run_config.py @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from data_designer.config.run_config import JinjaRenderingEngine, RunConfig + + +def test_run_config_defaults_to_secure_jinja_renderer() -> None: + assert JinjaRenderingEngine(RunConfig().jinja_rendering_engine) == JinjaRenderingEngine.SECURE + + +def test_run_config_accepts_native_renderer() -> None: + run_config = RunConfig(jinja_rendering_engine=JinjaRenderingEngine.NATIVE) + assert JinjaRenderingEngine(run_config.jinja_rendering_engine) == JinjaRenderingEngine.NATIVE diff --git a/packages/data-designer-engine/src/data_designer/engine/column_generators/generators/llm_completion.py b/packages/data-designer-engine/src/data_designer/engine/column_generators/generators/llm_completion.py index f6625f099..f0dbec4d9 100644 --- a/packages/data-designer-engine/src/data_designer/engine/column_generators/generators/llm_completion.py +++ b/packages/data-designer-engine/src/data_designer/engine/column_generators/generators/llm_completion.py @@ -57,6 +57,7 @@ def prompt_renderer(self) -> RecordBasedPromptRenderer: "column_type": self.config.column_type, "model_alias": self.config.model_alias, }, + jinja_rendering_engine=self.resource_provider.run_config.jinja_rendering_engine, ) def generate(self, data: dict) -> dict: diff --git a/packages/data-designer-engine/src/data_designer/engine/column_generators/generators/samplers.py b/packages/data-designer-engine/src/data_designer/engine/column_generators/generators/samplers.py index 7f327239f..08c0170fc 100644 --- a/packages/data-designer-engine/src/data_designer/engine/column_generators/generators/samplers.py +++ b/packages/data-designer-engine/src/data_designer/engine/column_generators/generators/samplers.py @@ -56,6 +56,7 @@ def _create_sampling_dataset_generator(self) -> SamplingDatasetGenerator: return SamplingDatasetGenerator( sampler_columns=self.config, person_generator_loader=(self._person_generator_loader if self._needs_person_generator else None), + jinja_rendering_engine=self.resource_provider.run_config.jinja_rendering_engine, ) def _log_person_generation_if_needed(self) -> None: diff --git a/packages/data-designer-engine/src/data_designer/engine/column_generators/utils/prompt_renderer.py b/packages/data-designer-engine/src/data_designer/engine/column_generators/utils/prompt_renderer.py index 2b19a6b3f..8a3338d44 100644 --- a/packages/data-designer-engine/src/data_designer/engine/column_generators/utils/prompt_renderer.py +++ b/packages/data-designer-engine/src/data_designer/engine/column_generators/utils/prompt_renderer.py @@ -9,6 +9,7 @@ from data_designer.config.base import SingleColumnConfig from data_designer.config.column_types import DataDesignerColumnType from data_designer.config.models import ModelConfig +from data_designer.config.run_config import JinjaRenderingEngine from data_designer.config.utils.code_lang import CodeLang from data_designer.config.utils.misc import extract_keywords_from_jinja2_template from data_designer.config.utils.type_helpers import StrEnum @@ -36,9 +37,16 @@ class PromptType(StrEnum): class RecordBasedPromptRenderer(WithJinja2UserTemplateRendering): - def __init__(self, response_recipe: ResponseRecipe, *, error_message_context: dict[str, str] | None = None): + def __init__( + self, + response_recipe: ResponseRecipe, + *, + error_message_context: dict[str, str] | None = None, + jinja_rendering_engine: JinjaRenderingEngine = JinjaRenderingEngine.SECURE, + ): self.response_recipe = response_recipe self._error_message_context = error_message_context + self._jinja_rendering_engine = jinja_rendering_engine def render(self, *, prompt_template: str | None, record: dict, prompt_type: PromptType) -> str | None: self._prepare_environment(prompt_template=prompt_template, record=record, prompt_type=prompt_type) diff --git a/packages/data-designer-engine/src/data_designer/engine/processing/ginja/environment.py b/packages/data-designer-engine/src/data_designer/engine/processing/ginja/environment.py index 98470d67d..072e09dc3 100644 --- a/packages/data-designer-engine/src/data_designer/engine/processing/ginja/environment.py +++ b/packages/data-designer-engine/src/data_designer/engine/processing/ginja/environment.py @@ -15,6 +15,7 @@ from jinja2.sandbox import ImmutableSandboxedEnvironment from jsonpath_rust_bindings import Finder +from data_designer.config.run_config import JinjaRenderingEngine from data_designer.engine.processing.ginja.ast import ( ast_count_name_references, ast_descendant_count, @@ -56,6 +57,7 @@ "trim", "truncate", "unique", + "upper", "urlencode", ## Custom Filters "jsonpath", @@ -278,9 +280,11 @@ def validate_template(self, user_template: str) -> None: self._assert_template_ast_complexity(ast) self._assert_template_has_no_self_reference(ast) self._assert_template_has_valid_references(ast) + except UserTemplateError: + raise except Exception as exception: maybe_handle_missing_filter_exception(exception, available_jinja_filters=list(self.filters.keys())) - raise exception + raise def _assert_rendered_text_length(self, rendered_text: str) -> None: """Check against the length of the rendered string.""" @@ -369,6 +373,18 @@ def safe_render( return rendered_text + def render_template( + self, + user_template: str, + record: dict, + skip_template_validation: bool = False, + ) -> str: + return self.safe_render( + user_template, + record, + skip_template_validation=skip_template_validation, + ) + def get_references(self, user_template: str) -> set[str]: """Get all referenced variables from the provided template. @@ -384,6 +400,60 @@ def get_references(self, user_template: str) -> set[str]: return meta.find_undeclared_variables(ast) +class NativeJinjaSandboxEnvironment(ImmutableSandboxedEnvironment): + """Jinja2's built-in sandbox with Data Designer's reference whitelist.""" + + allowed_references: list[str] + _prefer_dict_key_access: bool + + def __init__( + self, + allowed_references: list[str] | None = None, + prefer_dict_key_access: bool = False, + **kwargs, + ): + super().__init__(autoescape=False, **kwargs) + self.allowed_references = allowed_references if allowed_references else [] + self._prefer_dict_key_access = prefer_dict_key_access + self.filters["jsonpath"] = jsonpath_jinja_filter + + def getattr(self, obj: Any, attribute: str) -> Any: + if self._prefer_dict_key_access and isinstance(obj, dict) and attribute in obj: + return obj[attribute] + return super().getattr(obj, attribute) + + def validate_template(self, user_template: str) -> None: + try: + ast = self.parse(user_template) + template_vars = meta.find_undeclared_variables(ast) + unallowed_vars = set(template_vars) - set(self.allowed_references) + if len(unallowed_vars) > 0: + raise UserTemplateError(f"Unknown variable references in Jinja template: {unallowed_vars}") + except UserTemplateError: + raise + except Exception as exception: + maybe_handle_missing_filter_exception(exception, available_jinja_filters=list(self.filters.keys())) + raise + + def render_template( + self, + user_template: str, + record: dict, + skip_template_validation: bool = False, + ) -> str: + if not skip_template_validation: + self.validate_template(user_template) + + try: + template = self.from_string(user_template) + return template.render(record) + except SecurityError as exception: + raise UserTemplateError("Non-permitted operations in Jinja template.") from exception + except Exception as exception: + maybe_handle_missing_filter_exception(exception, available_jinja_filters=list(self.filters.keys())) + raise UserTemplateError(str(exception)) from exception + + def sanitize_user_exceptions(func): """Sanitize returned user-space exceptions.""" @@ -428,6 +498,33 @@ def my_func(self, user_template: str, records: list[dict]): _template_render_fn: Callable + def _get_jinja_rendering_engine(self) -> JinjaRenderingEngine: + engine = getattr(self, "_jinja_rendering_engine", None) + if engine is not None: + return JinjaRenderingEngine(engine) + + resource_provider = getattr(self, "_resource_provider", None) + if resource_provider is not None: + return JinjaRenderingEngine(resource_provider.run_config.jinja_rendering_engine) + + # The mixin predates the RunConfig toggle, so preserve the historical + # secure-by-default behavior when no explicit engine is wired in. + return JinjaRenderingEngine.SECURE + + def _create_render_environment( + self, + *, + dataset_variables: list[str], + record_str_fn: Callable[[Any], str] | None = None, + ) -> UserTemplateSandboxEnvironment | NativeJinjaSandboxEnvironment: + env_kwargs: dict[str, Any] = {} + if record_str_fn is not None: + env_kwargs["finalize"] = record_str_fn + env_kwargs["prefer_dict_key_access"] = True + if self._get_jinja_rendering_engine() == JinjaRenderingEngine.SECURE: + return UserTemplateSandboxEnvironment(allowed_references=dataset_variables, **env_kwargs) + return NativeJinjaSandboxEnvironment(allowed_references=dataset_variables, **env_kwargs) + @sanitize_user_exceptions def prepare_jinja2_template_renderer( self, @@ -445,14 +542,13 @@ def prepare_jinja2_template_renderer( and enables dict-key-priority attribute lookup for nested dot access ({{ col.sub.field }}). """ - env_kwargs: dict[str, Any] = {} - if record_str_fn is not None: - env_kwargs["finalize"] = record_str_fn - env_kwargs["prefer_dict_key_access"] = True - jinja_render_env = UserTemplateSandboxEnvironment(allowed_references=dataset_variables, **env_kwargs) + jinja_render_env = self._create_render_environment( + dataset_variables=dataset_variables, + record_str_fn=record_str_fn, + ) jinja_render_env.validate_template(prompt_template) self._template_render_fn = partial( - jinja_render_env.safe_render, + jinja_render_env.render_template, prompt_template, skip_template_validation=True, ) @@ -470,10 +566,10 @@ def prepare_jinja2_multi_template_renderer( ) -> None: if not self._template_prepared_in_multi_template_renderer(template_name): self._create_render_func_registry() - jinja_render_env = UserTemplateSandboxEnvironment(allowed_references=dataset_variables) + jinja_render_env = self._create_render_environment(dataset_variables=dataset_variables) jinja_render_env.validate_template(prompt_template) self._render_func_registry[template_name] = partial( - jinja_render_env.safe_render, + jinja_render_env.render_template, prompt_template, skip_template_validation=True, ) diff --git a/packages/data-designer-engine/src/data_designer/engine/sampling_gen/generator.py b/packages/data-designer-engine/src/data_designer/engine/sampling_gen/generator.py index 18bdeffd3..b316b5ba2 100644 --- a/packages/data-designer-engine/src/data_designer/engine/sampling_gen/generator.py +++ b/packages/data-designer-engine/src/data_designer/engine/sampling_gen/generator.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING import data_designer.lazy_heavy_imports as lazy +from data_designer.config.run_config import JinjaRenderingEngine from data_designer.engine.sampling_gen.data_sources.base import RadomStateT from data_designer.engine.sampling_gen.errors import RejectionSamplingError from data_designer.engine.sampling_gen.jinja_utils import JinjaDataFrame @@ -49,6 +50,7 @@ def __init__( *, schema: DataSchema | None = None, max_rejections_factor: int = 5, + jinja_rendering_engine: JinjaRenderingEngine = JinjaRenderingEngine.SECURE, ): # This is temporary while we need the legacy and refactored code to coexist. if schema is not None: @@ -63,6 +65,7 @@ def __init__( self.rng = check_random_state(random_state) self._dag = self.schema.dag.to_networkx() + self._jinja_rendering_engine = jinja_rendering_engine self._shared_sampler_kwargs = { "random_state": self.rng, "people_gen_resource": create_people_gen_resource(self.schema, person_generator_loader), @@ -81,7 +84,10 @@ def _run_rejection_sampling(self, df: pd.DataFrame, column: ConditionalDataColum while needs_samples.any(): for condition in column.conditions: - index = JinjaDataFrame(condition).select_index(df[needs_samples]) + index = JinjaDataFrame( + condition, + jinja_rendering_engine=self._jinja_rendering_engine, + ).select_index(df[needs_samples]) src = column.get_sampler(condition, **self._shared_sampler_kwargs) df = src.inject_data_column(df, name, index) diff --git a/packages/data-designer-engine/src/data_designer/engine/sampling_gen/jinja_utils.py b/packages/data-designer-engine/src/data_designer/engine/sampling_gen/jinja_utils.py index c10726116..ccf440ca2 100644 --- a/packages/data-designer-engine/src/data_designer/engine/sampling_gen/jinja_utils.py +++ b/packages/data-designer-engine/src/data_designer/engine/sampling_gen/jinja_utils.py @@ -6,9 +6,12 @@ import ast from typing import TYPE_CHECKING, Any +from jinja2 import meta + import data_designer.lazy_heavy_imports as lazy +from data_designer.config.run_config import JinjaRenderingEngine from data_designer.engine.processing.ginja.environment import ( - UserTemplateSandboxEnvironment, + NativeJinjaSandboxEnvironment, WithJinja2UserTemplateRendering, ) @@ -17,8 +20,14 @@ class JinjaDataFrame(WithJinja2UserTemplateRendering): - def __init__(self, expr: str): + def __init__( + self, + expr: str, + *, + jinja_rendering_engine: JinjaRenderingEngine = JinjaRenderingEngine.SECURE, + ): self.expr = expr + self._jinja_rendering_engine = jinja_rendering_engine def _jsonify(self, record) -> dict[str, Any]: for key, value in record.items(): @@ -61,4 +70,4 @@ def to_column(self, dataframe: pd.DataFrame) -> list[Any]: def extract_column_names_from_expression(expr: str) -> set[str]: """Extract valid column names from the given expression.""" - return UserTemplateSandboxEnvironment().get_references("{{ " + expr + " }}") + return meta.find_undeclared_variables(NativeJinjaSandboxEnvironment().parse("{{ " + expr + " }}")) diff --git a/packages/data-designer-engine/tests/engine/column_generators/generators/test_expression.py b/packages/data-designer-engine/tests/engine/column_generators/generators/test_expression.py index 63a8dc6b2..15080368f 100644 --- a/packages/data-designer-engine/tests/engine/column_generators/generators/test_expression.py +++ b/packages/data-designer-engine/tests/engine/column_generators/generators/test_expression.py @@ -9,8 +9,10 @@ import data_designer.lazy_heavy_imports as lazy from data_designer.config.column_configs import ExpressionColumnConfig +from data_designer.config.run_config import JinjaRenderingEngine, RunConfig from data_designer.engine.column_generators.generators.expression import ExpressionColumnGenerator from data_designer.engine.column_generators.utils.errors import ExpressionTemplateRenderError +from data_designer.engine.processing.ginja.exceptions import UserTemplateUnsupportedFiltersError from data_designer.engine.resources.resource_provider import ResourceProvider @@ -160,3 +162,26 @@ def test_generate_with_missing_columns(): match=r"There was an error preparing the Jinja2 expression template. The following columns \['col1'\] are missing!", ): generator.generate(df) + + +def test_generate_respects_run_config_jinja_rendering_engine() -> None: + df = lazy.pd.DataFrame({"col1": [["a", "b"]]}) + + native_provider = Mock(spec=ResourceProvider) + native_provider.run_config = RunConfig(jinja_rendering_engine=JinjaRenderingEngine.NATIVE) + native_generator = _create_test_generator( + _create_test_config("joined", "{{ col1 | join('-') }}", "str"), + native_provider, + ) + native_result = native_generator.generate(df) + assert native_result["joined"].tolist() == ["a-b"] + + secure_provider = Mock(spec=ResourceProvider) + secure_provider.run_config = RunConfig(jinja_rendering_engine=JinjaRenderingEngine.SECURE) + secure_generator = _create_test_generator( + _create_test_config("joined", "{{ col1 | join('-') }}", "str"), + secure_provider, + ) + + with pytest.raises(UserTemplateUnsupportedFiltersError): + secure_generator.generate(df) diff --git a/packages/data-designer-engine/tests/engine/column_generators/generators/test_image.py b/packages/data-designer-engine/tests/engine/column_generators/generators/test_image.py index fe42bb438..e60212f1b 100644 --- a/packages/data-designer-engine/tests/engine/column_generators/generators/test_image.py +++ b/packages/data-designer-engine/tests/engine/column_generators/generators/test_image.py @@ -114,13 +114,13 @@ def test_image_cell_generator_missing_columns_error(stub_image_column_config, st def test_image_cell_generator_empty_prompt_error(stub_resource_provider): - """Test that empty rendered prompt raises UserTemplateError.""" + """Test that empty rendered prompt is rejected by the secure renderer.""" # Create config with template that renders to empty string config = ImageColumnConfig(name="test_image", prompt="{{ empty }}", model_alias="test_model") generator = ImageCellGenerator(config=config, resource_provider=stub_resource_provider) - with pytest.raises(UserTemplateError): + with pytest.raises(UserTemplateError, match="invalid"): generator.generate(data={"empty": ""}) diff --git a/packages/data-designer-engine/tests/engine/column_generators/utils/test_prompt_renderer.py b/packages/data-designer-engine/tests/engine/column_generators/utils/test_prompt_renderer.py index 81f63f849..bdc7e6c8b 100644 --- a/packages/data-designer-engine/tests/engine/column_generators/utils/test_prompt_renderer.py +++ b/packages/data-designer-engine/tests/engine/column_generators/utils/test_prompt_renderer.py @@ -12,7 +12,9 @@ LLMTextColumnConfig, Score, ) +from data_designer.config.run_config import JinjaRenderingEngine from data_designer.config.utils.code_lang import CodeLang +from data_designer.engine.column_generators.utils.errors import PromptTemplateRenderError from data_designer.engine.column_generators.utils.prompt_renderer import ( PromptType, RecordBasedPromptRenderer, @@ -125,3 +127,33 @@ def test_prompt_renderer_render_prompt_template_error(): renderer.render( prompt_template="Test prompt: {{ invalid_template }}", record=data, prompt_type=PromptType.USER_PROMPT ) + + +def test_prompt_renderer_uses_secure_jinja_by_default() -> None: + config = LLMTextColumnConfig(name="test_column", prompt="Test prompt", model_alias="test_model") + recipe = create_response_recipe(config) + renderer = RecordBasedPromptRenderer(response_recipe=recipe) + + with pytest.raises(PromptTemplateRenderError, match=r"\| join"): + renderer.render( + prompt_template="Joined: {{ input | join('-') }}", + record={"input": ["Hello", "World"]}, + prompt_type=PromptType.USER_PROMPT, + ) + + +def test_prompt_renderer_can_opt_into_native_mode() -> None: + config = LLMTextColumnConfig(name="test_column", prompt="Test prompt", model_alias="test_model") + recipe = create_response_recipe(config) + renderer = RecordBasedPromptRenderer( + response_recipe=recipe, + jinja_rendering_engine=JinjaRenderingEngine.NATIVE, + ) + + result = renderer.render( + prompt_template="Joined: {{ input | join('-') }}", + record={"input": ["Hello", "World"]}, + prompt_type=PromptType.USER_PROMPT, + ) + + assert result == "Joined: Hello-World" diff --git a/packages/data-designer-engine/tests/engine/processing/ginja/test_environment.py b/packages/data-designer-engine/tests/engine/processing/ginja/test_environment.py index 69eb98206..2c79599ea 100644 --- a/packages/data-designer-engine/tests/engine/processing/ginja/test_environment.py +++ b/packages/data-designer-engine/tests/engine/processing/ginja/test_environment.py @@ -3,8 +3,10 @@ import pytest +from data_designer.config.run_config import JinjaRenderingEngine from data_designer.engine.processing.ginja.environment import ( ALLOWED_JINJA_FILTERS, + NativeJinjaSandboxEnvironment, UserTemplateSandboxEnvironment, WithJinja2UserTemplateRendering, is_jinja_template, @@ -96,6 +98,16 @@ def test_jsonpath_jinja_filter(jsonpath_query, expected_result): assert jsonpath_jinja_filter(TEST_RECORD, jsonpath_query) == expected_result +def test_native_jinja_sandbox_environment_supports_jsonpath_filter() -> None: + env = NativeJinjaSandboxEnvironment(allowed_references=list(TEST_RECORD.keys())) + + assert env.render_template('{{ field_c | jsonpath("$.sub_a.foo[:2]") }}', TEST_RECORD) == str([1, 2]) + + +def test_user_template_sandbox_environment_supports_upper_filter(stub_sandbox_env) -> None: + assert stub_sandbox_env.safe_render("{{ field_y | upper }}", TEST_RECORD) == "FOO" + + @pytest.mark.parametrize( "jinja_template,expected_result", [ @@ -180,6 +192,7 @@ def test_with_jinja2_user_template_rendering_mixin(test_case, template_1, templa class Foo(WithJinja2UserTemplateRendering): def __init__(self, template_1: str, template_2: str = None): + self._jinja_rendering_engine = JinjaRenderingEngine.SECURE if template_2 is None: # Single template self.prepare_jinja2_template_renderer(template_1, dataset_variables=["safe"]) @@ -211,3 +224,13 @@ def bar(self, record): else: with pytest.raises(expected_result): f = Foo(template_1, template_2) + + +def test_with_jinja2_user_template_rendering_defaults_to_secure_mode() -> None: + class Foo(WithJinja2UserTemplateRendering): + pass + + renderer = Foo() + + with pytest.raises(UserTemplateUnsupportedFiltersError): + renderer.prepare_jinja2_template_renderer("{{ items | join('-') }}", dataset_variables=["items"]) diff --git a/packages/data-designer-engine/tests/engine/sampling_gen/test_jinja_utils.py b/packages/data-designer-engine/tests/engine/sampling_gen/test_jinja_utils.py index e2a21c35a..bd050f7e1 100644 --- a/packages/data-designer-engine/tests/engine/sampling_gen/test_jinja_utils.py +++ b/packages/data-designer-engine/tests/engine/sampling_gen/test_jinja_utils.py @@ -8,6 +8,8 @@ import pytest import data_designer.lazy_heavy_imports as lazy +from data_designer.config.run_config import JinjaRenderingEngine +from data_designer.engine.processing.ginja.exceptions import UserTemplateUnsupportedFiltersError from data_designer.engine.sampling_gen.jinja_utils import JinjaDataFrame, extract_column_names_from_expression @@ -27,6 +29,7 @@ ("some_dude.age + 1", {"some_dude"}), ("'I\\'m a string' + i_am_a_var", {"i_am_a_var"}), ('"I am a string" + i_am_a_var', {"i_am_a_var"}), + ('data | jsonpath("$.key")', {"data"}), ], ) def test_extract_column_names_from_expression(expr: str, column_names: set[str]) -> None: @@ -113,3 +116,26 @@ def test_jinja_dataframe_to_column_scenarios(test_case, expr, df_data, mock_side jdf.render_template = Mock(side_effect=mock_side_effect) result = jdf.to_column(df) assert result == expected_result + + +def test_jinja_dataframe_can_switch_rendering_engines() -> None: + df = lazy.pd.DataFrame({"items": [["a", "b"]]}) + + with pytest.raises(UserTemplateUnsupportedFiltersError): + JinjaDataFrame( + "items | join('-')", + jinja_rendering_engine=JinjaRenderingEngine.SECURE, + ).to_column(df) + + native_result = JinjaDataFrame( + "items | join('-')", + jinja_rendering_engine=JinjaRenderingEngine.NATIVE, + ).to_column(df) + assert native_result == ["a-b"] + + +def test_jinja_dataframe_uses_secure_jinja_by_default() -> None: + df = lazy.pd.DataFrame({"items": [["a", "b"]]}) + + with pytest.raises(UserTemplateUnsupportedFiltersError): + JinjaDataFrame("items | join('-')").to_column(df) diff --git a/packages/data-designer/src/data_designer/interface/data_designer.py b/packages/data-designer/src/data_designer/interface/data_designer.py index e487074a5..7913bbbee 100644 --- a/packages/data-designer/src/data_designer/interface/data_designer.py +++ b/packages/data-designer/src/data_designer/interface/data_designer.py @@ -25,7 +25,7 @@ ModelProvider, ) from data_designer.config.preview_results import PreviewResults -from data_designer.config.run_config import RunConfig +from data_designer.config.run_config import JinjaRenderingEngine, RunConfig from data_designer.config.utils.constants import ( DEFAULT_NUM_RECORDS, MANAGED_ASSETS_PATH, @@ -65,7 +65,7 @@ DataDesignerProfilingError, ) from data_designer.interface.results import DatasetCreationResults -from data_designer.logging import RandomEmoji, configure_logging +from data_designer.logging import LOG_INDENT, RandomEmoji, configure_logging from data_designer.plugins.plugin import PluginType from data_designer.plugins.registry import PluginRegistry @@ -217,6 +217,7 @@ def create( DataDesignerProfilingError: If an error occurs during dataset profiling. """ logger.info("🎨 Creating Data Designer dataset") + self._log_jinja_rendering_engine_mode() resource_provider = self._create_resource_provider(dataset_name, config_builder) @@ -288,6 +289,7 @@ def preview( DataDesignerProfilingError: If an error occurs during preview dataset profiling. """ logger.info(f"{RandomEmoji.previewing()} Preview generation in progress") + self._log_jinja_rendering_engine_mode() resource_provider = self._create_resource_provider("preview-dataset", config_builder) try: @@ -333,6 +335,11 @@ def preview( dataset_metadata=dataset_metadata, ) + def _log_jinja_rendering_engine_mode(self) -> None: + engine = JinjaRenderingEngine(self._run_config.jinja_rendering_engine) + icon = "🔒" if engine == JinjaRenderingEngine.SECURE else "🏠" + logger.info(f"{LOG_INDENT}{icon} Jinja rendering engine: {engine.value}") + def validate(self, config_builder: DataDesignerConfigBuilder) -> None: """Validate the Data Designer configuration as defined by the DataDesignerConfigBuilder with the configured engine components (SecretResolver, SeedReaders, etc.). diff --git a/packages/data-designer/tests/interface/test_data_designer.py b/packages/data-designer/tests/interface/test_data_designer.py index 81db7722d..dc56b1a74 100644 --- a/packages/data-designer/tests/interface/test_data_designer.py +++ b/packages/data-designer/tests/interface/test_data_designer.py @@ -20,7 +20,7 @@ from data_designer.config.errors import InvalidConfigError from data_designer.config.models import ModelProvider from data_designer.config.processors import DropColumnsProcessorConfig -from data_designer.config.run_config import RunConfig +from data_designer.config.run_config import JinjaRenderingEngine, RunConfig from data_designer.config.sampler_params import CategorySamplerParams, DatetimeSamplerParams, SamplerType from data_designer.config.seed import IndexRange, PartitionBlock, SamplingStrategy from data_designer.config.seed_source import ( @@ -703,6 +703,86 @@ def test_preview_raises_generation_error_when_dataset_is_empty( data_designer.preview(stub_sampler_only_config_builder, num_records=1) +def test_create_logs_secure_jinja_rendering_mode( + stub_artifact_path: Path, + stub_model_providers: list[ModelProvider], + stub_sampler_only_config_builder: DataDesignerConfigBuilder, + stub_managed_assets_path: Path, +) -> None: + with patch.object(dd_mod, "get_default_provider_name", return_value="stub-model-provider"): + data_designer = DataDesigner( + artifact_path=stub_artifact_path, + model_providers=stub_model_providers, + secret_resolver=PlaintextResolver(), + managed_assets_path=stub_managed_assets_path, + ) + data_designer.set_run_config(RunConfig(jinja_rendering_engine=JinjaRenderingEngine.SECURE)) + + with ( + patch.object(dd_mod.logger, "info") as mock_info, + patch.object(data_designer, "_create_resource_provider") as mock_resource_provider_method, + patch.object(data_designer, "_create_dataset_builder") as mock_builder_method, + patch.object(data_designer, "_create_dataset_profiler") as mock_profiler_method, + ): + mock_resource_provider = MagicMock() + mock_resource_provider.get_dataset_metadata.return_value = {} + mock_resource_provider_method.return_value = mock_resource_provider + + mock_builder = MagicMock() + mock_builder.build.return_value = None + mock_builder.task_traces = [] + mock_builder.artifact_storage.load_dataset_with_dropped_columns.return_value = lazy.pd.DataFrame({"col": [1]}) + mock_builder_method.return_value = mock_builder + + mock_profiler = MagicMock() + mock_profiler.profile_dataset.return_value = None + mock_profiler_method.return_value = mock_profiler + + data_designer.create(stub_sampler_only_config_builder, num_records=1) + + assert any("🔒 Jinja rendering engine: secure" in call.args[0] for call in mock_info.call_args_list) + + +def test_preview_logs_native_jinja_rendering_mode( + stub_artifact_path: Path, + stub_model_providers: list[ModelProvider], + stub_sampler_only_config_builder: DataDesignerConfigBuilder, + stub_managed_assets_path: Path, +) -> None: + with patch.object(dd_mod, "get_default_provider_name", return_value="stub-model-provider"): + data_designer = DataDesigner( + artifact_path=stub_artifact_path, + model_providers=stub_model_providers, + secret_resolver=PlaintextResolver(), + managed_assets_path=stub_managed_assets_path, + ) + data_designer.set_run_config(RunConfig(jinja_rendering_engine=JinjaRenderingEngine.NATIVE)) + + with ( + patch.object(dd_mod.logger, "info") as mock_info, + patch.object(data_designer, "_create_resource_provider") as mock_resource_provider_method, + patch.object(data_designer, "_create_dataset_builder") as mock_builder_method, + patch.object(data_designer, "_create_dataset_profiler") as mock_profiler_method, + ): + mock_resource_provider = MagicMock() + mock_resource_provider.get_dataset_metadata.return_value = {} + mock_resource_provider_method.return_value = mock_resource_provider + + mock_builder = MagicMock() + mock_builder.build_preview.return_value = lazy.pd.DataFrame({"col": [1]}) + mock_builder.process_preview.return_value = lazy.pd.DataFrame({"col": [1]}) + mock_builder.artifact_storage.list_processor_names.return_value = [] + mock_builder_method.return_value = mock_builder + + mock_profiler = MagicMock() + mock_profiler.profile_dataset.return_value = None + mock_profiler_method.return_value = mock_profiler + + data_designer.preview(stub_sampler_only_config_builder, num_records=1) + + assert any("🏠 Jinja rendering engine: native" in call.args[0] for call in mock_info.call_args_list) + + def test_preview_datetime_single_record_returns_iso8601( stub_artifact_path, stub_model_providers, stub_model_configs, stub_managed_assets_path ):