From 201ecdc9ee45bc62752d2ff5d136341fa20dbf6d Mon Sep 17 00:00:00 2001 From: blob42 Date: Thu, 2 Mar 2023 22:33:17 +0100 Subject: [PATCH] fix run and exec_run default commands, actually use gVisor - run and exec_run need a separate default command. Run usually executes a script while exec_run simulates an interactive session. The image templates and run funcs have been upgraded to handle both types of commands. - test: make docker tests run when docker is installed and docker lib avaialble. - test that runsc runtime is used by default when gVisor is installed. (manually removing gVisor skips the test) --- langchain/utilities/docker/__init__.py | 42 ++++++++++++++++++-------- langchain/utilities/docker/images.py | 28 ++++++++--------- tests/unit_tests/test_docker.py | 32 +++++++++++++++++++- 3 files changed, 75 insertions(+), 27 deletions(-) diff --git a/langchain/utilities/docker/__init__.py b/langchain/utilities/docker/__init__.py index 59f01706..d5a97d67 100644 --- a/langchain/utilities/docker/__init__.py +++ b/langchain/utilities/docker/__init__.py @@ -158,11 +158,10 @@ class DockerSocket: return chunks -def _default_params() -> Dict: - return { - # the only required parameter to be able to attach. - 'stdin_open': True, - } +_default_params = { + # the only required parameter to be able to attach. + 'stdin_open': True, + } def _get_command(query: str, **kwargs: Dict) -> str: """Build an escaped command from a query string and keyword arguments.""" @@ -172,6 +171,7 @@ def _get_command(query: str, **kwargs: Dict) -> str: return cmd +#TEST: gVisor class DockerWrapper(BaseModel, extra=Extra.allow): """Executes arbitrary payloads and returns the output. @@ -187,6 +187,12 @@ class DockerWrapper(BaseModel, extra=Extra.allow): image: Union[str, Type[BaseImage]] = Field(default_factory=Shell,skip=True) from_env: Optional[bool] = Field(default=True, skip=True) + _default_params: dict = Field( + default = { + # needed to attach stdin + 'stdin_open': True, + } + ) # @property # def image_name(self) -> str: @@ -199,6 +205,10 @@ class DockerWrapper(BaseModel, extra=Extra.allow): if self.from_env: self._docker_client = docker.from_env() + if gvisor_runtime_available(docker.from_env()): + self._params['runtime'] = 'runsc' + + # if not isinstance(self.image, str) and issubclass(self.image, BaseImage): # self._params = {**self._params, **self.image().dict()} @@ -244,7 +254,7 @@ class DockerWrapper(BaseModel, extra=Extra.allow): _image = get_image_template(image) if isinstance(_image, str): # user wants a custom image, we should use default params - values["_params"] = {**_default_params(), **{'image': image}} + values["_params"] = {**_default_params, **{'image': image}} else: # user wants a pre registered image, we should use the image params values["_params"] = _image().dict() @@ -267,6 +277,7 @@ class DockerWrapper(BaseModel, extra=Extra.allow): return values + #FIX: default shell command should be different in run vs exec mode def run(self, query: str, **kwargs: Any) -> str: """Run arbitrary shell command inside a container. @@ -287,6 +298,7 @@ class DockerWrapper(BaseModel, extra=Extra.allow): del kwargs['image'] cmd = _get_command(query, **kwargs) kwargs.pop('default_command', None) + kwargs.pop('stdin_command', None) args['command'] = cmd # print(f"args: {args}") @@ -351,9 +363,13 @@ class DockerWrapper(BaseModel, extra=Extra.allow): kwargs = {**self._params, **kwargs} - if 'default_command' in kwargs: - kwargs['command'] = shlex.join([*kwargs['default_command'],query]) - del kwargs['default_command'] + kwargs.pop('default_command', None) + + # exec_run requires flags for stdin + if 'stdin_command' in kwargs: + assert isinstance(kwargs['stdin_command'], list) + kwargs['command'] = shlex.join(kwargs['stdin_command']) + del kwargs['stdin_command'] # kwargs.pop('default_command', None) # kwargs['command'] = cmd @@ -364,6 +380,8 @@ class DockerWrapper(BaseModel, extra=Extra.allow): # TODO: handle both output mode for tty=True/False logger.debug(f"running command {kwargs['command']}") logger.debug(f"creating container with params {kwargs}") + + #FIX: create container with base interactive command container = self._docker_client.containers.create(**kwargs) container.start() @@ -403,15 +421,15 @@ class DockerWrapper(BaseModel, extra=Extra.allow): # output is stored in a list of tuples (stream_type, payload) df = pd.DataFrame(output, columns=['stream_type', 'payload']) - df['payload'] = df['payload'].apply(lambda x: x.decode('utf-8')).apply(lambda x: x.strip()) + df['payload'] = df['payload'].apply(lambda x: x.decode('utf-8')) df['stream_type'] = df['stream_type'].apply(lambda x: 'stdout' if x == 1 else 'stderr') payload = df.groupby('stream_type')['payload'].apply(''.join).to_dict() logger.debug(f"payload: {payload}") #NOTE: stderr might just contain the prompt if 'stdout' in payload and 'stderr' in payload and with_stderr: - return f"STDOUT:\n {payload['stdout']}\nSTDERR:\n {payload['stderr']}" + return f"STDOUT:\n {payload['stdout'].strip()}\nSTDERR:\n {payload['stderr']}" elif 'stderr' in payload and not 'stdout' in payload: return f"STDERR: {payload['stderr']}" else: - return payload['stdout'] + return payload['stdout'].strip() diff --git a/langchain/utilities/docker/images.py b/langchain/utilities/docker/images.py index 882fd626..0c36c038 100644 --- a/langchain/utilities/docker/images.py +++ b/langchain/utilities/docker/images.py @@ -13,16 +13,14 @@ class BaseImage(BaseModel, extra=Extra.forbid): name: str tag: Optional[str] = 'latest' default_command: Optional[List[str]] = None + stdin_command: Optional[List[str]] = None def dict(self, *args, **kwargs): """Override the dict method to add the image name.""" d = super().dict(*args, **kwargs) del d['name'] del d['tag'] - # del d['default_command'] d['image'] = self.image_name - # if self.default_command: - # d['command'] = self.default_command return d @property @@ -46,16 +44,23 @@ class Shell(BaseImage): or by passing the full path to the shell binary. """ name: str = 'alpine' - default_command: List[str] = [ShellTypes.sh.value, '-s'] + default_command: List[str] = [ShellTypes.sh.value, '-c'] + stdin_command: List[str] = [ShellTypes.sh.value, '-i'] @validator('default_command') - def validate_shell(cls, value: str) -> str: + def validate_default_command(cls, value: str) -> str: + """Validate shell type.""" + val = getattr(ShellTypes, value, None) + if val: + return val.value + return value + + @validator('stdin_command') + def validate_stdin_command(cls, value: str) -> str: """Validate shell type.""" val = getattr(ShellTypes, value, None) if val: return val.value - # elif value in [v.value for v in list(ShellTypes.__members__.values())]: - # print(f"docker: overriding shell binary to: {value}") return value # example using base image to construct python image @@ -66,13 +71,8 @@ class Python(BaseImage): stdin open. """ name: str = 'python' - default_command: List[str] = ['python3', '-iq'] - - def __setattr__(self, name, value): - if name == 'default_command': - raise AttributeError(f'running this image with {self.default_command}' - ' is necessary to keep stdin open.') - super().__setattr__(name, value) + default_command: List[str] = ['python3', '-c'] + stdin_command: List[str] = ['python3', '-iq'] def get_image_template(image_name: str = 'shell') -> Union[str, Type[BaseImage]]: diff --git a/tests/unit_tests/test_docker.py b/tests/unit_tests/test_docker.py index 0f582344..58d4bdfc 100644 --- a/tests/unit_tests/test_docker.py +++ b/tests/unit_tests/test_docker.py @@ -1,6 +1,8 @@ """Test the docker wrapper utility.""" import pytest +import importlib +from typing import Any from langchain.utilities.docker import DockerWrapper, \ gvisor_runtime_available, _default_params from unittest.mock import MagicMock @@ -17,11 +19,31 @@ def docker_installed() -> bool: return True +def gvisor_installed() -> bool: + """return true if gvisor local runtime is installed""" + try: + docker_lib = importlib.import_module('docker') + client = docker_lib.from_env() + return gvisor_runtime_available(client) + except ImportError: + return False + + return False + + + +def docker_lib_installed() -> bool: + return importlib.util.find_spec('docker') is not None -@pytest.mark.skipif(not docker_installed(), reason="docker not installed") +def skip_docker_tests() -> bool: + return not docker_installed() or not docker_lib_installed() + + +@pytest.mark.skipif(skip_docker_tests(), reason="docker not installed") class TestDockerUtility: + def test_default_image(self) -> None: """Test running a command with the default alpine image.""" docker = DockerWrapper() @@ -62,6 +84,14 @@ class TestDockerUtility: mock_client.info.return_value = {'Runtimes': {'runc': {'path': 'runc'}}} assert not gvisor_runtime_available(mock_client) + + @pytest.mark.skipif(not gvisor_installed(), reason="gvisor not installed") + def test_run_with_runtime_runsc(self) -> None: + docker = DockerWrapper(image='shell') + output = docker.run('dmesg') + assert output.find('gVisor') != -1 + + def test_socket_read_timeout(self) -> None: """Test socket read timeout.""" docker = DockerWrapper(image='python', default_command=['python'])