Merge branch 'main' into tobitege/dyn-runtime-init

openhands/runtime/__init__.py

@@ -1,25 +1,78 @@
+from pydantic import BaseModel
+
 from openhands.runtime.e2b.sandbox import E2BBox
 
 
+class RuntimeInfo(BaseModel):
+    module: str
+    class_name: str
+
+
+# Core runtimes as default
+_registered_runtimes: dict[str, RuntimeInfo] = {
+    'eventstream': RuntimeInfo(module='client', class_name='EventStreamRuntime'),
+    'e2b': RuntimeInfo(module='e2b', class_name='E2bRuntime'),
+    'remote': RuntimeInfo(module='remote', class_name='RemoteRuntime'),
+}
+
+
+def register_runtime(name: str, module: str, class_name: str):
+    """
+    Registers a new runtime with the given name, module, and class name.
+
+    This function allows you to add a new runtime to the system. It takes the name
+    of the runtime, the module where the runtime class is defined, and the class name
+    of the runtime. The runtime information is stored in the _registered_runtimes
+    dictionary, which maps the runtime name to its corresponding RuntimeInfo object.
+
+    Example:
+    >>> register_runtime("new_runtime", "new_module", "NewRuntime")
+    """
+    _registered_runtimes[name] = RuntimeInfo(module=module, class_name=class_name)
+
+
 def get_runtime_cls(name: str):
-    # Local imports to avoid circular imports
-    if name == 'eventstream':
-        from openhands.runtime.client.runtime import EventStreamRuntime
+    """
+    Returns the runtime class based on the given name using the registered runtime information.
 
-        return EventStreamRuntime
-    elif name == 'e2b':
-        from openhands.runtime.e2b.runtime import E2BRuntime
+    This function dynamically imports and returns the runtime class corresponding to the
+    provided name. It uses the information stored in the _registered_runtimes dictionary
+    to determine the module and class name for the requested runtime.
 
-        return E2BRuntime
-    elif name == 'remote':
-        from openhands.runtime.remote.runtime import RemoteRuntime
+    The function performs the following steps:
+    1. Retrieves the RuntimeInfo object for the given name from _registered_runtimes.
+    2. Dynamically imports the module containing the runtime class.
+    3. Retrieves the runtime class from the imported module using getattr.
 
-        return RemoteRuntime
-    else:
+    If the runtime is not found or cannot be imported, it raises a ValueError.
+
+    Args:
+        name (str): The name of the runtime to retrieve.
+
+    Returns:
+        type: The runtime class corresponding to the given name.
+
+    Raises:
+        ValueError: If the specified runtime is not supported or cannot be imported.
+
+    Example:
+        >>> runtime_cls = get_runtime_cls("eventstream")
+        >>> runtime_instance = runtime_cls(config, event_stream)
+    """
+    import importlib
+
+    try:
+        runtime_info = _registered_runtimes[name]
+        module = importlib.import_module(
+            f'openhands.runtime.{runtime_info.module}.runtime'
+        )
+        return getattr(module, runtime_info.class_name)
+    except (ImportError, AttributeError, KeyError):
         raise ValueError(f'Runtime {name} not supported')
 
 
 __all__ = [
     'E2BBox',
     'get_runtime_cls',
+    'register_runtime',
 ]

tests/unit/test_dyn_runtime.py

@@ -0,0 +1,40 @@
+from unittest.mock import patch
+
+import pytest
+
+from openhands.runtime import (
+    RuntimeInfo,
+    _registered_runtimes,
+    get_runtime_cls,
+    register_runtime,
+)
+
+
+def test_register_and_get_runtime():
+    # Register a new runtime
+    register_runtime('test_runtime', 'test_module', 'TestRuntime')
+
+    # Check if it's registered
+    assert 'test_runtime' in _registered_runtimes
+    assert _registered_runtimes['test_runtime'] == RuntimeInfo(
+        module='test_module', class_name='TestRuntime'
+    )
+
+    # Mock the import process
+    with patch('importlib.import_module') as mock_import:
+        mock_module = mock_import.return_value
+        mock_module.TestRuntime = 'mock_runtime_class'  # Mock the runtime class
+
+        # Get the runtime class
+        runtime_cls = get_runtime_cls('test_runtime')
+
+        # Check if the correct module was imported
+        mock_import.assert_called_once_with('openhands.runtime.test_module.runtime')
+
+        # Check if the correct class was returned
+        assert runtime_cls == 'mock_runtime_class'
+
+
+def test_get_runtime_not_found():
+    with pytest.raises(ValueError, match='Runtime unknown_runtime not supported'):
+        get_runtime_cls('unknown_runtime')