Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit 2205be9

Browse files
ekzhuekzhu
authoredOct 17, 2024··
Merge branch '0.2' into couchbase-vectordb-bugfix
2 parents 7f806c8 + 02977ee commit 2205be9

7 files changed

+1550
-0
lines changed
 

‎LICENSE-CODE-KUBERNETES

+201
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,201 @@
1+
Apache License
2+
Version 2.0, January 2004
3+
http://www.apache.org/licenses/
4+
5+
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
6+
7+
1. Definitions.
8+
9+
"License" shall mean the terms and conditions for use, reproduction,
10+
and distribution as defined by Sections 1 through 9 of this document.
11+
12+
"Licensor" shall mean the copyright owner or entity authorized by
13+
the copyright owner that is granting the License.
14+
15+
"Legal Entity" shall mean the union of the acting entity and all
16+
other entities that control, are controlled by, or are under common
17+
control with that entity. For the purposes of this definition,
18+
"control" means (i) the power, direct or indirect, to cause the
19+
direction or management of such entity, whether by contract or
20+
otherwise, or (ii) ownership of fifty percent (50%) or more of the
21+
outstanding shares, or (iii) beneficial ownership of such entity.
22+
23+
"You" (or "Your") shall mean an individual or Legal Entity
24+
exercising permissions granted by this License.
25+
26+
"Source" form shall mean the preferred form for making modifications,
27+
including but not limited to software source code, documentation
28+
source, and configuration files.
29+
30+
"Object" form shall mean any form resulting from mechanical
31+
transformation or translation of a Source form, including but
32+
not limited to compiled object code, generated documentation,
33+
and conversions to other media types.
34+
35+
"Work" shall mean the work of authorship, whether in Source or
36+
Object form, made available under the License, as indicated by a
37+
copyright notice that is included in or attached to the work
38+
(an example is provided in the Appendix below).
39+
40+
"Derivative Works" shall mean any work, whether in Source or Object
41+
form, that is based on (or derived from) the Work and for which the
42+
editorial revisions, annotations, elaborations, or other modifications
43+
represent, as a whole, an original work of authorship. For the purposes
44+
of this License, Derivative Works shall not include works that remain
45+
separable from, or merely link (or bind by name) to the interfaces of,
46+
the Work and Derivative Works thereof.
47+
48+
"Contribution" shall mean any work of authorship, including
49+
the original version of the Work and any modifications or additions
50+
to that Work or Derivative Works thereof, that is intentionally
51+
submitted to Licensor for inclusion in the Work by the copyright owner
52+
or by an individual or Legal Entity authorized to submit on behalf of
53+
the copyright owner. For the purposes of this definition, "submitted"
54+
means any form of electronic, verbal, or written communication sent
55+
to the Licensor or its representatives, including but not limited to
56+
communication on electronic mailing lists, source code control systems,
57+
and issue tracking systems that are managed by, or on behalf of, the
58+
Licensor for the purpose of discussing and improving the Work, but
59+
excluding communication that is conspicuously marked or otherwise
60+
designated in writing by the copyright owner as "Not a Contribution."
61+
62+
"Contributor" shall mean Licensor and any individual or Legal Entity
63+
on behalf of whom a Contribution has been received by Licensor and
64+
subsequently incorporated within the Work.
65+
66+
2. Grant of Copyright License. Subject to the terms and conditions of
67+
this License, each Contributor hereby grants to You a perpetual,
68+
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
69+
copyright license to reproduce, prepare Derivative Works of,
70+
publicly display, publicly perform, sublicense, and distribute the
71+
Work and such Derivative Works in Source or Object form.
72+
73+
3. Grant of Patent License. Subject to the terms and conditions of
74+
this License, each Contributor hereby grants to You a perpetual,
75+
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
76+
(except as stated in this section) patent license to make, have made,
77+
use, offer to sell, sell, import, and otherwise transfer the Work,
78+
where such license applies only to those patent claims licensable
79+
by such Contributor that are necessarily infringed by their
80+
Contribution(s) alone or by combination of their Contribution(s)
81+
with the Work to which such Contribution(s) was submitted. If You
82+
institute patent litigation against any entity (including a
83+
cross-claim or counterclaim in a lawsuit) alleging that the Work
84+
or a Contribution incorporated within the Work constitutes direct
85+
or contributory patent infringement, then any patent licenses
86+
granted to You under this License for that Work shall terminate
87+
as of the date such litigation is filed.
88+
89+
4. Redistribution. You may reproduce and distribute copies of the
90+
Work or Derivative Works thereof in any medium, with or without
91+
modifications, and in Source or Object form, provided that You
92+
meet the following conditions:
93+
94+
(a) You must give any other recipients of the Work or
95+
Derivative Works a copy of this License; and
96+
97+
(b) You must cause any modified files to carry prominent notices
98+
stating that You changed the files; and
99+
100+
(c) You must retain, in the Source form of any Derivative Works
101+
that You distribute, all copyright, patent, trademark, and
102+
attribution notices from the Source form of the Work,
103+
excluding those notices that do not pertain to any part of
104+
the Derivative Works; and
105+
106+
(d) If the Work includes a "NOTICE" text file as part of its
107+
distribution, then any Derivative Works that You distribute must
108+
include a readable copy of the attribution notices contained
109+
within such NOTICE file, excluding those notices that do not
110+
pertain to any part of the Derivative Works, in at least one
111+
of the following places: within a NOTICE text file distributed
112+
as part of the Derivative Works; within the Source form or
113+
documentation, if provided along with the Derivative Works; or,
114+
within a display generated by the Derivative Works, if and
115+
wherever such third-party notices normally appear. The contents
116+
of the NOTICE file are for informational purposes only and
117+
do not modify the License. You may add Your own attribution
118+
notices within Derivative Works that You distribute, alongside
119+
or as an addendum to the NOTICE text from the Work, provided
120+
that such additional attribution notices cannot be construed
121+
as modifying the License.
122+
123+
You may add Your own copyright statement to Your modifications and
124+
may provide additional or different license terms and conditions
125+
for use, reproduction, or distribution of Your modifications, or
126+
for any such Derivative Works as a whole, provided Your use,
127+
reproduction, and distribution of the Work otherwise complies with
128+
the conditions stated in this License.
129+
130+
5. Submission of Contributions. Unless You explicitly state otherwise,
131+
any Contribution intentionally submitted for inclusion in the Work
132+
by You to the Licensor shall be under the terms and conditions of
133+
this License, without any additional terms or conditions.
134+
Notwithstanding the above, nothing herein shall supersede or modify
135+
the terms of any separate license agreement you may have executed
136+
with Licensor regarding such Contributions.
137+
138+
6. Trademarks. This License does not grant permission to use the trade
139+
names, trademarks, service marks, or product names of the Licensor,
140+
except as required for reasonable and customary use in describing the
141+
origin of the Work and reproducing the content of the NOTICE file.
142+
143+
7. Disclaimer of Warranty. Unless required by applicable law or
144+
agreed to in writing, Licensor provides the Work (and each
145+
Contributor provides its Contributions) on an "AS IS" BASIS,
146+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147+
implied, including, without limitation, any warranties or conditions
148+
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149+
PARTICULAR PURPOSE. You are solely responsible for determining the
150+
appropriateness of using or redistributing the Work and assume any
151+
risks associated with Your exercise of permissions under this License.
152+
153+
8. Limitation of Liability. In no event and under no legal theory,
154+
whether in tort (including negligence), contract, or otherwise,
155+
unless required by applicable law (such as deliberate and grossly
156+
negligent acts) or agreed to in writing, shall any Contributor be
157+
liable to You for damages, including any direct, indirect, special,
158+
incidental, or consequential damages of any character arising as a
159+
result of this License or out of the use or inability to use the
160+
Work (including but not limited to damages for loss of goodwill,
161+
work stoppage, computer failure or malfunction, or any and all
162+
other commercial damages or losses), even if such Contributor
163+
has been advised of the possibility of such damages.
164+
165+
9. Accepting Warranty or Additional Liability. While redistributing
166+
the Work or Derivative Works thereof, You may choose to offer,
167+
and charge a fee for, acceptance of support, warranty, indemnity,
168+
or other liability obligations and/or rights consistent with this
169+
License. However, in accepting such obligations, You may act only
170+
on Your own behalf and on Your sole responsibility, not on behalf
171+
of any other Contributor, and only if You agree to indemnify,
172+
defend, and hold each Contributor harmless for any liability
173+
incurred by, or claims asserted against, such Contributor by reason
174+
of your accepting any such warranty or additional liability.
175+
176+
END OF TERMS AND CONDITIONS
177+
178+
APPENDIX: How to apply the Apache License to your work.
179+
180+
To apply the Apache License to your work, attach the following
181+
boilerplate notice, with the fields enclosed by brackets "[]"
182+
replaced with your own identifying information. (Don't include
183+
the brackets!) The text should be enclosed in the appropriate
184+
comment syntax for the file format. We also recommend that a
185+
file or class name and description of purpose be included on the
186+
same "printed page" as the copyright notice for easier
187+
identification within third-party archives.
188+
189+
Copyright 2014 The Kubernetes Authors.
190+
191+
Licensed under the Apache License, Version 2.0 (the "License");
192+
you may not use this file except in compliance with the License.
193+
You may obtain a copy of the License at
194+
195+
http://www.apache.org/licenses/LICENSE-2.0
196+
197+
Unless required by applicable law or agreed to in writing, software
198+
distributed under the License is distributed on an "AS IS" BASIS,
199+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
200+
See the License for the specific language governing permissions and
201+
limitations under the License.

‎autogen/coding/kubernetes/__init__.py

+5
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
from .pod_commandline_code_executor import PodCommandLineCodeExecutor
2+
3+
__all__ = [
4+
"PodCommandLineCodeExecutor",
5+
]

‎autogen/coding/kubernetes/pod_commandline_code_executor.py

+323
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,323 @@
1+
from __future__ import annotations
2+
3+
import atexit
4+
import importlib
5+
import sys
6+
import textwrap
7+
import uuid
8+
from hashlib import md5
9+
from pathlib import Path
10+
from time import sleep
11+
from types import TracebackType
12+
from typing import Any, ClassVar, Dict, List, Optional, Type, Union
13+
14+
client = importlib.import_module("kubernetes.client")
15+
config = importlib.import_module("kubernetes.config")
16+
ApiException = importlib.import_module("kubernetes.client.rest").ApiException
17+
stream = importlib.import_module("kubernetes.stream").stream
18+
19+
from ...code_utils import TIMEOUT_MSG, _cmd
20+
from ..base import CodeBlock, CodeExecutor, CodeExtractor, CommandLineCodeResult
21+
from ..markdown_code_extractor import MarkdownCodeExtractor
22+
from ..utils import _get_file_name_from_content, silence_pip
23+
24+
if sys.version_info >= (3, 11):
25+
from typing import Self
26+
else:
27+
from typing_extensions import Self
28+
29+
30+
class PodCommandLineCodeExecutor(CodeExecutor):
31+
DEFAULT_EXECUTION_POLICY: ClassVar[Dict[str, bool]] = {
32+
"bash": True,
33+
"shell": True,
34+
"sh": True,
35+
"pwsh": False,
36+
"powershell": False,
37+
"ps1": False,
38+
"python": True,
39+
"javascript": False,
40+
"html": False,
41+
"css": False,
42+
}
43+
LANGUAGE_ALIASES: ClassVar[Dict[str, str]] = {
44+
"py": "python",
45+
"js": "javascript",
46+
}
47+
LANGUAGE_FILE_EXTENSION: ClassVar[Dict[str, str]] = {
48+
"python": "py",
49+
"javascript": "js",
50+
"bash": "sh",
51+
"shell": "sh",
52+
"sh": "sh",
53+
}
54+
55+
def __init__(
56+
self,
57+
image: str = "python:3-slim",
58+
pod_name: Optional[str] = None,
59+
namespace: Optional[str] = None,
60+
pod_spec: Optional[client.V1Pod] = None, # type: ignore
61+
container_name: Optional[str] = "autogen-code-exec",
62+
timeout: int = 60,
63+
work_dir: Union[Path, str] = Path("/workspace"),
64+
kube_config_file: Optional[str] = None,
65+
stop_container: bool = True,
66+
execution_policies: Optional[Dict[str, bool]] = None,
67+
):
68+
"""(Experimental) A code executor class that executes code through
69+
a command line environment in a kubernetes pod.
70+
71+
The executor first saves each code block in a file in the working
72+
directory, and then executes the code file in the container.
73+
The executor executes the code blocks in the order they are received.
74+
Currently, the executor only supports Python and shell scripts.
75+
For Python code, use the language "python" for the code block.
76+
For shell scripts, use the language "bash", "shell", or "sh" for the code
77+
block.
78+
79+
Args:
80+
image (_type_, optional): Docker image to use for code execution.
81+
Defaults to "python:3-slim".
82+
pod_name (Optional[str], optional): Name of the kubernetes pod
83+
which is created. If None, will autogenerate a name. Defaults to None.
84+
namespace (Optional[str], optional): Namespace of kubernetes pod
85+
which is created. If None, will use current namespace of this instance
86+
pod_spec (Optional[client.V1Pod], optional): Specification of kubernetes pod.
87+
custom pod spec can be provided with this param.
88+
if pod_spec is provided, params above(image, pod_name, namespace) are neglected.
89+
container_name (Optional[str], optional): Name of the container where code block will be
90+
executed. if pod_spec param is provided, container_name must be provided also.
91+
timeout (int, optional): The timeout for code execution. Defaults to 60.
92+
work_dir (Union[Path, str], optional): The working directory for the code
93+
execution. Defaults to Path("/workspace").
94+
kube_config_file (Optional[str], optional): kubernetes configuration file path.
95+
If None, will use KUBECONFIG environment variables or service account token(incluster config)
96+
stop_container (bool, optional): If true, will automatically stop the
97+
container when stop is called, when the context manager exits or when
98+
the Python process exits with atext. Defaults to True.
99+
execution_policies (dict[str, bool], optional): defines supported execution language
100+
101+
Raises:
102+
ValueError: On argument error, or if the container fails to start.
103+
"""
104+
if kube_config_file is None:
105+
config.load_config()
106+
else:
107+
config.load_config(config_file=kube_config_file)
108+
109+
self._api_client = client.CoreV1Api()
110+
111+
if timeout < 1:
112+
raise ValueError("Timeout must be greater than or equal to 1.")
113+
self._timeout = timeout
114+
115+
if isinstance(work_dir, str):
116+
work_dir = Path(work_dir)
117+
self._work_dir: Path = work_dir
118+
119+
if container_name is None:
120+
container_name = "autogen-code-exec"
121+
self._container_name = container_name
122+
123+
# Start a container from the image, read to exec commands later
124+
if pod_spec:
125+
pod = pod_spec
126+
else:
127+
if pod_name is None:
128+
pod_name = f"autogen-code-exec-{uuid.uuid4()}"
129+
if namespace is None:
130+
namespace_path = "/var/run/secrets/kubernetes.io/serviceaccount/namespace"
131+
if not Path(namespace_path).is_file():
132+
raise ValueError("Namespace where the pod will be launched must be provided")
133+
with open(namespace_path, "r") as f:
134+
namespace = f.read()
135+
136+
pod = client.V1Pod(
137+
metadata=client.V1ObjectMeta(name=pod_name, namespace=namespace),
138+
spec=client.V1PodSpec(
139+
restart_policy="Never",
140+
containers=[
141+
client.V1Container(
142+
args=["-c", "while true;do sleep 5; done"],
143+
command=["/bin/sh"],
144+
name=container_name,
145+
image=image,
146+
)
147+
],
148+
),
149+
)
150+
151+
try:
152+
pod_name = pod.metadata.name
153+
namespace = pod.metadata.namespace
154+
self._pod = self._api_client.create_namespaced_pod(namespace=namespace, body=pod)
155+
except ApiException as e:
156+
raise ValueError(f"Creating pod failed: {e}")
157+
158+
self._wait_for_ready()
159+
160+
def cleanup() -> None:
161+
try:
162+
self._api_client.delete_namespaced_pod(pod_name, namespace)
163+
except ApiException:
164+
pass
165+
atexit.unregister(cleanup)
166+
167+
self._cleanup = cleanup
168+
169+
if stop_container:
170+
atexit.register(cleanup)
171+
172+
self.execution_policies = self.DEFAULT_EXECUTION_POLICY.copy()
173+
if execution_policies is not None:
174+
self.execution_policies.update(execution_policies)
175+
176+
def _wait_for_ready(self, stop_time: float = 0.1) -> None:
177+
elapsed_time = 0.0
178+
name = self._pod.metadata.name
179+
namespace = self._pod.metadata.namespace
180+
while True:
181+
sleep(stop_time)
182+
elapsed_time += stop_time
183+
if elapsed_time > self._timeout:
184+
raise ValueError(
185+
f"pod name {name} on namespace {namespace} is not Ready after timeout {self._timeout} seconds"
186+
)
187+
try:
188+
pod_status = self._api_client.read_namespaced_pod_status(name, namespace)
189+
if pod_status.status.phase == "Running":
190+
break
191+
except ApiException as e:
192+
raise ValueError(f"reading pod status failed: {e}")
193+
194+
@property
195+
def timeout(self) -> int:
196+
"""(Experimental) The timeout for code execution."""
197+
return self._timeout
198+
199+
@property
200+
def work_dir(self) -> Path:
201+
"""(Experimental) The working directory for the code execution."""
202+
return self._work_dir
203+
204+
@property
205+
def code_extractor(self) -> CodeExtractor:
206+
"""(Experimental) Export a code extractor that can be used by an agent."""
207+
return MarkdownCodeExtractor()
208+
209+
def execute_code_blocks(self, code_blocks: List[CodeBlock]) -> CommandLineCodeResult:
210+
"""(Experimental) Execute the code blocks and return the result.
211+
212+
Args:
213+
code_blocks (List[CodeBlock]): The code blocks to execute.
214+
215+
Returns:
216+
CommandlineCodeResult: The result of the code execution."""
217+
218+
if len(code_blocks) == 0:
219+
raise ValueError("No code blocks to execute.")
220+
221+
outputs = []
222+
files = []
223+
last_exit_code = 0
224+
for code_block in code_blocks:
225+
lang = self.LANGUAGE_ALIASES.get(code_block.language.lower(), code_block.language.lower())
226+
if lang not in self.DEFAULT_EXECUTION_POLICY:
227+
outputs.append(f"Unsupported language {lang}\n")
228+
last_exit_code = 1
229+
break
230+
231+
execute_code = self.execution_policies.get(lang, False)
232+
code = silence_pip(code_block.code, lang)
233+
if lang in ["bash", "shell", "sh"]:
234+
code = "\n".join(["#!/bin/bash", code])
235+
236+
try:
237+
filename = _get_file_name_from_content(code, self._work_dir)
238+
except ValueError:
239+
outputs.append("Filename is not in the workspace")
240+
last_exit_code = 1
241+
break
242+
243+
if not filename:
244+
extension = self.LANGUAGE_FILE_EXTENSION.get(lang, lang)
245+
filename = f"tmp_code_{md5(code.encode()).hexdigest()}.{extension}"
246+
247+
code_path = self._work_dir / filename
248+
249+
exec_script = textwrap.dedent(
250+
"""
251+
if [ ! -d "{workspace}" ]; then
252+
mkdir {workspace}
253+
fi
254+
cat <<EOM >{code_path}\n
255+
{code}
256+
EOM
257+
chmod +x {code_path}"""
258+
)
259+
exec_script = exec_script.format(workspace=str(self._work_dir), code_path=code_path, code=code)
260+
stream(
261+
self._api_client.connect_get_namespaced_pod_exec,
262+
self._pod.metadata.name,
263+
self._pod.metadata.namespace,
264+
command=["/bin/sh", "-c", exec_script],
265+
container=self._container_name,
266+
stderr=True,
267+
stdin=False,
268+
stdout=True,
269+
tty=False,
270+
)
271+
272+
files.append(code_path)
273+
274+
if not execute_code:
275+
outputs.append(f"Code saved to {str(code_path)}\n")
276+
continue
277+
278+
resp = stream(
279+
self._api_client.connect_get_namespaced_pod_exec,
280+
self._pod.metadata.name,
281+
self._pod.metadata.namespace,
282+
command=["timeout", str(self._timeout), _cmd(lang), str(code_path)],
283+
container=self._container_name,
284+
stderr=True,
285+
stdin=False,
286+
stdout=True,
287+
tty=False,
288+
_preload_content=False,
289+
)
290+
291+
stdout_messages = []
292+
stderr_messages = []
293+
while resp.is_open():
294+
resp.update(timeout=1)
295+
if resp.peek_stderr():
296+
stderr_messages.append(resp.read_stderr())
297+
if resp.peek_stdout():
298+
stdout_messages.append(resp.read_stdout())
299+
outputs.extend(stdout_messages + stderr_messages)
300+
exit_code = resp.returncode
301+
resp.close()
302+
303+
if exit_code == 124:
304+
outputs.append("\n" + TIMEOUT_MSG)
305+
306+
last_exit_code = exit_code
307+
if exit_code != 0:
308+
break
309+
310+
code_file = str(files[0]) if files else None
311+
return CommandLineCodeResult(exit_code=last_exit_code, output="".join(outputs), code_file=code_file)
312+
313+
def stop(self) -> None:
314+
"""(Experimental) Stop the code executor."""
315+
self._cleanup()
316+
317+
def __enter__(self) -> Self:
318+
return self
319+
320+
def __exit__(
321+
self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
322+
) -> None:
323+
self.stop()

‎setup.py

+1
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,7 @@
107107
"cohere": ["cohere>=5.5.8"],
108108
"ollama": ["ollama>=0.3.3", "fix_busted_json>=0.0.18"],
109109
"bedrock": ["boto3>=1.34.149"],
110+
"kubernetes": ["kubernetes>=27.2.0"],
110111
}
111112

112113
setuptools.setup(

‎test/coding/test_kubernetes.commandline_code_executor.md

+44
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
# Test Environment for autogen.coding.kubernetes.PodCommandLineCodeExecutor
2+
3+
To test PodCommandLineCodeExecutor, the following environment is required.
4+
- kubernetes cluster config file
5+
- autogen package
6+
7+
## kubernetes cluster config file
8+
9+
kubernetes cluster config file, kubeconfig file's location should be set on environment variable `KUBECONFIG` or
10+
It must be located in the .kube/config path of your home directory.
11+
12+
For Windows, `C:\Users\<<user>>\.kube\config`,
13+
For Linux or MacOS, place the kubeconfig file in the `/home/<<user>>/.kube/config` directory.
14+
15+
## package install
16+
17+
Clone autogen github repository for package install and testing
18+
19+
Clone the repository with the command below.
20+
21+
before contribution
22+
```sh
23+
git clone -b k8s-code-executor https://github.com/questcollector/autogen.git
24+
```
25+
26+
after contribution
27+
```sh
28+
git clone https://github.com/microsoft/autogen.git
29+
```
30+
31+
install autogen with kubernetes >= 27.0.2
32+
33+
```sh
34+
cd autogen
35+
pip install .[kubernetes] -U
36+
```
37+
38+
## test execution
39+
40+
Perform the test with the following command
41+
42+
```sh
43+
pytest test/coding/test_kubernetes_commandline_code_executor.py
44+
```

‎test/coding/test_kubernetes_commandline_code_executor.py

+203
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,203 @@
1+
import importlib
2+
import os
3+
import sys
4+
from pathlib import Path
5+
6+
import pytest
7+
8+
from autogen.code_utils import TIMEOUT_MSG
9+
from autogen.coding.base import CodeBlock, CodeExecutor
10+
11+
try:
12+
from autogen.coding.kubernetes import PodCommandLineCodeExecutor
13+
14+
client = importlib.import_module("kubernetes.client")
15+
config = importlib.import_module("kubernetes.config")
16+
17+
kubeconfig = Path(".kube/config")
18+
if os.environ.get("KUBECONFIG", None):
19+
kubeconfig = Path(os.environ["KUBECONFIG"])
20+
elif sys.platform == "win32":
21+
kubeconfig = os.environ["userprofile"] / kubeconfig
22+
else:
23+
kubeconfig = os.environ["HOME"] / kubeconfig
24+
25+
if kubeconfig.is_file():
26+
config.load_config(config_file=str(kubeconfig))
27+
api_client = client.CoreV1Api()
28+
api_client.list_namespace()
29+
skip_kubernetes_tests = False
30+
else:
31+
skip_kubernetes_tests = True
32+
33+
pod_spec = client.V1Pod(
34+
metadata=client.V1ObjectMeta(
35+
name="abcd", namespace="default", annotations={"sidecar.istio.io/inject": "false"}
36+
),
37+
spec=client.V1PodSpec(
38+
restart_policy="Never",
39+
containers=[
40+
client.V1Container(
41+
args=["-c", "while true;do sleep 5; done"],
42+
command=["/bin/sh"],
43+
name="abcd",
44+
image="python:3.11-slim",
45+
env=[
46+
client.V1EnvVar(name="TEST", value="TEST"),
47+
client.V1EnvVar(
48+
name="POD_NAME",
49+
value_from=client.V1EnvVarSource(
50+
field_ref=client.V1ObjectFieldSelector(field_path="metadata.name")
51+
),
52+
),
53+
],
54+
)
55+
],
56+
),
57+
)
58+
except Exception:
59+
skip_kubernetes_tests = True
60+
61+
62+
@pytest.mark.skipif(skip_kubernetes_tests, reason="kubernetes not accessible")
63+
def test_create_default_pod_executor():
64+
with PodCommandLineCodeExecutor(namespace="default", kube_config_file=str(kubeconfig)) as executor:
65+
assert executor.timeout == 60
66+
assert executor.work_dir == Path("/workspace")
67+
assert executor._container_name == "autogen-code-exec"
68+
assert executor._pod.metadata.name.startswith("autogen-code-exec-")
69+
_test_execute_code(executor)
70+
71+
72+
@pytest.mark.skipif(skip_kubernetes_tests, reason="kubernetes not accessible")
73+
def test_create_node_pod_executor():
74+
with PodCommandLineCodeExecutor(
75+
image="node:22-alpine",
76+
namespace="default",
77+
work_dir="./app",
78+
timeout=30,
79+
kube_config_file=str(kubeconfig),
80+
execution_policies={"javascript": True},
81+
) as executor:
82+
assert executor.timeout == 30
83+
assert executor.work_dir == Path("./app")
84+
assert executor._container_name == "autogen-code-exec"
85+
assert executor._pod.metadata.name.startswith("autogen-code-exec-")
86+
assert executor.execution_policies["javascript"]
87+
88+
# Test single code block.
89+
code_blocks = [CodeBlock(code="console.log('hello world!')", language="javascript")]
90+
code_result = executor.execute_code_blocks(code_blocks)
91+
assert code_result.exit_code == 0 and "hello world!" in code_result.output and code_result.code_file is not None
92+
93+
# Test multiple code blocks.
94+
code_blocks = [
95+
CodeBlock(code="console.log('hello world!')", language="javascript"),
96+
CodeBlock(code="let a = 100 + 100; console.log(a)", language="javascript"),
97+
]
98+
code_result = executor.execute_code_blocks(code_blocks)
99+
assert (
100+
code_result.exit_code == 0
101+
and "hello world!" in code_result.output
102+
and "200" in code_result.output
103+
and code_result.code_file is not None
104+
)
105+
106+
# Test running code.
107+
file_lines = ["console.log('hello world!')", "let a = 100 + 100", "console.log(a)"]
108+
code_blocks = [CodeBlock(code="\n".join(file_lines), language="javascript")]
109+
code_result = executor.execute_code_blocks(code_blocks)
110+
assert (
111+
code_result.exit_code == 0
112+
and "hello world!" in code_result.output
113+
and "200" in code_result.output
114+
and code_result.code_file is not None
115+
)
116+
117+
118+
@pytest.mark.skipif(skip_kubernetes_tests, reason="kubernetes not accessible")
119+
def test_create_pod_spec_pod_executor():
120+
with PodCommandLineCodeExecutor(
121+
pod_spec=pod_spec, container_name="abcd", kube_config_file=str(kubeconfig)
122+
) as executor:
123+
assert executor.timeout == 60
124+
assert executor._container_name == "abcd"
125+
assert executor._pod.metadata.name == pod_spec.metadata.name
126+
assert executor._pod.metadata.namespace == pod_spec.metadata.namespace
127+
_test_execute_code(executor)
128+
129+
# Test bash script.
130+
if sys.platform not in ["win32"]:
131+
code_blocks = [CodeBlock(code="echo $TEST $POD_NAME", language="bash")]
132+
code_result = executor.execute_code_blocks(code_blocks)
133+
assert (
134+
code_result.exit_code == 0 and "TEST abcd" in code_result.output and code_result.code_file is not None
135+
)
136+
137+
138+
@pytest.mark.skipif(skip_kubernetes_tests, reason="kubernetes not accessible")
139+
def test_pod_executor_timeout():
140+
with PodCommandLineCodeExecutor(namespace="default", timeout=5, kube_config_file=str(kubeconfig)) as executor:
141+
assert executor.timeout == 5
142+
assert executor.work_dir == Path("/workspace")
143+
assert executor._container_name == "autogen-code-exec"
144+
assert executor._pod.metadata.name.startswith("autogen-code-exec-")
145+
# Test running code.
146+
file_lines = ["import time", "time.sleep(10)", "a = 100 + 100", "print(a)"]
147+
code_blocks = [CodeBlock(code="\n".join(file_lines), language="python")]
148+
code_result = executor.execute_code_blocks(code_blocks)
149+
assert code_result.exit_code == 124 and TIMEOUT_MSG in code_result.output and code_result.code_file is not None
150+
151+
152+
def _test_execute_code(executor: CodeExecutor) -> None:
153+
# Test single code block.
154+
code_blocks = [CodeBlock(code="import sys; print('hello world!')", language="python")]
155+
code_result = executor.execute_code_blocks(code_blocks)
156+
assert code_result.exit_code == 0 and "hello world!" in code_result.output and code_result.code_file is not None
157+
158+
# Test multiple code blocks.
159+
code_blocks = [
160+
CodeBlock(code="import sys; print('hello world!')", language="python"),
161+
CodeBlock(code="a = 100 + 100; print(a)", language="python"),
162+
]
163+
code_result = executor.execute_code_blocks(code_blocks)
164+
assert (
165+
code_result.exit_code == 0
166+
and "hello world!" in code_result.output
167+
and "200" in code_result.output
168+
and code_result.code_file is not None
169+
)
170+
171+
# Test bash script.
172+
if sys.platform not in ["win32"]:
173+
code_blocks = [CodeBlock(code="echo 'hello world!'", language="bash")]
174+
code_result = executor.execute_code_blocks(code_blocks)
175+
assert code_result.exit_code == 0 and "hello world!" in code_result.output and code_result.code_file is not None
176+
177+
# Test running code.
178+
file_lines = ["import sys", "print('hello world!')", "a = 100 + 100", "print(a)"]
179+
code_blocks = [CodeBlock(code="\n".join(file_lines), language="python")]
180+
code_result = executor.execute_code_blocks(code_blocks)
181+
assert (
182+
code_result.exit_code == 0
183+
and "hello world!" in code_result.output
184+
and "200" in code_result.output
185+
and code_result.code_file is not None
186+
)
187+
188+
# Test running code has filename.
189+
file_lines = ["# filename: test.py", "import sys", "print('hello world!')", "a = 100 + 100", "print(a)"]
190+
code_blocks = [CodeBlock(code="\n".join(file_lines), language="python")]
191+
code_result = executor.execute_code_blocks(code_blocks)
192+
print(code_result.code_file)
193+
assert (
194+
code_result.exit_code == 0
195+
and "hello world!" in code_result.output
196+
and "200" in code_result.output
197+
and code_result.code_file.find("test.py") > 0
198+
)
199+
200+
# Test error code.
201+
code_blocks = [CodeBlock(code="print(sys.platform)", language="python")]
202+
code_result = executor.execute_code_blocks(code_blocks)
203+
assert code_result.exit_code == 1 and "Traceback" in code_result.output and code_result.code_file is not None

‎website/docs/topics/code-execution/kubernetes-pod-commandline-code-executor.ipynb

+773
Large diffs are not rendered by default.

0 commit comments

Comments
 (0)
Please sign in to comment.