langfun 0.1.2.dev202509120804__py3-none-any.whl → 0.1.2.dev202512040805__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- langfun/__init__.py +1 -1
- langfun/core/__init__.py +7 -1
- langfun/core/agentic/__init__.py +8 -1
- langfun/core/agentic/action.py +740 -112
- langfun/core/agentic/action_eval.py +9 -2
- langfun/core/agentic/action_test.py +189 -24
- langfun/core/async_support.py +104 -5
- langfun/core/async_support_test.py +23 -0
- langfun/core/coding/python/correction.py +19 -9
- langfun/core/coding/python/execution.py +14 -12
- langfun/core/coding/python/generation.py +21 -16
- langfun/core/coding/python/sandboxing.py +23 -3
- langfun/core/component.py +42 -3
- langfun/core/concurrent.py +70 -6
- langfun/core/concurrent_test.py +9 -2
- langfun/core/console.py +1 -1
- langfun/core/data/conversion/anthropic.py +12 -3
- langfun/core/data/conversion/anthropic_test.py +8 -6
- langfun/core/data/conversion/gemini.py +11 -2
- langfun/core/data/conversion/gemini_test.py +48 -9
- langfun/core/data/conversion/openai.py +145 -31
- langfun/core/data/conversion/openai_test.py +161 -17
- langfun/core/eval/base.py +48 -44
- langfun/core/eval/base_test.py +5 -5
- langfun/core/eval/matching.py +5 -2
- langfun/core/eval/patching.py +3 -3
- langfun/core/eval/scoring.py +4 -3
- langfun/core/eval/v2/__init__.py +2 -0
- langfun/core/eval/v2/checkpointing.py +76 -7
- langfun/core/eval/v2/checkpointing_test.py +9 -2
- langfun/core/eval/v2/config_saver.py +37 -0
- langfun/core/eval/v2/config_saver_test.py +36 -0
- langfun/core/eval/v2/eval_test_helper.py +104 -3
- langfun/core/eval/v2/evaluation.py +92 -17
- langfun/core/eval/v2/evaluation_test.py +9 -3
- langfun/core/eval/v2/example.py +50 -40
- langfun/core/eval/v2/example_test.py +16 -8
- langfun/core/eval/v2/experiment.py +84 -15
- langfun/core/eval/v2/experiment_test.py +19 -0
- langfun/core/eval/v2/metric_values.py +31 -3
- langfun/core/eval/v2/metric_values_test.py +32 -0
- langfun/core/eval/v2/metrics.py +157 -44
- langfun/core/eval/v2/metrics_test.py +39 -18
- langfun/core/eval/v2/progress.py +31 -1
- langfun/core/eval/v2/progress_test.py +27 -0
- langfun/core/eval/v2/progress_tracking.py +13 -5
- langfun/core/eval/v2/progress_tracking_test.py +9 -1
- langfun/core/eval/v2/reporting.py +90 -71
- langfun/core/eval/v2/reporting_test.py +24 -6
- langfun/core/eval/v2/runners/__init__.py +30 -0
- langfun/core/eval/v2/{runners.py → runners/base.py} +72 -180
- langfun/core/eval/v2/runners/beam.py +354 -0
- langfun/core/eval/v2/runners/beam_test.py +153 -0
- langfun/core/eval/v2/runners/ckpt_monitor.py +294 -0
- langfun/core/eval/v2/runners/ckpt_monitor_test.py +162 -0
- langfun/core/eval/v2/runners/debug.py +40 -0
- langfun/core/eval/v2/runners/debug_test.py +76 -0
- langfun/core/eval/v2/runners/parallel.py +243 -0
- langfun/core/eval/v2/runners/parallel_test.py +182 -0
- langfun/core/eval/v2/runners/sequential.py +47 -0
- langfun/core/eval/v2/runners/sequential_test.py +169 -0
- langfun/core/langfunc.py +45 -130
- langfun/core/langfunc_test.py +7 -5
- langfun/core/language_model.py +189 -36
- langfun/core/language_model_test.py +54 -3
- langfun/core/llms/__init__.py +12 -1
- langfun/core/llms/anthropic.py +157 -2
- langfun/core/llms/azure_openai.py +29 -17
- langfun/core/llms/cache/base.py +25 -3
- langfun/core/llms/cache/in_memory.py +48 -7
- langfun/core/llms/cache/in_memory_test.py +14 -4
- langfun/core/llms/compositional.py +25 -1
- langfun/core/llms/deepseek.py +30 -2
- langfun/core/llms/fake.py +32 -1
- langfun/core/llms/gemini.py +64 -12
- langfun/core/llms/gemini_test.py +110 -0
- langfun/core/llms/google_genai.py +34 -1
- langfun/core/llms/groq.py +28 -3
- langfun/core/llms/llama_cpp.py +23 -4
- langfun/core/llms/openai.py +120 -3
- langfun/core/llms/openai_compatible.py +148 -27
- langfun/core/llms/openai_compatible_test.py +207 -20
- langfun/core/llms/openai_test.py +0 -2
- langfun/core/llms/rest.py +16 -1
- langfun/core/llms/vertexai.py +58 -8
- langfun/core/logging.py +1 -1
- langfun/core/mcp/__init__.py +10 -0
- langfun/core/mcp/client.py +177 -0
- langfun/core/mcp/client_test.py +71 -0
- langfun/core/mcp/session.py +241 -0
- langfun/core/mcp/session_test.py +54 -0
- langfun/core/mcp/testing/simple_mcp_client.py +33 -0
- langfun/core/mcp/testing/simple_mcp_server.py +33 -0
- langfun/core/mcp/tool.py +254 -0
- langfun/core/mcp/tool_test.py +197 -0
- langfun/core/memory.py +1 -0
- langfun/core/message.py +160 -55
- langfun/core/message_test.py +65 -81
- langfun/core/modalities/__init__.py +8 -0
- langfun/core/modalities/audio.py +21 -1
- langfun/core/modalities/image.py +73 -3
- langfun/core/modalities/image_test.py +116 -0
- langfun/core/modalities/mime.py +64 -3
- langfun/core/modalities/mime_test.py +11 -0
- langfun/core/modalities/pdf.py +19 -1
- langfun/core/modalities/video.py +21 -1
- langfun/core/modality.py +167 -29
- langfun/core/modality_test.py +42 -12
- langfun/core/natural_language.py +1 -1
- langfun/core/sampling.py +4 -4
- langfun/core/sampling_test.py +20 -4
- langfun/core/structured/__init__.py +2 -24
- langfun/core/structured/completion.py +34 -44
- langfun/core/structured/completion_test.py +23 -43
- langfun/core/structured/description.py +54 -50
- langfun/core/structured/function_generation.py +29 -12
- langfun/core/structured/mapping.py +81 -37
- langfun/core/structured/parsing.py +95 -79
- langfun/core/structured/parsing_test.py +0 -3
- langfun/core/structured/querying.py +230 -154
- langfun/core/structured/querying_test.py +69 -33
- langfun/core/structured/schema/__init__.py +49 -0
- langfun/core/structured/schema/base.py +664 -0
- langfun/core/structured/schema/base_test.py +531 -0
- langfun/core/structured/schema/json.py +174 -0
- langfun/core/structured/schema/json_test.py +121 -0
- langfun/core/structured/schema/python.py +316 -0
- langfun/core/structured/schema/python_test.py +410 -0
- langfun/core/structured/schema_generation.py +33 -14
- langfun/core/structured/scoring.py +47 -36
- langfun/core/structured/tokenization.py +26 -11
- langfun/core/subscription.py +2 -2
- langfun/core/template.py +175 -50
- langfun/core/template_test.py +123 -17
- langfun/env/__init__.py +43 -0
- langfun/env/base_environment.py +827 -0
- langfun/env/base_environment_test.py +473 -0
- langfun/env/base_feature.py +304 -0
- langfun/env/base_feature_test.py +228 -0
- langfun/env/base_sandbox.py +842 -0
- langfun/env/base_sandbox_test.py +1235 -0
- langfun/env/event_handlers/__init__.py +14 -0
- langfun/env/event_handlers/chain.py +233 -0
- langfun/env/event_handlers/chain_test.py +253 -0
- langfun/env/event_handlers/event_logger.py +472 -0
- langfun/env/event_handlers/event_logger_test.py +304 -0
- langfun/env/event_handlers/metric_writer.py +726 -0
- langfun/env/event_handlers/metric_writer_test.py +214 -0
- langfun/env/interface.py +1640 -0
- langfun/env/interface_test.py +153 -0
- langfun/env/load_balancers.py +59 -0
- langfun/env/load_balancers_test.py +141 -0
- langfun/env/test_utils.py +507 -0
- {langfun-0.1.2.dev202509120804.dist-info → langfun-0.1.2.dev202512040805.dist-info}/METADATA +7 -3
- langfun-0.1.2.dev202512040805.dist-info/RECORD +217 -0
- langfun/core/eval/v2/runners_test.py +0 -343
- langfun/core/structured/schema.py +0 -987
- langfun/core/structured/schema_test.py +0 -982
- langfun-0.1.2.dev202509120804.dist-info/RECORD +0 -172
- {langfun-0.1.2.dev202509120804.dist-info → langfun-0.1.2.dev202512040805.dist-info}/WHEEL +0 -0
- {langfun-0.1.2.dev202509120804.dist-info → langfun-0.1.2.dev202512040805.dist-info}/licenses/LICENSE +0 -0
- {langfun-0.1.2.dev202509120804.dist-info → langfun-0.1.2.dev202512040805.dist-info}/top_level.txt +0 -0
langfun/env/interface.py
ADDED
|
@@ -0,0 +1,1640 @@
|
|
|
1
|
+
# Copyright 2025 The Langfun Authors
|
|
2
|
+
#
|
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
4
|
+
# you may not use this file except in compliance with the License.
|
|
5
|
+
# You may obtain a copy of the License at
|
|
6
|
+
#
|
|
7
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
8
|
+
#
|
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
12
|
+
# See the License for the specific language governing permissions and
|
|
13
|
+
# limitations under the License.
|
|
14
|
+
"""Interfaces for environments, sandboxes and features."""
|
|
15
|
+
|
|
16
|
+
import abc
|
|
17
|
+
import contextlib
|
|
18
|
+
import dataclasses
|
|
19
|
+
import enum
|
|
20
|
+
import functools
|
|
21
|
+
import os
|
|
22
|
+
from typing import Annotated, Any, Callable, ContextManager, ClassVar, Iterator, Optional, Sequence, Type
|
|
23
|
+
|
|
24
|
+
import pyglove as pg
|
|
25
|
+
|
|
26
|
+
#
|
|
27
|
+
# Environment errors.
|
|
28
|
+
#
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
class EnvironmentError(RuntimeError): # pylint: disable=redefined-builtin
|
|
32
|
+
"""Base class for environment errors."""
|
|
33
|
+
|
|
34
|
+
def __init__(
|
|
35
|
+
self,
|
|
36
|
+
message: str,
|
|
37
|
+
*args,
|
|
38
|
+
environment: 'Environment',
|
|
39
|
+
**kwargs
|
|
40
|
+
) -> None:
|
|
41
|
+
self.environment = environment
|
|
42
|
+
super().__init__(f'[{environment.id}] {message}.', *args, **kwargs)
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
class EnvironmentOutageError(EnvironmentError):
|
|
46
|
+
"""Error that indicates environment is offline."""
|
|
47
|
+
|
|
48
|
+
def __init__(
|
|
49
|
+
self,
|
|
50
|
+
message: str | None = None,
|
|
51
|
+
*args,
|
|
52
|
+
offline_duration: float,
|
|
53
|
+
**kwargs
|
|
54
|
+
):
|
|
55
|
+
self.offline_duration = offline_duration
|
|
56
|
+
super().__init__(
|
|
57
|
+
message or f'Environment is offline for {offline_duration} seconds.',
|
|
58
|
+
*args,
|
|
59
|
+
**kwargs
|
|
60
|
+
)
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
class EnvironmentOverloadError(EnvironmentError):
|
|
64
|
+
"""Error that indicates environment is overloaded."""
|
|
65
|
+
|
|
66
|
+
def __init__(
|
|
67
|
+
self,
|
|
68
|
+
message: str | None = None,
|
|
69
|
+
*args,
|
|
70
|
+
**kwargs
|
|
71
|
+
):
|
|
72
|
+
super().__init__(
|
|
73
|
+
message or 'All sandboxes in the pool are either busy or dead.',
|
|
74
|
+
*args, **kwargs
|
|
75
|
+
)
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
class SandboxError(RuntimeError):
|
|
79
|
+
"""Base class for sandbox errors."""
|
|
80
|
+
|
|
81
|
+
def __init__(
|
|
82
|
+
self,
|
|
83
|
+
message: str,
|
|
84
|
+
*args,
|
|
85
|
+
sandbox: 'Sandbox',
|
|
86
|
+
**kwargs
|
|
87
|
+
) -> None:
|
|
88
|
+
self.sandbox = sandbox
|
|
89
|
+
super().__init__(f'[{sandbox.id}] {message}.', *args, **kwargs)
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
class SandboxStateError(SandboxError):
|
|
93
|
+
"""Error that indicates sandbox is in an unexpected state.
|
|
94
|
+
|
|
95
|
+
This error is raised when the sandbox is in an unexpected state and cannot
|
|
96
|
+
be recovered. As a result, the sandbox will be shutdown and user session
|
|
97
|
+
will be terminated.
|
|
98
|
+
"""
|
|
99
|
+
|
|
100
|
+
def __init__(
|
|
101
|
+
self,
|
|
102
|
+
message: str | None = None,
|
|
103
|
+
*args,
|
|
104
|
+
code: str | None = None,
|
|
105
|
+
**kwargs
|
|
106
|
+
):
|
|
107
|
+
default_message = 'Sandbox is in an unexpected state'
|
|
108
|
+
if code is not None:
|
|
109
|
+
default_message = (
|
|
110
|
+
f'Sandbox is in an unexpected state after executing code: {code!r}'
|
|
111
|
+
)
|
|
112
|
+
super().__init__(message or default_message, *args, **kwargs)
|
|
113
|
+
|
|
114
|
+
|
|
115
|
+
class FeatureTeardownError(SandboxError):
|
|
116
|
+
"""Base class for feature errors."""
|
|
117
|
+
|
|
118
|
+
def __init__(
|
|
119
|
+
self,
|
|
120
|
+
message: str | None = None,
|
|
121
|
+
*args,
|
|
122
|
+
errors: dict[str, BaseException],
|
|
123
|
+
**kwargs
|
|
124
|
+
):
|
|
125
|
+
self.errors = errors
|
|
126
|
+
super().__init__(
|
|
127
|
+
(message or
|
|
128
|
+
f'Feature teardown failed with user-defined errors: {errors}.'),
|
|
129
|
+
*args,
|
|
130
|
+
**kwargs
|
|
131
|
+
)
|
|
132
|
+
|
|
133
|
+
@property
|
|
134
|
+
def has_non_sandbox_state_error(self) -> bool:
|
|
135
|
+
"""Returns True if the feature teardown error has non-sandbox state error."""
|
|
136
|
+
return any(
|
|
137
|
+
not isinstance(e, SandboxStateError) for e in self.errors.values()
|
|
138
|
+
)
|
|
139
|
+
|
|
140
|
+
|
|
141
|
+
class SessionTeardownError(SandboxError):
|
|
142
|
+
"""Base class for session errors."""
|
|
143
|
+
|
|
144
|
+
def __init__(
|
|
145
|
+
self,
|
|
146
|
+
message: str | None = None,
|
|
147
|
+
*args,
|
|
148
|
+
errors: dict[str, BaseException],
|
|
149
|
+
**kwargs
|
|
150
|
+
):
|
|
151
|
+
self.errors = errors
|
|
152
|
+
super().__init__(
|
|
153
|
+
(message or
|
|
154
|
+
f'Session teardown failed with user-defined errors: {errors}.'),
|
|
155
|
+
*args,
|
|
156
|
+
**kwargs
|
|
157
|
+
)
|
|
158
|
+
|
|
159
|
+
@property
|
|
160
|
+
def has_non_sandbox_state_error(self) -> bool:
|
|
161
|
+
"""Returns True if the feature teardown error has non-sandbox state error."""
|
|
162
|
+
return any(
|
|
163
|
+
not isinstance(e, SandboxStateError) for e in self.errors.values()
|
|
164
|
+
)
|
|
165
|
+
|
|
166
|
+
#
|
|
167
|
+
# Interface for sandbox-based environment.
|
|
168
|
+
#
|
|
169
|
+
|
|
170
|
+
|
|
171
|
+
class Environment(pg.Object):
|
|
172
|
+
"""Base class for an environment.
|
|
173
|
+
|
|
174
|
+
An **Environment** is the central component for managing sandboxes and
|
|
175
|
+
**Features**. It acts as an abstraction layer, hiding the implementation
|
|
176
|
+
details of the underlying container/sandboxing system
|
|
177
|
+
(e.g., Docker, virtual machines).
|
|
178
|
+
|
|
179
|
+
The core goal is to enable the development of features that are **agnostic**
|
|
180
|
+
to how or where they are executed (sandboxed or not).
|
|
181
|
+
|
|
182
|
+
-----------------------------------------------------------------------------
|
|
183
|
+
|
|
184
|
+
## Core Concepts
|
|
185
|
+
|
|
186
|
+
1. **Sandbox-Based Features:** Features that require isolated execution
|
|
187
|
+
contexts (sandboxes). They become available as properties of the sandbox
|
|
188
|
+
object (e.g., `sandbox.feature1`). Applicability is determined by an image
|
|
189
|
+
ID regex.
|
|
190
|
+
2. **Non-Sandbox-Based Features:** Features that run directly within the host
|
|
191
|
+
process or manage their own execution context outside of the Environment's
|
|
192
|
+
managed sandboxes. They are accessed directly via the Environment (e.g.,
|
|
193
|
+
`env.feature3()`).
|
|
194
|
+
|
|
195
|
+
How to Use:
|
|
196
|
+
|
|
197
|
+
The primary usage patterns are creating a **sandbox session** or directly
|
|
198
|
+
accessing a specific **Feature**, which transparently handles sandbox creation
|
|
199
|
+
if needed.
|
|
200
|
+
|
|
201
|
+
```python
|
|
202
|
+
env = MyEnvironment(
|
|
203
|
+
image_ids=['image1', 'image2'],
|
|
204
|
+
features={
|
|
205
|
+
'feature1': Feature1(applicable_images=['image1.*']),
|
|
206
|
+
'feature2': Feature2(applicable_images=['image2.*']),
|
|
207
|
+
'feature3': Feature3(is_sandbox_based=False)
|
|
208
|
+
})
|
|
209
|
+
|
|
210
|
+
# Context manager for the Environment lifetime.
|
|
211
|
+
with env:
|
|
212
|
+
|
|
213
|
+
# 1. Access a specific sandbox directly.
|
|
214
|
+
# Upon exiting the context, the sandbox will be shutdown or returned to
|
|
215
|
+
# the pool for reuse.
|
|
216
|
+
with env.sandbox(image_id='image1') as sandbox:
|
|
217
|
+
# Execute a shell command inside the sandbox.
|
|
218
|
+
sandbox.shell('echo "hello world"')
|
|
219
|
+
|
|
220
|
+
# Access a sandbox-based feature (feature1 is applicable to image1).
|
|
221
|
+
sandbox.feature1.feature_method()
|
|
222
|
+
|
|
223
|
+
# Attempts to access inapplicable features will raise AttributeError:
|
|
224
|
+
# sandbox.feature2 # Not applicable to image1.
|
|
225
|
+
# sandbox.feature3 # Not sandbox-based.
|
|
226
|
+
|
|
227
|
+
# 2. Access a sandbox-based feature and let the Environment manage the
|
|
228
|
+
# sandbox. A suitable sandbox (e.g., one built from an image matching
|
|
229
|
+
# 'image1.*') will be provisioned, and the feature instance will be yielded.
|
|
230
|
+
with env.feature1() as feature1:
|
|
231
|
+
feature1.feature_method()
|
|
232
|
+
|
|
233
|
+
# 3. Access a non-sandbox-based feature.
|
|
234
|
+
with env.feature3() as feature3:
|
|
235
|
+
feature3.feature_method()
|
|
236
|
+
```
|
|
237
|
+
|
|
238
|
+
-----------------------------------------------------------------------------
|
|
239
|
+
|
|
240
|
+
## Multi-tenancy and Pooling
|
|
241
|
+
|
|
242
|
+
The Environment supports multi-tenancy (working with multiple image types) and
|
|
243
|
+
pooling (reusing sandboxes) to amortize setup costs across different user
|
|
244
|
+
requests. Pooling is configured via the pool_size parameter.
|
|
245
|
+
|
|
246
|
+
| pool_size Value | Behavior |
|
|
247
|
+
|---------------------------|------------------------------------------------|
|
|
248
|
+
| 0 or (0, 0) | No pooling. Sandboxes are created and shut down|
|
|
249
|
+
| | on demand (useful for local development). |
|
|
250
|
+
|---------------------------|------------------------------------------------|
|
|
251
|
+
| (MIN, MAX) tuple | Global Pool: Applies the same minimum and |
|
|
252
|
+
| | maximum pool size to sandboxes created from all|
|
|
253
|
+
| | specified images. |
|
|
254
|
+
|---------------------------|------------------------------------------------|
|
|
255
|
+
| {image_regex: (MIN, MAX)} | Per-Image Pool: Allows customizing pool |
|
|
256
|
+
| | settings based on image ID regular expressions.|
|
|
257
|
+
| | (e.g., 'image1.*': (64, 256), '.*': (0, 256)). |
|
|
258
|
+
|
|
259
|
+
|
|
260
|
+
**Example 1: No Pooling (pool_size=0)**
|
|
261
|
+
Sandboxes are created and shutdown immediately upon session end.
|
|
262
|
+
|
|
263
|
+
```python
|
|
264
|
+
env = MyEnvironment(image_ids=['image1', 'image2'], pool_size=0)
|
|
265
|
+
|
|
266
|
+
# Sandbox created and shutdown on demand.
|
|
267
|
+
with env.sandbox(image_id='image1') as sandbox1:
|
|
268
|
+
...
|
|
269
|
+
```
|
|
270
|
+
|
|
271
|
+
**Exaxmple 2: Global Pooling (pool_size=(0, 256))**
|
|
272
|
+
Up to 256 sandboxes will be created and pooled across both images as needed.
|
|
273
|
+
None are created initially.
|
|
274
|
+
|
|
275
|
+
```python
|
|
276
|
+
env = MyEnvironment(
|
|
277
|
+
image_ids=['image1', 'image2'],
|
|
278
|
+
pool_size=(0, 256)
|
|
279
|
+
)
|
|
280
|
+
```
|
|
281
|
+
|
|
282
|
+
**Example 3: Per-Image Custom Pooling**:
|
|
283
|
+
For images matching 'image1.*': 64 sandboxes are pre-created (MIN=64) and
|
|
284
|
+
pooled, up to a MAX of 256.
|
|
285
|
+
For all other images ('.*'): Sandboxes are created and pooled on demand
|
|
286
|
+
(MIN=0), up to a MAX of 256.
|
|
287
|
+
|
|
288
|
+
```python
|
|
289
|
+
env = MyEnvironment(
|
|
290
|
+
image_ids=['image1', 'image2'],
|
|
291
|
+
pool_size={
|
|
292
|
+
'image1.*': (64, 256),
|
|
293
|
+
'.*': (0, 256),
|
|
294
|
+
}
|
|
295
|
+
)
|
|
296
|
+
```
|
|
297
|
+
|
|
298
|
+
## Handling Sandbox Failures
|
|
299
|
+
|
|
300
|
+
Sandboxes often run in distributed, ephemeral environments and must be treated
|
|
301
|
+
as fault-tolerant. Langfun provides a protocol for handling unexpected sandbox
|
|
302
|
+
state issues.
|
|
303
|
+
|
|
304
|
+
### Communicating Errors
|
|
305
|
+
If a feature encounters an unexpected state in its sandbox (e.g., a process
|
|
306
|
+
died), it should raise `lf.env.SandboxStateError`.
|
|
307
|
+
|
|
308
|
+
The sandbox will be automatically shut down when its context manager exits.
|
|
309
|
+
The Environment handles replacement and ensures future requests are routed to
|
|
310
|
+
healthy sandboxes.
|
|
311
|
+
|
|
312
|
+
For example:
|
|
313
|
+
```
|
|
314
|
+
with env:
|
|
315
|
+
with env.sandbox() as sb:
|
|
316
|
+
# If SandboxStateError is raised within this block, the sandbox
|
|
317
|
+
# will be forcefully shut down upon block exit.
|
|
318
|
+
sb.shell('echo hi')
|
|
319
|
+
```
|
|
320
|
+
|
|
321
|
+
### Robust User Code
|
|
322
|
+
A simple strategy for robust user code is to wrap critical operations in a
|
|
323
|
+
retry loop:
|
|
324
|
+
```
|
|
325
|
+
while True:
|
|
326
|
+
try:
|
|
327
|
+
result = do_something_that_involves_sandbox()
|
|
328
|
+
break # Success!
|
|
329
|
+
except lf.env.SandboxStateError:
|
|
330
|
+
# The sandbox failed; a new, healthy one will be provisioned on the next
|
|
331
|
+
# iteration.
|
|
332
|
+
# Wait briefly to avoid resource thrashing.
|
|
333
|
+
time.sleep(1)
|
|
334
|
+
except lf.env.EnvironmentOutageError:
|
|
335
|
+
# If the Environment is down for too long
|
|
336
|
+
# (past BaseEnvironment.outage_grace_period)
|
|
337
|
+
# and cannot provision a healthy replacement, this error is raised.
|
|
338
|
+
# The retry loop should be broken or an outer failure reported.
|
|
339
|
+
raise
|
|
340
|
+
```
|
|
341
|
+
"""
|
|
342
|
+
|
|
343
|
+
# Disable symbolic comparison and hashing for environment objects.
|
|
344
|
+
use_symbolic_comparison = False
|
|
345
|
+
|
|
346
|
+
@dataclasses.dataclass(frozen=True)
|
|
347
|
+
class Id:
|
|
348
|
+
"""Identifier for an environment."""
|
|
349
|
+
environment_id: str
|
|
350
|
+
|
|
351
|
+
def __str__(self) -> str:
|
|
352
|
+
return self.environment_id
|
|
353
|
+
|
|
354
|
+
def working_dir(self, root_dir: str | None) -> str | None:
|
|
355
|
+
"""Returns the download directory for the service."""
|
|
356
|
+
if root_dir is None:
|
|
357
|
+
return None
|
|
358
|
+
return os.path.join(root_dir, _make_path_compatible(self.environment_id))
|
|
359
|
+
|
|
360
|
+
class Status(enum.Enum):
|
|
361
|
+
"""Environment state.
|
|
362
|
+
|
|
363
|
+
State transitions:
|
|
364
|
+
|
|
365
|
+
+---------------+ +-----------+
|
|
366
|
+
| <CREATED> | --(start)---> | ONLINE | -(shutdown or outage detected)
|
|
367
|
+
+---------------+ | +-----------+ |
|
|
368
|
+
| +-----------+ |
|
|
369
|
+
+------> | <OFFLINE> | <------------+
|
|
370
|
+
+-----------+
|
|
371
|
+
"""
|
|
372
|
+
CREATED = 'created'
|
|
373
|
+
ONLINE = 'online'
|
|
374
|
+
SHUTTING_DOWN = 'shutting_down'
|
|
375
|
+
OFFLINE = 'offline'
|
|
376
|
+
|
|
377
|
+
features: Annotated[
|
|
378
|
+
dict[str, 'Feature'],
|
|
379
|
+
'Features to be exposed by the environment.'
|
|
380
|
+
] = {}
|
|
381
|
+
|
|
382
|
+
_ENV_STACK: Annotated[
|
|
383
|
+
ClassVar[list['Environment']],
|
|
384
|
+
'Recording the environments stacked through context managers.'
|
|
385
|
+
] = []
|
|
386
|
+
|
|
387
|
+
#
|
|
388
|
+
# Subclasses must implement:
|
|
389
|
+
#
|
|
390
|
+
|
|
391
|
+
@property
|
|
392
|
+
@abc.abstractmethod
|
|
393
|
+
def id(self) -> Id:
|
|
394
|
+
"""Returns the identifier for the environment."""
|
|
395
|
+
|
|
396
|
+
@property
|
|
397
|
+
@abc.abstractmethod
|
|
398
|
+
def image_ids(self) -> list[str]:
|
|
399
|
+
"""Returns the non-dynamic image IDs served by the environment."""
|
|
400
|
+
|
|
401
|
+
def image_id_for(self, feature: 'Feature') -> str:
|
|
402
|
+
"""Returns the default image ID for the environment."""
|
|
403
|
+
for image_id in self.image_ids:
|
|
404
|
+
if feature.is_applicable(image_id):
|
|
405
|
+
return image_id
|
|
406
|
+
raise ValueError(
|
|
407
|
+
f'No image ID found for feature {feature.name} in {self.image_ids}.'
|
|
408
|
+
)
|
|
409
|
+
|
|
410
|
+
def non_sandbox_based_features(self) -> Iterator['Feature']:
|
|
411
|
+
"""Returns non-sandbox-based features."""
|
|
412
|
+
for feature in self.features.values():
|
|
413
|
+
if not feature.is_sandbox_based:
|
|
414
|
+
yield feature
|
|
415
|
+
|
|
416
|
+
@property
|
|
417
|
+
@abc.abstractmethod
|
|
418
|
+
def event_handler(self) -> 'EventHandler':
|
|
419
|
+
"""Returns the event handler for the environment."""
|
|
420
|
+
|
|
421
|
+
@property
|
|
422
|
+
@abc.abstractmethod
|
|
423
|
+
def status(self) -> Status:
|
|
424
|
+
"""Returns the status of the environment."""
|
|
425
|
+
|
|
426
|
+
@abc.abstractmethod
|
|
427
|
+
def stats(self) -> dict[str, Any]:
|
|
428
|
+
"""Returns the stats of the environment."""
|
|
429
|
+
|
|
430
|
+
@abc.abstractmethod
|
|
431
|
+
def start(self) -> None:
|
|
432
|
+
"""Starts the environment.
|
|
433
|
+
|
|
434
|
+
Raises:
|
|
435
|
+
EnvironmentError: If the environment is not available.
|
|
436
|
+
"""
|
|
437
|
+
|
|
438
|
+
@abc.abstractmethod
|
|
439
|
+
def shutdown(self) -> None:
|
|
440
|
+
"""Shuts down the environment.
|
|
441
|
+
|
|
442
|
+
IMPORTANT: This method shall not raise any exceptions.
|
|
443
|
+
"""
|
|
444
|
+
|
|
445
|
+
@abc.abstractmethod
|
|
446
|
+
def acquire(
|
|
447
|
+
self,
|
|
448
|
+
image_id: str | None = None,
|
|
449
|
+
) -> 'Sandbox':
|
|
450
|
+
"""Acquires a free sandbox from the environment.
|
|
451
|
+
|
|
452
|
+
Args:
|
|
453
|
+
image_id: The image ID to use for the sandbox. If None, it will be
|
|
454
|
+
automatically determined by the environment.
|
|
455
|
+
|
|
456
|
+
Returns:
|
|
457
|
+
A free sandbox from the environment.
|
|
458
|
+
|
|
459
|
+
Raises:
|
|
460
|
+
EnvironmentOutageError: If the environment is out of service.
|
|
461
|
+
EnvironmentOverloadError: If the environment is overloaded.
|
|
462
|
+
"""
|
|
463
|
+
|
|
464
|
+
@abc.abstractmethod
|
|
465
|
+
def new_session_id(self, feature_hint: str | None = None) -> str:
|
|
466
|
+
"""Generates a new session ID."""
|
|
467
|
+
|
|
468
|
+
#
|
|
469
|
+
# Environment lifecycle.
|
|
470
|
+
#
|
|
471
|
+
|
|
472
|
+
@property
|
|
473
|
+
def is_online(self) -> bool:
|
|
474
|
+
"""Returns True if the environment is alive."""
|
|
475
|
+
return self.status == self.Status.ONLINE
|
|
476
|
+
|
|
477
|
+
def __enter__(self) -> 'Environment':
|
|
478
|
+
"""Enters the environment and sets it as the current environment."""
|
|
479
|
+
self.start()
|
|
480
|
+
Environment._ENV_STACK.append(self)
|
|
481
|
+
return self
|
|
482
|
+
|
|
483
|
+
def __exit__(self, exc_type, exc_value, traceback):
|
|
484
|
+
"""Exits the environment and reset the current environment."""
|
|
485
|
+
assert Environment._ENV_STACK
|
|
486
|
+
Environment._ENV_STACK.pop()
|
|
487
|
+
self.shutdown()
|
|
488
|
+
|
|
489
|
+
@classmethod
|
|
490
|
+
def current(cls) -> Optional['Environment']:
|
|
491
|
+
"""Returns the current environment."""
|
|
492
|
+
if not Environment._ENV_STACK:
|
|
493
|
+
return None
|
|
494
|
+
return Environment._ENV_STACK[-1]
|
|
495
|
+
|
|
496
|
+
#
|
|
497
|
+
# Environment operations.
|
|
498
|
+
#
|
|
499
|
+
|
|
500
|
+
def sandbox(
|
|
501
|
+
self,
|
|
502
|
+
session_id: str | None = None,
|
|
503
|
+
image_id: str | None = None,
|
|
504
|
+
) -> ContextManager['Sandbox']:
|
|
505
|
+
"""Gets a sandbox from the environment and starts a new user session."""
|
|
506
|
+
return self.acquire(image_id=image_id).new_session(
|
|
507
|
+
session_id or self.new_session_id()
|
|
508
|
+
)
|
|
509
|
+
|
|
510
|
+
def __getattr__(self, name: str) -> Any:
|
|
511
|
+
"""Gets a feature session from a free sandbox from the environment.
|
|
512
|
+
|
|
513
|
+
Example:
|
|
514
|
+
```
|
|
515
|
+
with XboxEnvironment(
|
|
516
|
+
features={'selenium': SeleniumFeature()}
|
|
517
|
+
) as env:
|
|
518
|
+
with env.selenium() as selenium:
|
|
519
|
+
driver = selenium.get_driver()
|
|
520
|
+
```
|
|
521
|
+
|
|
522
|
+
Args:
|
|
523
|
+
name: The name of the feature.
|
|
524
|
+
|
|
525
|
+
Returns:
|
|
526
|
+
A callable `(image_id, *, session_id) -> ContextManager[Feature]` that
|
|
527
|
+
creates a context manager for the requested feature under a new client
|
|
528
|
+
session.
|
|
529
|
+
"""
|
|
530
|
+
if name in self.features:
|
|
531
|
+
return _feature_session_creator(self, self.features[name])
|
|
532
|
+
raise AttributeError(name)
|
|
533
|
+
|
|
534
|
+
|
|
535
|
+
@contextlib.contextmanager
|
|
536
|
+
def _sandbox_session_for_feature(
|
|
537
|
+
environment: Environment,
|
|
538
|
+
feature: 'Feature',
|
|
539
|
+
image_id: str | None = None,
|
|
540
|
+
session_id: str | None = None,
|
|
541
|
+
) -> Iterator['Feature']:
|
|
542
|
+
"""Returns a context manager for a session for a feature."""
|
|
543
|
+
assert feature.is_sandbox_based
|
|
544
|
+
if image_id is None:
|
|
545
|
+
image_id = environment.image_id_for(feature)
|
|
546
|
+
elif not feature.is_applicable(image_id):
|
|
547
|
+
raise ValueError(
|
|
548
|
+
f'Feature {feature.name!r} is not applicable to image {image_id!r}.'
|
|
549
|
+
)
|
|
550
|
+
sandbox = environment.acquire(image_id=image_id)
|
|
551
|
+
with sandbox.new_session(
|
|
552
|
+
session_id=session_id or environment.new_session_id(feature.name)
|
|
553
|
+
):
|
|
554
|
+
yield sandbox.features[feature.name]
|
|
555
|
+
|
|
556
|
+
|
|
557
|
+
def _feature_session_creator(environment: Environment, feature: 'Feature'):
|
|
558
|
+
"""Returns a callable that returns a context manager for a feature session."""
|
|
559
|
+
def fn(session_id: str | None = None, image_id: str | None = None):
|
|
560
|
+
if feature.is_sandbox_based:
|
|
561
|
+
return _sandbox_session_for_feature(
|
|
562
|
+
environment, feature, image_id, session_id
|
|
563
|
+
)
|
|
564
|
+
assert image_id is None, (
|
|
565
|
+
'Non-sandbox based feature does not support image ID.'
|
|
566
|
+
)
|
|
567
|
+
return feature.new_session(
|
|
568
|
+
session_id or environment.new_session_id(feature.name)
|
|
569
|
+
)
|
|
570
|
+
return fn
|
|
571
|
+
|
|
572
|
+
|
|
573
|
+
# Enable automatic conversion from str to Environment.Id.
|
|
574
|
+
pg.typing.register_converter(str, Environment.Id, Environment.Id)
|
|
575
|
+
|
|
576
|
+
|
|
577
|
+
class Sandbox(pg.Object):
|
|
578
|
+
"""Interface for sandboxes.
|
|
579
|
+
|
|
580
|
+
A sandbox is a container that runs a single image with a set of features.
|
|
581
|
+
It will be brought up by the environment, setup the features, fullfill user
|
|
582
|
+
requests, and then tear down features and finally the sandbox itself.
|
|
583
|
+
"""
|
|
584
|
+
|
|
585
|
+
# Disable symbolic comparison and hashing for sandbox objects.
|
|
586
|
+
use_symbolic_comparison = False
|
|
587
|
+
|
|
588
|
+
@dataclasses.dataclass(frozen=True, slots=True)
|
|
589
|
+
class Id:
|
|
590
|
+
"""Identifier for a sandbox."""
|
|
591
|
+
environment_id: Environment.Id
|
|
592
|
+
image_id: str
|
|
593
|
+
sandbox_id: str
|
|
594
|
+
|
|
595
|
+
def __str__(self) -> str:
|
|
596
|
+
return f'{self.environment_id}/{self.image_id}:{self.sandbox_id}'
|
|
597
|
+
|
|
598
|
+
def working_dir(self, root_dir: str | None) -> str | None:
|
|
599
|
+
"""Returns the download directory for the sandbox."""
|
|
600
|
+
if root_dir is None:
|
|
601
|
+
return None
|
|
602
|
+
return os.path.join(
|
|
603
|
+
self.environment_id.working_dir(root_dir),
|
|
604
|
+
_make_path_compatible(self.image_id),
|
|
605
|
+
_make_path_compatible(self.sandbox_id)
|
|
606
|
+
)
|
|
607
|
+
|
|
608
|
+
class Status(enum.Enum):
|
|
609
|
+
r"""Sandbox state.
|
|
610
|
+
|
|
611
|
+
State transitions:
|
|
612
|
+
|
|
613
|
+
(sandbox / feature
|
|
614
|
+
+------------+ teardown) +---------------+
|
|
615
|
+
| <OFFLINE> | <--------------------- | SHUTTING_DOWN |
|
|
616
|
+
+------------+ +---------------+
|
|
617
|
+
^ ^
|
|
618
|
+
/ \
|
|
619
|
+
(setup failed) / \
|
|
620
|
+
/ \
|
|
621
|
+
+-----------+ (start) +------------+ \
|
|
622
|
+
| <CREATED> | --------> | SETTING_UP | \
|
|
623
|
+
+-----------+ ^ +------------+ \
|
|
624
|
+
/ | \
|
|
625
|
+
/ | (sandbox / \
|
|
626
|
+
/ | feature /session \
|
|
627
|
+
/ v setup succeeded) \
|
|
628
|
+
/ +---------+ \
|
|
629
|
+
/ | READY | \
|
|
630
|
+
/ +---------+ \
|
|
631
|
+
/ | \
|
|
632
|
+
/ | (acquire) \
|
|
633
|
+
/ v \
|
|
634
|
+
/ +----------+ \
|
|
635
|
+
| | ACQUIRED | \
|
|
636
|
+
| +----------+ |
|
|
637
|
+
| | |
|
|
638
|
+
| | (start_session) |
|
|
639
|
+
| +------------+ |
|
|
640
|
+
| | SETTING_UP |-- (setup failed) ------>+
|
|
641
|
+
| +------------+ |
|
|
642
|
+
| | |
|
|
643
|
+
| v (succeeded) |
|
|
644
|
+
| +--------------+ |
|
|
645
|
+
| | IN_SESSION |- (op failed) -------->+
|
|
646
|
+
| +--------------+ |
|
|
647
|
+
| | |
|
|
648
|
+
| | (end_session) |
|
|
649
|
+
| | |
|
|
650
|
+
| v (session teardown |
|
|
651
|
+
(setup next +-----------------+ failed OR |
|
|
652
|
+
session for <---------| EXITING_SESSION |- non-reusable -----+
|
|
653
|
+
reusable sandbox) +-----------------+ sandbox)
|
|
654
|
+
|
|
655
|
+
"""
|
|
656
|
+
|
|
657
|
+
# The sandbox is created, but not yet started.
|
|
658
|
+
CREATED = 'created'
|
|
659
|
+
|
|
660
|
+
# The sandbox is being setting up to serve user sessions.
|
|
661
|
+
SETTING_UP = 'setting_up'
|
|
662
|
+
|
|
663
|
+
# The sandbox is set up and free to be acquired by the user.
|
|
664
|
+
READY = 'ready'
|
|
665
|
+
|
|
666
|
+
# The sandbox is acquired by a thread, but not yet in a user session.
|
|
667
|
+
ACQUIRED = 'acquired'
|
|
668
|
+
|
|
669
|
+
# The sandbox is in a user session.
|
|
670
|
+
IN_SESSION = 'in_session'
|
|
671
|
+
|
|
672
|
+
# The sandbox is exiting a user session.
|
|
673
|
+
EXITING_SESSION = 'exiting_session'
|
|
674
|
+
|
|
675
|
+
# The sandbox is being shut down.
|
|
676
|
+
SHUTTING_DOWN = 'shutting_down'
|
|
677
|
+
|
|
678
|
+
# The sandbox is offline.
|
|
679
|
+
OFFLINE = 'offline'
|
|
680
|
+
|
|
681
|
+
@property
|
|
682
|
+
def is_online(self) -> bool:
|
|
683
|
+
"""Returns True if the sandbox is online."""
|
|
684
|
+
return self in (
|
|
685
|
+
Sandbox.Status.SETTING_UP,
|
|
686
|
+
Sandbox.Status.READY,
|
|
687
|
+
Sandbox.Status.ACQUIRED,
|
|
688
|
+
Sandbox.Status.IN_SESSION,
|
|
689
|
+
)
|
|
690
|
+
|
|
691
|
+
@property
|
|
692
|
+
@abc.abstractmethod
|
|
693
|
+
def id(self) -> Id:
|
|
694
|
+
"""Returns the identifier for the sandbox."""
|
|
695
|
+
|
|
696
|
+
@property
|
|
697
|
+
@abc.abstractmethod
|
|
698
|
+
def image_id(self) -> str:
|
|
699
|
+
"""Returns the image ID used for bootstrapping the sandbox."""
|
|
700
|
+
|
|
701
|
+
@property
|
|
702
|
+
@abc.abstractmethod
|
|
703
|
+
def environment(self) -> Environment:
|
|
704
|
+
"""Returns the environment for the sandbox."""
|
|
705
|
+
|
|
706
|
+
@property
|
|
707
|
+
@abc.abstractmethod
|
|
708
|
+
def features(self) -> dict[str, 'Feature']:
|
|
709
|
+
"""Returns the features in the sandbox."""
|
|
710
|
+
|
|
711
|
+
@property
|
|
712
|
+
@abc.abstractmethod
|
|
713
|
+
def status(self) -> Status:
|
|
714
|
+
"""Returns the status of the sandbox."""
|
|
715
|
+
|
|
716
|
+
@property
|
|
717
|
+
def is_online(self) -> bool:
|
|
718
|
+
"""Returns True if the sandbox is online."""
|
|
719
|
+
return self.status.is_online
|
|
720
|
+
|
|
721
|
+
@abc.abstractmethod
|
|
722
|
+
def report_state_error(self, error: SandboxStateError) -> None:
|
|
723
|
+
"""Reports state error the sandbox.
|
|
724
|
+
|
|
725
|
+
If state errors are reported, the sandbox will be forcefully shutdown when
|
|
726
|
+
`Sandbox.end_session()` is called, even if the sandbox is set to be
|
|
727
|
+
reusable.
|
|
728
|
+
|
|
729
|
+
Args:
|
|
730
|
+
error: SandboxStateError to report.
|
|
731
|
+
"""
|
|
732
|
+
|
|
733
|
+
@abc.abstractmethod
|
|
734
|
+
def start(self) -> None:
|
|
735
|
+
"""Starts the sandbox.
|
|
736
|
+
|
|
737
|
+
State transitions:
|
|
738
|
+
CREATED -> SETTING_UP -> READY: When all sandbox and feature setup
|
|
739
|
+
succeeds.
|
|
740
|
+
CREATED -> SETTING_UP -> SHUTTING_DOWN -> OFFLINE: When sandbox or feature
|
|
741
|
+
setup fails.
|
|
742
|
+
|
|
743
|
+
`start` and `shutdown` should be called in pairs, even when the sandbox
|
|
744
|
+
fails to start. This ensures proper cleanup.
|
|
745
|
+
|
|
746
|
+
Start may fail with two sources of errors:
|
|
747
|
+
|
|
748
|
+
1. SandboxStateError: If sandbox or feature setup fail due to enviroment
|
|
749
|
+
outage or sandbox state errors.
|
|
750
|
+
2. BaseException: If feature setup failed with user-defined errors, this
|
|
751
|
+
could happen when there is bug in the user code or non-environment code
|
|
752
|
+
failure.
|
|
753
|
+
|
|
754
|
+
In both cases, the sandbox will be shutdown automatically, and the error
|
|
755
|
+
will be add to `errors`. The sandbox is considered dead and will not be
|
|
756
|
+
further used.
|
|
757
|
+
|
|
758
|
+
Raises:
|
|
759
|
+
SandboxStateError: If the sandbox is in a bad state.
|
|
760
|
+
BaseException: If feature setup failed with user-defined errors.
|
|
761
|
+
"""
|
|
762
|
+
|
|
763
|
+
@abc.abstractmethod
|
|
764
|
+
def shutdown(self) -> None:
|
|
765
|
+
"""Shuts down the sandbox.
|
|
766
|
+
|
|
767
|
+
State transitions:
|
|
768
|
+
SHUTTING_DOWN -> SHUTTING_DOWN: No operation.
|
|
769
|
+
OFFLINE -> OFFLINE: No operation.
|
|
770
|
+
SETTING_UP -> SHUTTING_DOWN -> OFFLINE: When sandbox and feature
|
|
771
|
+
setup fails.
|
|
772
|
+
IN_SESSION -> SHUTTING_DOWN -> OFFLINE: When user session exits while
|
|
773
|
+
sandbox is set not to reuse, or session teardown fails.
|
|
774
|
+
FREE -> SHUTTING_DOWN -> OFFLINE: When sandbox is shutdown when the
|
|
775
|
+
environment is shutting down, or housekeeping loop shuts down the
|
|
776
|
+
sandbox due to housekeeping failures.
|
|
777
|
+
|
|
778
|
+
|
|
779
|
+
Please be aware that `shutdown` will be called whenever an operation on the
|
|
780
|
+
sandbox encounters a critical error. This means, `shutdown` should not make
|
|
781
|
+
the assumption that the sandbox is in a healthy state, even `start` could
|
|
782
|
+
fail. As a result, `shutdown` must allow re-entry and be thread-safe with
|
|
783
|
+
other sandbox operations.
|
|
784
|
+
|
|
785
|
+
Shutdown may fail with two sources of errors:
|
|
786
|
+
|
|
787
|
+
1. SandboxStateError: If the sandbox is in a bad state, and feature teardown
|
|
788
|
+
logic depending on a healthy sandbox may fail. In such case, we do not
|
|
789
|
+
raise error to the user as the user session is considered completed. The
|
|
790
|
+
sandbox is abandoned and new user sessions will be served on other
|
|
791
|
+
sandboxes.
|
|
792
|
+
|
|
793
|
+
2. BaseException: The sandbox is in good state, but user code raises error
|
|
794
|
+
due to bug or non-environment code failure. In such case, errors will be
|
|
795
|
+
raised to the user so the error could be surfaced and handled properly.
|
|
796
|
+
The sandbox is treated as shutdown and will not be further used.
|
|
797
|
+
|
|
798
|
+
Raises:
|
|
799
|
+
BaseException: If feature teardown failed with user-defined errors.
|
|
800
|
+
"""
|
|
801
|
+
|
|
802
|
+
@abc.abstractmethod
|
|
803
|
+
def start_session(
|
|
804
|
+
self,
|
|
805
|
+
session_id: str,
|
|
806
|
+
) -> None:
|
|
807
|
+
"""Begins a user session with the sandbox.
|
|
808
|
+
|
|
809
|
+
State transitions:
|
|
810
|
+
ACQUIRED -> SETTING_UP -> IN_SESSION: When session setup succeeds.
|
|
811
|
+
ACQUIRED -> SETTING_UP -> SHUTTING_DOWN -> OFFLINE: When session setup
|
|
812
|
+
fails.
|
|
813
|
+
|
|
814
|
+
A session is a sequence of stateful interactions with the sandbox.
|
|
815
|
+
Across different sessions the sandbox are considered stateless.
|
|
816
|
+
`start_session` and `end_session` should always be called in pairs, even
|
|
817
|
+
when the session fails to start. `Sandbox.new_session` context manager is
|
|
818
|
+
the recommended way to use `start_session` and `end_session` in pairs.
|
|
819
|
+
|
|
820
|
+
Starting a session may fail with two sources of errors:
|
|
821
|
+
|
|
822
|
+
1. SandboxStateError: If the sandbox is in a bad state or session setup
|
|
823
|
+
failed.
|
|
824
|
+
|
|
825
|
+
2. BaseException: If session setup failed with user-defined errors.
|
|
826
|
+
|
|
827
|
+
In both cases, the sandbox will be shutdown automatically and the
|
|
828
|
+
session will be considered ended. The error will be added to `errors`.
|
|
829
|
+
Future session will be served on other sandboxes.
|
|
830
|
+
|
|
831
|
+
Args:
|
|
832
|
+
session_id: The identifier for the user session.
|
|
833
|
+
|
|
834
|
+
Raises:
|
|
835
|
+
SandboxStateError: If the sandbox is already in a bad state or session
|
|
836
|
+
setup failed.
|
|
837
|
+
BaseException: If session setup failed with user-defined errors.
|
|
838
|
+
"""
|
|
839
|
+
|
|
840
|
+
@abc.abstractmethod
|
|
841
|
+
def end_session(self) -> None:
|
|
842
|
+
"""Ends the user session with the sandbox.
|
|
843
|
+
|
|
844
|
+
State transitions:
|
|
845
|
+
IN_SESSION -> EXITING_SESSION -> READY: When user session exits normally,
|
|
846
|
+
and sandbox is set to reuse.
|
|
847
|
+
IN_SESSION -> EXITING_SESSION -> SHUTTING_DOWN -> OFFLINE: When user
|
|
848
|
+
session exits while
|
|
849
|
+
sandbox is set not to reuse, or session teardown fails.
|
|
850
|
+
IN_SESSION -> EXITING_SESSION -> SETTING_UP -> READY: When user session
|
|
851
|
+
exits normally, and sandbox is set to reuse, and proactive session setup
|
|
852
|
+
is enabled.
|
|
853
|
+
IN_SESSION -> EXITING_SESSION -> SETTING_UP -> SHUTTING_DOWN -> OFFLINE:
|
|
854
|
+
When user session exits normally, and proactive session setup is enabled
|
|
855
|
+
but fails.
|
|
856
|
+
EXITING_SESSION -> EXITING_SESSION: No operation.
|
|
857
|
+
not IN_SESSION -> same state: No operation
|
|
858
|
+
|
|
859
|
+
`end_session` should always be called for each `start_session` call, even
|
|
860
|
+
when the session fails to start, to ensure proper cleanup.
|
|
861
|
+
|
|
862
|
+
When `end_session` is called with state errors reported, the sandbox will be
|
|
863
|
+
forcefully shutdown even if the sandbox is set to be reusable.
|
|
864
|
+
|
|
865
|
+
`end_session` may fail with two sources of errors:
|
|
866
|
+
|
|
867
|
+
1. SandboxStateError: If the sandbox is in a bad state or session teardown
|
|
868
|
+
failed.
|
|
869
|
+
|
|
870
|
+
2. BaseException: If session teardown failed with user-defined errors.
|
|
871
|
+
|
|
872
|
+
In both cases, the sandbox will be shutdown automatically and the
|
|
873
|
+
session will be considered ended. The error will be added to `errors`.
|
|
874
|
+
Future session will be served on other sandboxes.
|
|
875
|
+
|
|
876
|
+
However, SandboxStateError encountered during `end_session` will NOT be
|
|
877
|
+
raised to the user as the user session is considered completed.
|
|
878
|
+
|
|
879
|
+
Raises:
|
|
880
|
+
BaseException: If session teardown failed with user-defined errors.
|
|
881
|
+
"""
|
|
882
|
+
|
|
883
|
+
@property
|
|
884
|
+
@abc.abstractmethod
|
|
885
|
+
def session_id(self) -> str | None:
|
|
886
|
+
"""Returns the current user session identifier."""
|
|
887
|
+
|
|
888
|
+
@abc.abstractmethod
|
|
889
|
+
def track_activity(
|
|
890
|
+
self,
|
|
891
|
+
name: str,
|
|
892
|
+
**kwargs: Any
|
|
893
|
+
) -> ContextManager[None]:
|
|
894
|
+
"""Context manager that tracks a sandbox activity.
|
|
895
|
+
|
|
896
|
+
Args:
|
|
897
|
+
name: The name of the activity.
|
|
898
|
+
**kwargs: Additional keyword arguments to pass to the activity handler.
|
|
899
|
+
|
|
900
|
+
Returns:
|
|
901
|
+
A context manager that tracks the activity, including duration and error.
|
|
902
|
+
"""
|
|
903
|
+
|
|
904
|
+
#
|
|
905
|
+
# API related to a user session.
|
|
906
|
+
# A sandbox could be reused across different user sessions.
|
|
907
|
+
# A user session is a sequence of stateful interactions with the sandbox,
|
|
908
|
+
# Across different sessions the sandbox are considered stateless.
|
|
909
|
+
#
|
|
910
|
+
|
|
911
|
+
@contextlib.contextmanager
|
|
912
|
+
def new_session(self, session_id: str) -> Iterator['Sandbox']:
|
|
913
|
+
"""Context manager for obtaining a sandbox for a user session.
|
|
914
|
+
|
|
915
|
+
State transitions:
|
|
916
|
+
ACQUIRED -> IN_SESSION -> READY: When session setup and teardown succeed.
|
|
917
|
+
ACQUIRED -> IN_SESSINO -> OFFLINE: When session setup or teardown fails.
|
|
918
|
+
|
|
919
|
+
Args:
|
|
920
|
+
session_id: The identifier for the user session.
|
|
921
|
+
|
|
922
|
+
Yields:
|
|
923
|
+
The sandbox for the user session.
|
|
924
|
+
|
|
925
|
+
Raises:
|
|
926
|
+
SandboxStateError: If a session cannot be started on the sandbox.
|
|
927
|
+
BaseException: If session setup or teardown failed with user-defined
|
|
928
|
+
errors.
|
|
929
|
+
"""
|
|
930
|
+
self.start_session(session_id)
|
|
931
|
+
try:
|
|
932
|
+
yield self
|
|
933
|
+
except SandboxStateError as e:
|
|
934
|
+
self.report_state_error(e)
|
|
935
|
+
raise
|
|
936
|
+
finally:
|
|
937
|
+
self.end_session()
|
|
938
|
+
|
|
939
|
+
def __getattr__(self, name: str) -> Any:
|
|
940
|
+
"""Gets a feature from current sandbox.
|
|
941
|
+
|
|
942
|
+
Example:
|
|
943
|
+
```
|
|
944
|
+
with MyEnvironment(
|
|
945
|
+
features={'feature1': Feature1()}
|
|
946
|
+
) as env:
|
|
947
|
+
with env.sandbox('session1') as sb:
|
|
948
|
+
driver = sb.feature1.feature_method()
|
|
949
|
+
```
|
|
950
|
+
|
|
951
|
+
Args:
|
|
952
|
+
name: The name of the feature.
|
|
953
|
+
|
|
954
|
+
Returns:
|
|
955
|
+
A feature from current sandbox.
|
|
956
|
+
"""
|
|
957
|
+
if name in self.features:
|
|
958
|
+
return self.features[name]
|
|
959
|
+
raise AttributeError(name)
|
|
960
|
+
|
|
961
|
+
|
|
962
|
+
class Feature(pg.Object):
|
|
963
|
+
"""Interface for features that run in a Langfun environment.
|
|
964
|
+
|
|
965
|
+
There are two type of features: sandbox-based and non-sandbox-based.
|
|
966
|
+
Sandbox-based features run in a sandbox, which is emulated in a separate
|
|
967
|
+
process. Non-sandbox-based features do not run in a sandbox.
|
|
968
|
+
|
|
969
|
+
Features can be directly accessed through the environment, for example:
|
|
970
|
+
|
|
971
|
+
```python
|
|
972
|
+
env = MyEnvironment(
|
|
973
|
+
feature={
|
|
974
|
+
'feature1': SandboxBasedFeature(),
|
|
975
|
+
'feature2': NonSandboxBasedFeature(),
|
|
976
|
+
}
|
|
977
|
+
)
|
|
978
|
+
# Start the environment.
|
|
979
|
+
with env:
|
|
980
|
+
# Access feature1, which involves acquiring a sandbox and return the feature
|
|
981
|
+
# associated with the sandbox.
|
|
982
|
+
with env.feature1() as f1:
|
|
983
|
+
f1.feature_method()
|
|
984
|
+
|
|
985
|
+
# Access feature2, which does not involve acquiring a sandbox.
|
|
986
|
+
with env.feature2() as f2:
|
|
987
|
+
f2.feature_method()
|
|
988
|
+
```
|
|
989
|
+
|
|
990
|
+
Sandbox-based features can also be accessed through the sandbox, for example:
|
|
991
|
+
|
|
992
|
+
```python
|
|
993
|
+
with env.sandbox('session1') as sb:
|
|
994
|
+
# Access feature1 within the sandbox.
|
|
995
|
+
sb.feature1.feature_method()
|
|
996
|
+
|
|
997
|
+
# Attribute error.
|
|
998
|
+
sb.feature2
|
|
999
|
+
```
|
|
1000
|
+
"""
|
|
1001
|
+
|
|
1002
|
+
@dataclasses.dataclass
|
|
1003
|
+
class Id:
|
|
1004
|
+
container_id: Environment.Id | Sandbox.Id
|
|
1005
|
+
feature_name: str
|
|
1006
|
+
|
|
1007
|
+
def __str__(self) -> str:
|
|
1008
|
+
return f'{self.container_id}/{self.feature_name}'
|
|
1009
|
+
|
|
1010
|
+
def working_dir(self, root_dir: str | None) -> str | None:
|
|
1011
|
+
"""Returns the working directory for the feature."""
|
|
1012
|
+
if root_dir is None:
|
|
1013
|
+
return None
|
|
1014
|
+
return os.path.join(
|
|
1015
|
+
self.container_id.working_dir(root_dir),
|
|
1016
|
+
_make_path_compatible(self.feature_name)
|
|
1017
|
+
)
|
|
1018
|
+
|
|
1019
|
+
# Disable symbolic comparison and hashing for sandbox objects.
|
|
1020
|
+
allow_symbolic_comparison = False
|
|
1021
|
+
|
|
1022
|
+
@property
|
|
1023
|
+
@abc.abstractmethod
|
|
1024
|
+
def name(self) -> str:
|
|
1025
|
+
"""Name of the feature, which will be used as key to access the feature."""
|
|
1026
|
+
|
|
1027
|
+
@functools.cached_property
|
|
1028
|
+
def id(self) -> Id:
|
|
1029
|
+
"""Returns the identifier of the feature."""
|
|
1030
|
+
if self.is_sandbox_based:
|
|
1031
|
+
return Feature.Id(self.sandbox.id, self.name)
|
|
1032
|
+
return Feature.Id(self.environment.id, self.name)
|
|
1033
|
+
|
|
1034
|
+
@property
|
|
1035
|
+
@abc.abstractmethod
|
|
1036
|
+
def environment(self) -> Environment:
|
|
1037
|
+
"""Returns the environment that the feature is running in."""
|
|
1038
|
+
|
|
1039
|
+
@property
|
|
1040
|
+
@abc.abstractmethod
|
|
1041
|
+
def is_sandbox_based(self) -> bool:
|
|
1042
|
+
"""Returns True if the feature is sandbox-based."""
|
|
1043
|
+
|
|
1044
|
+
@property
|
|
1045
|
+
@abc.abstractmethod
|
|
1046
|
+
def sandbox(self) -> Sandbox | None:
|
|
1047
|
+
"""Returns the sandbox that the feature is running in.
|
|
1048
|
+
|
|
1049
|
+
Returns:
|
|
1050
|
+
The sandbox that the feature is running in. None if the feature is not
|
|
1051
|
+
sandbox-based or not yet bound with a sandbox.
|
|
1052
|
+
"""
|
|
1053
|
+
|
|
1054
|
+
@abc.abstractmethod
|
|
1055
|
+
def is_applicable(self, image_id: str) -> bool:
|
|
1056
|
+
"""Returns True if the feature is applicable to the given image."""
|
|
1057
|
+
|
|
1058
|
+
@abc.abstractmethod
|
|
1059
|
+
def setup(self, sandbox: Sandbox | None = None) -> None:
|
|
1060
|
+
"""Sets up the feature.
|
|
1061
|
+
|
|
1062
|
+
For sandbox-based features, the setup will be called when a sandbox is
|
|
1063
|
+
started for the first time.
|
|
1064
|
+
|
|
1065
|
+
For non-sandbox-based features, the setup will be called when the
|
|
1066
|
+
environment starts.
|
|
1067
|
+
|
|
1068
|
+
When a feature's `setup` is called, its `teardown` is guaranteed to be
|
|
1069
|
+
called.
|
|
1070
|
+
|
|
1071
|
+
State transitions:
|
|
1072
|
+
SETTING_UP -> READY: When setup succeeds.
|
|
1073
|
+
SETTING_UP -> SHUTTING_DOWN -> OFFLINE: When setup fails.
|
|
1074
|
+
|
|
1075
|
+
`setup` is called when a sandbox is started for the first time. When a
|
|
1076
|
+
feature's `setup` is called, its `teardown` is guaranteed to be called.
|
|
1077
|
+
|
|
1078
|
+
Args:
|
|
1079
|
+
sandbox: The sandbox that the feature is running in.
|
|
1080
|
+
|
|
1081
|
+
Raises:
|
|
1082
|
+
SandboxStateError: If setup failed due to sandbox state errors.
|
|
1083
|
+
BaseException: If setup failed with user-defined errors.
|
|
1084
|
+
"""
|
|
1085
|
+
|
|
1086
|
+
@abc.abstractmethod
|
|
1087
|
+
def teardown(self) -> None:
|
|
1088
|
+
"""Teardowns the feature, which is called once when the sandbox is down.
|
|
1089
|
+
|
|
1090
|
+
State transitions:
|
|
1091
|
+
SHUTTING_DOWN -> OFFLINE: When teardown succeeds or fails.
|
|
1092
|
+
|
|
1093
|
+
Raises:
|
|
1094
|
+
SandboxStateError: If the sandbox is in a bad state.
|
|
1095
|
+
"""
|
|
1096
|
+
|
|
1097
|
+
@abc.abstractmethod
|
|
1098
|
+
def setup_session(self) -> None:
|
|
1099
|
+
"""Sets up the feature for the upcoming user session.
|
|
1100
|
+
|
|
1101
|
+
State transitions:
|
|
1102
|
+
SETTING_UP -> READY: When session setup succeeds.
|
|
1103
|
+
SETTING_UP -> SHUTTING_DOWN -> OFFLINE: When sessino setup fails.
|
|
1104
|
+
|
|
1105
|
+
`setup_session` is called when a new user session starts for per-session
|
|
1106
|
+
setup. `setup_session` and `teardown_session` will be called in pairs, even
|
|
1107
|
+
`setup_session` fails.
|
|
1108
|
+
|
|
1109
|
+
`setup_session` may fail with two sources of errors:
|
|
1110
|
+
|
|
1111
|
+
1. SandboxStateError: If the sandbox is in a bad state or session setup
|
|
1112
|
+
failed.
|
|
1113
|
+
2. BaseException: If session setup failed with user-defined errors.
|
|
1114
|
+
|
|
1115
|
+
In both cases, the error will be raised to the user and the session will be
|
|
1116
|
+
ended. The sandbox will shutdown automatically and will not be further used.
|
|
1117
|
+
|
|
1118
|
+
Raises:
|
|
1119
|
+
SandboxStateError: If the sandbox is in a bad state or session setup
|
|
1120
|
+
failed.
|
|
1121
|
+
BaseException: If session setup failed with user-defined errors.
|
|
1122
|
+
"""
|
|
1123
|
+
|
|
1124
|
+
@abc.abstractmethod
|
|
1125
|
+
def teardown_session(self) -> None:
|
|
1126
|
+
"""Teardowns the feature for an ending user session.
|
|
1127
|
+
|
|
1128
|
+
State transitions:
|
|
1129
|
+
SHUTTING_DOWN -> OFFLINE: When session teardown succeeds or fails.
|
|
1130
|
+
|
|
1131
|
+
`teardown_session` is called when a user session ends for per-
|
|
1132
|
+
session teardown. `teardown_session` will always be called upon a feature
|
|
1133
|
+
whose `setup_session` is called.
|
|
1134
|
+
|
|
1135
|
+
`teardown_session` may fail with two sources of errors:
|
|
1136
|
+
|
|
1137
|
+
1. SandboxStateError: If the sandbox is in a bad state or session teardown
|
|
1138
|
+
failed.
|
|
1139
|
+
2. BaseException: If session teardown failed with user-defined errors.
|
|
1140
|
+
|
|
1141
|
+
In both cases, the session will be closed and the sandbox will be shutdown.
|
|
1142
|
+
However, SandboxStateError encountered during `teardown_session` will NOT be
|
|
1143
|
+
raised to the user as the user session is considered completed. Other errors
|
|
1144
|
+
will be raised to the user for proper error handling.
|
|
1145
|
+
|
|
1146
|
+
Raises:
|
|
1147
|
+
BaseException: If session teardown failed with user-defined errors.
|
|
1148
|
+
"""
|
|
1149
|
+
|
|
1150
|
+
@abc.abstractmethod
|
|
1151
|
+
def housekeep(self) -> None:
|
|
1152
|
+
"""Performs housekeeping for the feature.
|
|
1153
|
+
|
|
1154
|
+
State transitions:
|
|
1155
|
+
(no state change): When housekeeping succeeds.
|
|
1156
|
+
original state -> SHUTTING_DOWN -> OFFLINE: When housekeeping fails.
|
|
1157
|
+
|
|
1158
|
+
Raises:
|
|
1159
|
+
SandboxStateError: If the sandbox is in a bad state.
|
|
1160
|
+
"""
|
|
1161
|
+
|
|
1162
|
+
@property
|
|
1163
|
+
@abc.abstractmethod
|
|
1164
|
+
def housekeep_interval(self) -> int | None:
|
|
1165
|
+
"""Returns the interval in seconds for feature housekeeping.
|
|
1166
|
+
|
|
1167
|
+
Returns:
|
|
1168
|
+
The interval in seconds for feature housekeeping. If None, the feature
|
|
1169
|
+
will not be housekeeping.
|
|
1170
|
+
"""
|
|
1171
|
+
|
|
1172
|
+
@abc.abstractmethod
|
|
1173
|
+
def track_activity(
|
|
1174
|
+
self,
|
|
1175
|
+
name: str,
|
|
1176
|
+
**kwargs: Any
|
|
1177
|
+
) -> ContextManager[None]:
|
|
1178
|
+
"""Context manager that tracks a feature activity.
|
|
1179
|
+
|
|
1180
|
+
Args:
|
|
1181
|
+
name: The name of the activity.
|
|
1182
|
+
**kwargs: Additional keyword arguments to pass to the activity handler.
|
|
1183
|
+
|
|
1184
|
+
Returns:
|
|
1185
|
+
A context manager that tracks the activity, including duration and error.
|
|
1186
|
+
"""
|
|
1187
|
+
|
|
1188
|
+
@property
|
|
1189
|
+
def session_id(self) -> str | None:
|
|
1190
|
+
"""Returns the current user session identifier."""
|
|
1191
|
+
if self.is_sandbox_based:
|
|
1192
|
+
return self.sandbox.session_id
|
|
1193
|
+
return self._non_sandbox_based_session_id
|
|
1194
|
+
|
|
1195
|
+
@contextlib.contextmanager
|
|
1196
|
+
def new_session(self, session_id: str) -> Iterator['Feature']:
|
|
1197
|
+
"""Context manager for obtaining a non-sandbox-based feature session."""
|
|
1198
|
+
assert not self.is_sandbox_based, (
|
|
1199
|
+
'Applicable only to non-sandbox-based features. '
|
|
1200
|
+
'For sandbox-based features, use `Sandbox.new_session` instead.'
|
|
1201
|
+
)
|
|
1202
|
+
try:
|
|
1203
|
+
self._non_sandbox_based_session_id = session_id
|
|
1204
|
+
self.setup_session()
|
|
1205
|
+
yield self
|
|
1206
|
+
finally:
|
|
1207
|
+
try:
|
|
1208
|
+
# Since the session is ended, we don't want to raise any errors during
|
|
1209
|
+
# session teardown to the user. So we catch all exceptions here.
|
|
1210
|
+
# However the event handler will still be notified and log the error.
|
|
1211
|
+
self.teardown_session()
|
|
1212
|
+
except BaseException: # pylint: disable=broad-except
|
|
1213
|
+
pass
|
|
1214
|
+
self._non_sandbox_based_session_id = None
|
|
1215
|
+
|
|
1216
|
+
def _on_bound(self) -> None:
|
|
1217
|
+
"""Called when the feature is bound to a sandbox."""
|
|
1218
|
+
super()._on_bound()
|
|
1219
|
+
self._non_sandbox_based_session_id = None
|
|
1220
|
+
|
|
1221
|
+
|
|
1222
|
+
def _make_path_compatible(id_str: str) -> str:
|
|
1223
|
+
"""Makes a path compatible with CNS."""
|
|
1224
|
+
return id_str.translate(
|
|
1225
|
+
str.maketrans({
|
|
1226
|
+
'@': '_',
|
|
1227
|
+
':': '_',
|
|
1228
|
+
'#': '_',
|
|
1229
|
+
' ': '',
|
|
1230
|
+
'<': '',
|
|
1231
|
+
'>': '',
|
|
1232
|
+
})
|
|
1233
|
+
)
|
|
1234
|
+
|
|
1235
|
+
|
|
1236
|
+
def treat_as_sandbox_state_error(
|
|
1237
|
+
errors: Sequence[
|
|
1238
|
+
Type[BaseException] | tuple[Type[BaseException], str]
|
|
1239
|
+
] | None = None
|
|
1240
|
+
) -> Callable[..., Any]:
|
|
1241
|
+
"""Decorator for Sandbox/Feature methods to convert errors to SandboxStateError.
|
|
1242
|
+
|
|
1243
|
+
Args:
|
|
1244
|
+
errors: A sequence of exception types or tuples of (error_type, msg_regex).
|
|
1245
|
+
when matched, treat the error as SandboxStateError, which will lead to
|
|
1246
|
+
a sandbox shutdown when caught by `Sandbox.new_session()` context manager.
|
|
1247
|
+
|
|
1248
|
+
Returns:
|
|
1249
|
+
The decorator function.
|
|
1250
|
+
"""
|
|
1251
|
+
|
|
1252
|
+
def decorator(func):
|
|
1253
|
+
@functools.wraps(func)
|
|
1254
|
+
def method_wrapper(self, *args, **kwargs) -> Any:
|
|
1255
|
+
"""Helper function to safely execute logics in the sandbox."""
|
|
1256
|
+
|
|
1257
|
+
assert isinstance(self, (Sandbox, Feature)), self
|
|
1258
|
+
sandbox = self.sandbox if isinstance(self, Feature) else self
|
|
1259
|
+
|
|
1260
|
+
try:
|
|
1261
|
+
# Execute the service function.
|
|
1262
|
+
return func(self, *args, **kwargs)
|
|
1263
|
+
except BaseException as e:
|
|
1264
|
+
if pg.match_error(e, errors):
|
|
1265
|
+
state_error = SandboxStateError(
|
|
1266
|
+
'Sandbox encountered an unexpected error executing '
|
|
1267
|
+
f'`{func.__name__}` (args={args!r}, kwargs={kwargs!r}): {e}',
|
|
1268
|
+
sandbox=sandbox
|
|
1269
|
+
)
|
|
1270
|
+
raise state_error from e
|
|
1271
|
+
raise
|
|
1272
|
+
return method_wrapper
|
|
1273
|
+
return decorator
|
|
1274
|
+
|
|
1275
|
+
|
|
1276
|
+
def log_activity(name: str | None = None):
|
|
1277
|
+
"""Decorator for Sandbox/Feature methods to log sandbox/feature activity."""
|
|
1278
|
+
|
|
1279
|
+
def decorator(func):
|
|
1280
|
+
signature = pg.typing.get_signature(func)
|
|
1281
|
+
def to_kwargs(*args, **kwargs):
|
|
1282
|
+
num_non_self_args = len(signature.arg_names) - 1
|
|
1283
|
+
if len(args) > num_non_self_args:
|
|
1284
|
+
assert signature.varargs is not None, (signature, args)
|
|
1285
|
+
kwargs[signature.varargs.name] = tuple(args[num_non_self_args:])
|
|
1286
|
+
args = args[:num_non_self_args]
|
|
1287
|
+
for i in range(len(args)):
|
|
1288
|
+
# The first argument is `self`.
|
|
1289
|
+
kwargs[signature.arg_names[i + 1]] = args[i]
|
|
1290
|
+
return kwargs
|
|
1291
|
+
|
|
1292
|
+
@functools.wraps(func)
|
|
1293
|
+
def method_wrapper(self, *args, **kwargs) -> Any:
|
|
1294
|
+
"""Helper function to safely execute logics in the sandbox."""
|
|
1295
|
+
|
|
1296
|
+
assert isinstance(self, (Sandbox, Feature)), self
|
|
1297
|
+
with self.track_activity(
|
|
1298
|
+
name or func.__name__,
|
|
1299
|
+
**to_kwargs(*args, **kwargs)
|
|
1300
|
+
):
|
|
1301
|
+
return func(self, *args, **kwargs)
|
|
1302
|
+
return method_wrapper
|
|
1303
|
+
return decorator
|
|
1304
|
+
|
|
1305
|
+
|
|
1306
|
+
#
|
|
1307
|
+
# Interface for event handlers.
|
|
1308
|
+
#
|
|
1309
|
+
|
|
1310
|
+
|
|
1311
|
+
class _EnvironmentEventHandler:
|
|
1312
|
+
"""Base class for event handlers of an environment."""
|
|
1313
|
+
|
|
1314
|
+
def on_environment_starting(self, environment: Environment) -> None:
|
|
1315
|
+
"""Called when the environment is getting started.
|
|
1316
|
+
|
|
1317
|
+
Args:
|
|
1318
|
+
environment: The environment.
|
|
1319
|
+
"""
|
|
1320
|
+
|
|
1321
|
+
def on_environment_start(
|
|
1322
|
+
self,
|
|
1323
|
+
environment: Environment,
|
|
1324
|
+
duration: float,
|
|
1325
|
+
error: BaseException | None
|
|
1326
|
+
) -> None:
|
|
1327
|
+
"""Called when the environment is started.
|
|
1328
|
+
|
|
1329
|
+
Args:
|
|
1330
|
+
environment: The environment.
|
|
1331
|
+
duration: The environment start duration in seconds.
|
|
1332
|
+
error: The error that failed the environment start. If None, the
|
|
1333
|
+
environment started normally.
|
|
1334
|
+
"""
|
|
1335
|
+
|
|
1336
|
+
def on_environment_housekeep(
|
|
1337
|
+
self,
|
|
1338
|
+
environment: Environment,
|
|
1339
|
+
counter: int,
|
|
1340
|
+
duration: float,
|
|
1341
|
+
error: BaseException | None,
|
|
1342
|
+
**kwargs
|
|
1343
|
+
) -> None:
|
|
1344
|
+
"""Called when the environment finishes a round of housekeeping.
|
|
1345
|
+
|
|
1346
|
+
Args:
|
|
1347
|
+
environment: The environment.
|
|
1348
|
+
counter: Zero-based counter of the housekeeping round.
|
|
1349
|
+
duration: The environment start duration in seconds.
|
|
1350
|
+
error: The error that failed the housekeeping. If None, the
|
|
1351
|
+
housekeeping succeeded.
|
|
1352
|
+
**kwargs: Environment-specific properties computed during housekeeping.
|
|
1353
|
+
"""
|
|
1354
|
+
|
|
1355
|
+
def on_environment_shutting_down(
|
|
1356
|
+
self,
|
|
1357
|
+
environment: Environment,
|
|
1358
|
+
offline_duration: float,
|
|
1359
|
+
) -> None:
|
|
1360
|
+
"""Called when the environment is shutting down.
|
|
1361
|
+
|
|
1362
|
+
Args:
|
|
1363
|
+
environment: The environment.
|
|
1364
|
+
offline_duration: The environment offline duration in seconds.
|
|
1365
|
+
"""
|
|
1366
|
+
|
|
1367
|
+
def on_environment_shutdown(
|
|
1368
|
+
self,
|
|
1369
|
+
environment: Environment,
|
|
1370
|
+
duration: float,
|
|
1371
|
+
lifetime: float,
|
|
1372
|
+
error: BaseException | None
|
|
1373
|
+
) -> None:
|
|
1374
|
+
"""Called when the environment is shutdown.
|
|
1375
|
+
|
|
1376
|
+
Args:
|
|
1377
|
+
environment: The environment.
|
|
1378
|
+
duration: The environment shutdown duration in seconds.
|
|
1379
|
+
lifetime: The environment lifetime in seconds.
|
|
1380
|
+
error: The error that caused the environment to shutdown. If None, the
|
|
1381
|
+
environment shutdown normally.
|
|
1382
|
+
"""
|
|
1383
|
+
|
|
1384
|
+
|
|
1385
|
+
class _SandboxEventHandler:
|
|
1386
|
+
"""Base class for sandbox event handlers."""
|
|
1387
|
+
|
|
1388
|
+
def on_sandbox_start(
|
|
1389
|
+
self,
|
|
1390
|
+
sandbox: Sandbox,
|
|
1391
|
+
duration: float,
|
|
1392
|
+
error: BaseException | None
|
|
1393
|
+
) -> None:
|
|
1394
|
+
"""Called when a sandbox is started.
|
|
1395
|
+
|
|
1396
|
+
Args:
|
|
1397
|
+
sandbox: The sandbox.
|
|
1398
|
+
duration: The time spent on starting the sandbox.
|
|
1399
|
+
error: The error that caused the sandbox to start. If None, the sandbox
|
|
1400
|
+
started normally.
|
|
1401
|
+
"""
|
|
1402
|
+
|
|
1403
|
+
def on_sandbox_status_change(
|
|
1404
|
+
self,
|
|
1405
|
+
sandbox: Sandbox,
|
|
1406
|
+
old_status: 'Sandbox.Status',
|
|
1407
|
+
new_status: 'Sandbox.Status',
|
|
1408
|
+
span: float,
|
|
1409
|
+
) -> None:
|
|
1410
|
+
"""Called when a sandbox status changes.
|
|
1411
|
+
|
|
1412
|
+
Args:
|
|
1413
|
+
sandbox: The sandbox.
|
|
1414
|
+
old_status: The old sandbox status.
|
|
1415
|
+
new_status: The new sandbox status.
|
|
1416
|
+
span: Time spent on the old status in seconds.
|
|
1417
|
+
"""
|
|
1418
|
+
|
|
1419
|
+
def on_sandbox_shutdown(
|
|
1420
|
+
self,
|
|
1421
|
+
sandbox: Sandbox,
|
|
1422
|
+
duration: float,
|
|
1423
|
+
lifetime: float,
|
|
1424
|
+
error: BaseException | None
|
|
1425
|
+
) -> None:
|
|
1426
|
+
"""Called when a sandbox is shutdown.
|
|
1427
|
+
|
|
1428
|
+
Args:
|
|
1429
|
+
sandbox: The sandbox.
|
|
1430
|
+
duration: The time spent on shutting down the sandbox.
|
|
1431
|
+
lifetime: The sandbox lifetime in seconds.
|
|
1432
|
+
error: The error that caused the sandbox to shutdown. If None, the
|
|
1433
|
+
sandbox shutdown normally.
|
|
1434
|
+
"""
|
|
1435
|
+
|
|
1436
|
+
def on_sandbox_session_start(
|
|
1437
|
+
self,
|
|
1438
|
+
sandbox: Sandbox,
|
|
1439
|
+
session_id: str,
|
|
1440
|
+
duration: float,
|
|
1441
|
+
error: BaseException | None
|
|
1442
|
+
) -> None:
|
|
1443
|
+
"""Called when a sandbox session starts.
|
|
1444
|
+
|
|
1445
|
+
Args:
|
|
1446
|
+
sandbox: The sandbox.
|
|
1447
|
+
session_id: The session ID.
|
|
1448
|
+
duration: The time spent on starting the session.
|
|
1449
|
+
error: The error that caused the session to start. If None, the session
|
|
1450
|
+
started normally.
|
|
1451
|
+
"""
|
|
1452
|
+
|
|
1453
|
+
def on_sandbox_session_end(
|
|
1454
|
+
self,
|
|
1455
|
+
sandbox: Sandbox,
|
|
1456
|
+
session_id: str,
|
|
1457
|
+
duration: float,
|
|
1458
|
+
lifetime: float,
|
|
1459
|
+
error: BaseException | None
|
|
1460
|
+
) -> None:
|
|
1461
|
+
"""Called when a sandbox session ends.
|
|
1462
|
+
|
|
1463
|
+
Args:
|
|
1464
|
+
sandbox: The sandbox.
|
|
1465
|
+
session_id: The session ID.
|
|
1466
|
+
duration: The time spent on ending the session.
|
|
1467
|
+
lifetime: The session lifetime in seconds.
|
|
1468
|
+
error: The error that caused the session to end. If None, the session
|
|
1469
|
+
ended normally.
|
|
1470
|
+
"""
|
|
1471
|
+
|
|
1472
|
+
def on_sandbox_activity(
|
|
1473
|
+
self,
|
|
1474
|
+
name: str,
|
|
1475
|
+
sandbox: Sandbox,
|
|
1476
|
+
session_id: str | None,
|
|
1477
|
+
duration: float,
|
|
1478
|
+
error: BaseException | None,
|
|
1479
|
+
**kwargs
|
|
1480
|
+
) -> None:
|
|
1481
|
+
"""Called when a sandbox activity is performed.
|
|
1482
|
+
|
|
1483
|
+
Args:
|
|
1484
|
+
name: The name of the sandbox activity.
|
|
1485
|
+
sandbox: The sandbox.
|
|
1486
|
+
session_id: The session ID.
|
|
1487
|
+
duration: The sandbox activity duration in seconds.
|
|
1488
|
+
error: The error that caused the sandbox activity to perform. If None,
|
|
1489
|
+
the sandbox activity performed normally.
|
|
1490
|
+
**kwargs: The keyword arguments of the sandbox activity.
|
|
1491
|
+
"""
|
|
1492
|
+
|
|
1493
|
+
def on_sandbox_housekeep(
|
|
1494
|
+
self,
|
|
1495
|
+
sandbox: Sandbox,
|
|
1496
|
+
counter: int,
|
|
1497
|
+
duration: float,
|
|
1498
|
+
error: BaseException | None,
|
|
1499
|
+
**kwargs
|
|
1500
|
+
) -> None:
|
|
1501
|
+
"""Called when a sandbox finishes a round of housekeeping.
|
|
1502
|
+
|
|
1503
|
+
Args:
|
|
1504
|
+
sandbox: The sandbox.
|
|
1505
|
+
counter: Zero-based counter of the housekeeping round.
|
|
1506
|
+
duration: The sandbox housekeeping duration in seconds.
|
|
1507
|
+
error: The error that caused the sandbox to housekeeping. If None, the
|
|
1508
|
+
sandbox housekeeping normally.
|
|
1509
|
+
**kwargs: Sandbox-specific properties computed during housekeeping.
|
|
1510
|
+
"""
|
|
1511
|
+
|
|
1512
|
+
|
|
1513
|
+
class _FeatureEventHandler:
|
|
1514
|
+
"""Base class for feature event handlers."""
|
|
1515
|
+
|
|
1516
|
+
def on_feature_setup(
|
|
1517
|
+
self,
|
|
1518
|
+
feature: Feature,
|
|
1519
|
+
duration: float,
|
|
1520
|
+
error: BaseException | None
|
|
1521
|
+
) -> None:
|
|
1522
|
+
"""Called when a sandbox feature is setup.
|
|
1523
|
+
|
|
1524
|
+
Applicable to both sandbox-based and non-sandbox-based features.
|
|
1525
|
+
|
|
1526
|
+
Args:
|
|
1527
|
+
feature: The feature.
|
|
1528
|
+
duration: The feature setup duration in seconds.
|
|
1529
|
+
error: The error happened during the feature setup. If None,
|
|
1530
|
+
the feature setup performed normally.
|
|
1531
|
+
"""
|
|
1532
|
+
|
|
1533
|
+
def on_feature_teardown(
|
|
1534
|
+
self,
|
|
1535
|
+
feature: Feature,
|
|
1536
|
+
duration: float,
|
|
1537
|
+
error: BaseException | None
|
|
1538
|
+
) -> None:
|
|
1539
|
+
"""Called when a sandbox feature is teardown.
|
|
1540
|
+
|
|
1541
|
+
Applicable to both sandbox-based and non-sandbox-based features.
|
|
1542
|
+
|
|
1543
|
+
Args:
|
|
1544
|
+
feature: The feature.
|
|
1545
|
+
duration: The feature teardown duration in seconds.
|
|
1546
|
+
error: The error happened during the feature teardown. If None,
|
|
1547
|
+
the feature teardown performed normally.
|
|
1548
|
+
"""
|
|
1549
|
+
|
|
1550
|
+
def on_feature_teardown_session(
|
|
1551
|
+
self,
|
|
1552
|
+
feature: Feature,
|
|
1553
|
+
session_id: str,
|
|
1554
|
+
duration: float,
|
|
1555
|
+
error: BaseException | None
|
|
1556
|
+
) -> None:
|
|
1557
|
+
"""Called when a feature is teardown with a session.
|
|
1558
|
+
|
|
1559
|
+
Applicable to both sandbox-based and non-sandbox-based features.
|
|
1560
|
+
|
|
1561
|
+
Args:
|
|
1562
|
+
feature: The feature.
|
|
1563
|
+
session_id: The session ID.
|
|
1564
|
+
duration: The feature teardown session duration in seconds.
|
|
1565
|
+
error: The error happened during the feature teardown session. If
|
|
1566
|
+
None, the feature teardown session performed normally.
|
|
1567
|
+
"""
|
|
1568
|
+
|
|
1569
|
+
def on_feature_setup_session(
|
|
1570
|
+
self,
|
|
1571
|
+
feature: Feature,
|
|
1572
|
+
session_id: str | None,
|
|
1573
|
+
duration: float,
|
|
1574
|
+
error: BaseException | None,
|
|
1575
|
+
) -> None:
|
|
1576
|
+
"""Called when a feature is setup with a session.
|
|
1577
|
+
|
|
1578
|
+
Applicable to both sandbox-based and non-sandbox-based features.
|
|
1579
|
+
|
|
1580
|
+
Args:
|
|
1581
|
+
feature: The feature.
|
|
1582
|
+
session_id: The session ID.
|
|
1583
|
+
duration: The feature setup session duration in seconds.
|
|
1584
|
+
error: The error happened during the feature setup session. If
|
|
1585
|
+
None, the feature setup session performed normally.
|
|
1586
|
+
"""
|
|
1587
|
+
|
|
1588
|
+
def on_feature_activity(
|
|
1589
|
+
self,
|
|
1590
|
+
name: str,
|
|
1591
|
+
feature: Feature,
|
|
1592
|
+
session_id: str | None,
|
|
1593
|
+
duration: float,
|
|
1594
|
+
error: BaseException | None,
|
|
1595
|
+
**kwargs
|
|
1596
|
+
) -> None:
|
|
1597
|
+
"""Called when a feature activity is performed.
|
|
1598
|
+
|
|
1599
|
+
Applicable to both sandbox-based and non-sandbox-based features.
|
|
1600
|
+
|
|
1601
|
+
Args:
|
|
1602
|
+
name: The name of the feature activity.
|
|
1603
|
+
feature: The feature.
|
|
1604
|
+
session_id: The session ID. Session ID could be None if a feature
|
|
1605
|
+
activity is performed when setting up a session
|
|
1606
|
+
(e.g. BaseEnvironment.proactive_session_setup is on)
|
|
1607
|
+
duration: The feature activity duration in seconds.
|
|
1608
|
+
error: The error happened during the feature activity. If None,
|
|
1609
|
+
the feature activity performed normally.
|
|
1610
|
+
**kwargs: The keyword arguments of the feature activity.
|
|
1611
|
+
"""
|
|
1612
|
+
|
|
1613
|
+
def on_feature_housekeep(
|
|
1614
|
+
self,
|
|
1615
|
+
feature: Feature,
|
|
1616
|
+
counter: int,
|
|
1617
|
+
duration: float,
|
|
1618
|
+
error: BaseException | None,
|
|
1619
|
+
**kwargs,
|
|
1620
|
+
) -> None:
|
|
1621
|
+
"""Called when a sandbox feature is housekeeping.
|
|
1622
|
+
|
|
1623
|
+
Applicable to both sandbox-based and non-sandbox-based features.
|
|
1624
|
+
|
|
1625
|
+
Args:
|
|
1626
|
+
feature: The feature.
|
|
1627
|
+
counter: Zero-based counter of the housekeeping round.
|
|
1628
|
+
duration: The feature housekeeping duration in seconds.
|
|
1629
|
+
error: The error happened during the feature housekeeping. If None, the
|
|
1630
|
+
feature housekeeping normally.
|
|
1631
|
+
**kwargs: Feature-specific properties computed during housekeeping.
|
|
1632
|
+
"""
|
|
1633
|
+
|
|
1634
|
+
|
|
1635
|
+
class EventHandler(
|
|
1636
|
+
_EnvironmentEventHandler,
|
|
1637
|
+
_SandboxEventHandler,
|
|
1638
|
+
_FeatureEventHandler,
|
|
1639
|
+
):
|
|
1640
|
+
"""Base class for langfun/env handlers."""
|