Add ability to define custom runtime classes (#6955)

Co-authored-by: Engel Nyst <enyst@users.noreply.github.com>

openhands/runtime/README.md

@@ -52,6 +52,11 @@ There are currently four implementations:
 * Modal (uses the Modal API)
 * Runloop (uses the Runloop API)
 
+You may also add your own `Runtime` subclass to the classpath and configure it like this:
+
+```toml
+runtime = "app.my.CustomRuntime"
+```
 
 ## Workflow Description
 

openhands/runtime/__init__.py

@@ -1,41 +1,54 @@
-from openhands.core.logger import openhands_logger as logger
+from typing import Type
+
+from openhands.runtime.base import Runtime
 from openhands.runtime.impl.daytona.daytona_runtime import DaytonaRuntime
 from openhands.runtime.impl.docker.docker_runtime import (
     DockerRuntime,
 )
-from openhands.runtime.impl.e2b.sandbox import E2BBox
+from openhands.runtime.impl.e2b.e2b_runtime import E2BRuntime
 from openhands.runtime.impl.local.local_runtime import LocalRuntime
 from openhands.runtime.impl.modal.modal_runtime import ModalRuntime
 from openhands.runtime.impl.remote.remote_runtime import RemoteRuntime
 from openhands.runtime.impl.runloop.runloop_runtime import RunloopRuntime
+from openhands.utils.import_utils import get_impl
+
+# mypy: disable-error-code="type-abstract"
+_DEFAULT_RUNTIME_CLASSES: dict[str, Type[Runtime]] = {
+    'eventstream': DockerRuntime,
+    'docker': DockerRuntime,
+    'e2b': E2BRuntime,
+    'remote': RemoteRuntime,
+    'modal': ModalRuntime,
+    'runloop': RunloopRuntime,
+    'local': LocalRuntime,
+    'daytona': DaytonaRuntime,
+}
 
 
-def get_runtime_cls(name: str):
-    # Local imports to avoid circular imports
-    if name == 'eventstream' or name == 'docker':
-        return DockerRuntime
-    elif name == 'e2b':
-        return E2BBox
-    elif name == 'remote':
-        return RemoteRuntime
-    elif name == 'modal':
-        logger.debug('Using ModalRuntime')
-        return ModalRuntime
-    elif name == 'runloop':
-        return RunloopRuntime
-    elif name == 'local':
-        return LocalRuntime
-    elif name == 'daytona':
-        return DaytonaRuntime
-    else:
-        raise ValueError(f'Runtime {name} not supported')
+def get_runtime_cls(name: str) -> Type[Runtime]:
+    """
+    If name is one of the predefined runtime names (e.g. 'docker'), return its class.
+    Otherwise attempt to resolve name as subclass of Runtime and return it.
+    Raise on invalid selections.
+    """
+    if name in _DEFAULT_RUNTIME_CLASSES:
+        return _DEFAULT_RUNTIME_CLASSES[name]
+    try:
+        return get_impl(Runtime, name)
+    except Exception as e:
+        known_keys = _DEFAULT_RUNTIME_CLASSES.keys()
+        raise ValueError(
+            f'Runtime {name} not supported, known are: {known_keys}'
+        ) from e
 
 
 __all__ = [
-    'E2BBox',
+    'Runtime',
+    'E2BRuntime',
     'RemoteRuntime',
     'ModalRuntime',
     'RunloopRuntime',
     'DockerRuntime',
+    'DaytonaRuntime',
     'get_runtime_cls',
 ]

openhands/server/session/agent_session.py

@@ -221,6 +221,7 @@ class AgentSession:
             plugins=agent.sandbox_plugins,
             status_callback=self._status_callback,
             headless_mode=False,
+            attach_to_existing=False,
             env_vars=env_vars,
             **kwargs,
         )