Skip to content

Commit

Permalink
Support passing through all modules (#737)
Browse files Browse the repository at this point in the history
Fixes #691
  • Loading branch information
cretz authored Jan 23, 2025
1 parent 1d89d75 commit 4892714
Show file tree
Hide file tree
Showing 4 changed files with 51 additions and 2 deletions.
15 changes: 15 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -1029,6 +1029,21 @@ my_worker = Worker(
In both of these cases, now the `pydantic` module will be passed through from outside of the sandbox instead of
being reloaded for every workflow run.

If users are sure that no imports they use in workflow files will ever need to be sandboxed (meaning all calls within
are deterministic and never mutate shared, global state), the `passthrough_all_modules` option can be set on the
restrictions or the `with_passthrough_all_modules` helper can by used, for example:

```python
my_worker = Worker(
...,
workflow_runner=SandboxedWorkflowRunner(
restrictions=SandboxRestrictions.default.with_passthrough_all_modules()
)
)
```

Note, some calls from the module may still be checked for invalid calls at runtime for certain builtins.

###### Invalid Module Members

`SandboxRestrictions.invalid_module_members` contains a root matcher that applies to all module members. This already
Expand Down
5 changes: 3 additions & 2 deletions temporalio/worker/workflow_sandbox/_importer.py
Original file line number Diff line number Diff line change
Expand Up @@ -252,10 +252,11 @@ def _assert_valid_module(self, name: str) -> None:
raise RestrictedWorkflowAccessError(name)

def _maybe_passthrough_module(self, name: str) -> Optional[types.ModuleType]:
# If imports not passed through and name not in passthrough modules,
# check parents
# If imports not passed through and all modules are not passed through
# and name not in passthrough modules, check parents
if (
not temporalio.workflow.unsafe.is_imports_passed_through()
and not self.restrictions.passthrough_all_modules
and name not in self.restrictions.passthrough_modules
):
end_dot = -1
Expand Down
17 changes: 17 additions & 0 deletions temporalio/worker/workflow_sandbox/_restrictions.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,17 @@ class methods (including __init__, etc). The check compares the against the
fully qualified path to the item.
"""

passthrough_all_modules: bool = False
"""
Pass through all modules, do not sandbox any modules. This is the equivalent
of setting :py:attr:`passthrough_modules` to a list of all modules imported
by the workflow. This is unsafe. This means modules are never reloaded per
workflow run which means workflow authors have to be careful that they don't
import modules that do non-deterministic things. Note, just because a module
is passed through from outside the sandbox doesn't mean runtime restrictions
on invalid calls are not still applied.
"""

passthrough_modules_minimum: ClassVar[Set[str]]
"""Set of modules that must be passed through at the minimum."""

Expand Down Expand Up @@ -133,6 +144,12 @@ def with_passthrough_modules(self, *modules: str) -> SandboxRestrictions:
self, passthrough_modules=self.passthrough_modules | set(modules)
)

def with_passthrough_all_modules(self) -> SandboxRestrictions:
"""Create a new restriction set with :py:attr:`passthrough_all_modules`
as true.
"""
return dataclasses.replace(self, passthrough_all_modules=True)


# We intentionally use specific fields instead of generic "matcher" callbacks
# for optimization reasons.
Expand Down
16 changes: 16 additions & 0 deletions tests/worker/workflow_sandbox/test_importer.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,22 @@ def test_workflow_sandbox_importer_passthough_context_manager():
assert id(outside) == id(inside)


def test_workflow_sandbox_importer_passthrough_all_modules():
import tests.worker.workflow_sandbox.testmodules.stateful_module as outside

# Confirm regular restrictions does re-import
with Importer(restrictions, RestrictionContext()).applied():
import tests.worker.workflow_sandbox.testmodules.stateful_module as inside1
assert id(outside) != id(inside1)

# But that one with all modules passed through does not
with Importer(
restrictions.with_passthrough_all_modules(), RestrictionContext()
).applied():
import tests.worker.workflow_sandbox.testmodules.stateful_module as inside2
assert id(outside) == id(inside2)


def test_workflow_sandbox_importer_invalid_module_members():
importer = Importer(restrictions, RestrictionContext())
# Can access the function, no problem
Expand Down

0 comments on commit 4892714

Please sign in to comment.