Fix: Circular imports between controllers and orchestrator

src/aleph/vm/orchestrator/conf.py → src/aleph/vm/conf.py

@@ -11,7 +11,7 @@
 
 from pydantic import BaseSettings, Field
 
-from aleph.vm.orchestrator.utils import is_command_available
+from aleph.vm.utils import is_command_available
 
 logger = logging.getLogger(__name__)
 

src/aleph/vm/controllers/firecracker/executable.py

@@ -15,13 +15,13 @@
 from aleph_message.models import ExecutableContent, ItemHash
 from aleph_message.models.execution.environment import MachineResources
 
+from aleph.vm.conf import settings
+from aleph.vm.controllers.firecracker.snapshots import CompressedDiskVolumeSnapshot
 from aleph.vm.guest_api.__main__ import run_guest_api
 from aleph.vm.hypervisors.firecracker.microvm import FirecrackerConfig, MicroVM
-from aleph.vm.orchestrator.conf import settings
-from aleph.vm.orchestrator.network.firewall import teardown_nftables_for_vm
-from aleph.vm.orchestrator.network.interfaces import TapInterface
-from aleph.vm.orchestrator.snapshots import CompressedDiskVolumeSnapshot
-from aleph.vm.orchestrator.storage import get_volume_path
+from aleph.vm.network.firewall import teardown_nftables_for_vm
+from aleph.vm.network.interfaces import TapInterface
+from aleph.vm.storage import get_volume_path
 
 try:
     import psutil  # type: ignore [no-redef]

src/aleph/vm/controllers/firecracker/instance.py

@@ -10,6 +10,7 @@
 from aleph_message.models import ItemHash
 from aleph_message.models.execution.environment import MachineResources
 
+from aleph.vm.conf import settings
 from aleph.vm.hypervisors.firecracker.config import (
     BootSource,
     Drive,
@@ -19,26 +20,22 @@
     Vsock,
 )
 from aleph.vm.hypervisors.firecracker.microvm import setfacl
-from aleph.vm.orchestrator.conf import settings
-from aleph.vm.orchestrator.network.interfaces import TapInterface
-from aleph.vm.orchestrator.snapshots import (
-    CompressedDiskVolumeSnapshot,
-    DiskVolume,
-    DiskVolumeSnapshot,
-)
-from aleph.vm.orchestrator.storage import (
+from aleph.vm.network.interfaces import TapInterface
+from aleph.vm.utils import (
+    HostNotFoundError,
     NotEnoughDiskSpace,
     check_disk_space,
-    create_devmapper,
-    create_volume_file,
+    ping,
+    run_in_subprocess,
 )
-from aleph.vm.orchestrator.utils import HostNotFoundError, ping, run_in_subprocess
 
+from ...storage import create_devmapper, create_volume_file
 from .executable import (
     AlephFirecrackerExecutable,
     AlephFirecrackerResources,
     BaseConfiguration,
 )
+from .snapshots import CompressedDiskVolumeSnapshot, DiskVolume, DiskVolumeSnapshot
 
 logger = logging.getLogger(__name__)
 

src/aleph/vm/controllers/firecracker/program.py

@@ -15,6 +15,7 @@
 from aleph_message.models.execution.base import Encoding
 from aleph_message.models.execution.environment import MachineResources
 
+from aleph.vm.conf import settings
 from aleph.vm.hypervisors.firecracker.config import (
     BootSource,
     Drive,
@@ -24,10 +25,9 @@
     Vsock,
 )
 from aleph.vm.hypervisors.firecracker.microvm import RuntimeConfiguration, setfacl
-from aleph.vm.orchestrator.conf import settings
-from aleph.vm.orchestrator.network.interfaces import TapInterface
-from aleph.vm.orchestrator.storage import get_code_path, get_data_path, get_runtime_path
-from aleph.vm.orchestrator.utils import MsgpackSerializable
+from aleph.vm.network.interfaces import TapInterface
+from aleph.vm.storage import get_code_path, get_data_path, get_runtime_path
+from aleph.vm.utils import MsgpackSerializable
 
 from .executable import (
     AlephFirecrackerExecutable,

src/aleph/vm/orchestrator/snapshot_manager.py → src/aleph/vm/controllers/firecracker/snapshot_manager.py

@@ -7,8 +7,9 @@
 from aleph_message.models import ItemHash
 from schedule import Job, Scheduler
 
-from .conf import settings
-from .models import VmExecution
+from aleph.vm.conf import settings
+from aleph.vm.orchestrator.models import VmExecution
+
 from .snapshots import CompressedDiskVolumeSnapshot
 
 logger = logging.getLogger(__name__)
@@ -99,15 +100,14 @@ async def start_for(self, execution: VmExecution, frequency: Optional[int] = Non
             msg = "Snapshots are not implemented for programs."
             raise NotImplementedError(msg)
 
-        if not frequency:
-            frequency = settings.SNAPSHOT_FREQUENCY
+        default_frequency = frequency or settings.SNAPSHOT_FREQUENCY
 
         vm_hash = execution.vm_hash
         snapshot_execution = SnapshotExecution(
             scheduler=self._scheduler,
             vm_hash=vm_hash,
             execution=execution,
-            frequency=frequency,
+            frequency=default_frequency,
         )
         self.executions[vm_hash] = snapshot_execution
         await snapshot_execution.start()

src/aleph/vm/orchestrator/snapshots.py → src/aleph/vm/controllers/firecracker/snapshots.py

@@ -4,8 +4,8 @@
 
 from aleph_message.models import ItemHash
 
-from .conf import SnapshotCompressionAlgorithm
-from .storage import compress_volume_snapshot, create_volume_snapshot
+from aleph.vm.conf import SnapshotCompressionAlgorithm
+from aleph.vm.storage import compress_volume_snapshot, create_volume_snapshot
 
 logger = logging.getLogger(__name__)
 

src/aleph/vm/orchestrator/network/firewall.py → src/aleph/vm/network/firewall.py

@@ -4,7 +4,7 @@
 
 from nftables import Nftables
 
-from aleph.vm.orchestrator.conf import settings
+from aleph.vm.conf import settings
 
 from .interfaces import TapInterface
 

src/aleph/vm/orchestrator/network/hostnetwork.py → src/aleph/vm/network/hostnetwork.py

@@ -5,7 +5,7 @@
 
 from aleph_message.models import ItemHash
 
-from aleph.vm.orchestrator.conf import IPv6AllocationPolicy
+from aleph.vm.conf import IPv6AllocationPolicy
 from aleph.vm.orchestrator.vm.vm_type import VmType
 
 from .firewall import initialize_nftables, setup_nftables_for_vm, teardown_nftables

src/aleph/vm/orchestrator/network/ndp_proxy.py → src/aleph/vm/network/ndp_proxy.py

@@ -14,7 +14,7 @@
 from ipaddress import IPv6Network
 from pathlib import Path
 
-from aleph.vm.orchestrator.utils import run_in_subprocess
+from aleph.vm.utils import run_in_subprocess
 
 logger = logging.getLogger(__name__)
 

src/aleph/vm/orchestrator/__init__.py

@@ -1,5 +1,4 @@
 from . import (
-    conf,
     messages,
     metrics,
     models,
@@ -9,10 +8,8 @@
     resources,
     run,
     status,
-    storage,
     supervisor,
     tasks,
-    utils,
     version,
     views,
     vm,
@@ -21,7 +18,6 @@
 __version__ = version.__version__
 
 __all__ = (
-    "conf",
     "messages",
     "metrics",
     "models",
@@ -31,10 +27,8 @@
     "resources",
     "run",
     "status",
-    "storage",
     "supervisor",
     "tasks",
-    "utils",
     "version",
     "views",
     "vm",

src/aleph/vm/orchestrator/__main__.py

@@ -1,4 +1,4 @@
-from . import cli
+from .cli import main
 
 if __name__ == "__main__":
-    cli.main()
+    main()

src/aleph/vm/orchestrator/cli.py

@@ -20,8 +20,8 @@
 import alembic.config
 from aleph_message.models import ItemHash
 
+from ..conf import ALLOW_DEVELOPER_SSH_KEYS, make_db_url, settings
 from . import metrics, supervisor
-from .conf import ALLOW_DEVELOPER_SSH_KEYS, make_db_url, settings
 from .pubsub import PubSub
 from .run import run_code_on_event, run_code_on_request, start_persistent_vm
 

src/aleph/vm/orchestrator/messages.py

@@ -5,7 +5,7 @@
 from aiohttp.web_exceptions import HTTPNotFound, HTTPServiceUnavailable
 from aleph_message.models import ExecutableMessage, ItemHash, MessageType
 
-from .storage import get_latest_amend, get_message
+from aleph.vm.storage import get_latest_amend, get_message
 
 
 async def try_get_message(ref: str) -> ExecutableMessage:

src/aleph/vm/orchestrator/metrics.py

@@ -14,7 +14,7 @@
     from sqlalchemy.ext.declarative import declarative_base
 
 
-from .conf import make_db_url, settings
+from ..conf import make_db_url, settings
 
 Session: sessionmaker
 

src/aleph/vm/orchestrator/migrations/env.py

@@ -1,7 +1,7 @@
 from alembic import context
 from sqlalchemy import create_engine
 
-from aleph.vm.orchestrator.conf import make_db_url
+from aleph.vm.conf import make_db_url
 
 # Auto-generate migrations
 from aleph.vm.orchestrator.metrics import Base

src/aleph/vm/orchestrator/migrations/versions/0001_bbb12a12372e_execution_records.py

@@ -12,7 +12,7 @@
 from sqlalchemy import create_engine
 from sqlalchemy.engine import reflection
 
-from aleph.vm.orchestrator.conf import make_db_url
+from aleph.vm.conf import make_db_url
 
 revision = "bbb12a12372e"
 down_revision = None

src/aleph/vm/orchestrator/models.py

@@ -21,16 +21,16 @@
     AlephFirecrackerResources,
     AlephProgramResources,
 )
+from aleph.vm.network.interfaces import TapInterface
+from aleph.vm.utils import create_task_log_exceptions, dumps_for_json
 
-from .conf import settings
+from ..conf import settings
 from .metrics import ExecutionRecord, save_execution_data, save_record
-from .network.interfaces import TapInterface
 from .pubsub import PubSub
-from .utils import create_task_log_exceptions, dumps_for_json
 from .vm import AlephFirecrackerInstance
 
 if TYPE_CHECKING:
-    from .snapshot_manager import SnapshotManager
+    from aleph.vm.controllers.firecracker.snapshot_manager import SnapshotManager
 
 
 logger = logging.getLogger(__name__)

src/aleph/vm/orchestrator/pool.py

@@ -6,10 +6,11 @@
 from aleph_message.models import ExecutableMessage, ItemHash
 from aleph_message.models.execution.instance import InstanceContent
 
-from .conf import settings
+from aleph.vm.controllers.firecracker.snapshot_manager import SnapshotManager
+
+from ..conf import settings
+from ..network.hostnetwork import Network, make_ipv6_allocator
 from .models import ExecutableContent, VmExecution
-from .network.hostnetwork import Network, make_ipv6_allocator
-from .snapshot_manager import SnapshotManager
 from .vm.vm_type import VmType
 
 logger = logging.getLogger(__name__)

src/aleph/vm/orchestrator/reactor.py

@@ -4,9 +4,10 @@
 from aleph_message.models import AlephMessage
 from aleph_message.models.execution.environment import Subscription
 
+from aleph.vm.utils import create_task_log_exceptions
+
 from .pubsub import PubSub
 from .run import run_code_on_event
-from .utils import create_task_log_exceptions
 
 logger = logging.getLogger(__name__)
 

src/aleph/vm/orchestrator/resources.py

@@ -9,7 +9,7 @@
 from aleph_message.models.execution.environment import CpuProperties
 from pydantic import BaseModel, Field
 
-from .conf import settings
+from ..conf import settings
 
 
 class Period(BaseModel):

src/aleph/vm/orchestrator/run.py

@@ -15,13 +15,13 @@
     VmSetupError,
 )
 from aleph.vm.hypervisors.firecracker.microvm import MicroVMFailedInit
+from aleph.vm.utils import HostNotFoundError
 
-from .conf import settings
+from ..conf import settings
 from .messages import load_updated_message
 from .models import VmExecution
 from .pool import VmPool
 from .pubsub import PubSub
-from .utils import HostNotFoundError
 
 logger = logging.getLogger(__name__)
 

src/aleph/vm/orchestrator/status.py

@@ -7,7 +7,7 @@
 
 from aiohttp import ClientResponseError, ClientSession
 
-from .conf import settings
+from ..conf import settings
 
 logger = logging.getLogger(__name__)
 

src/aleph/vm/orchestrator/supervisor.py

@@ -12,8 +12,8 @@
 
 from aiohttp import web
 
-from . import metrics
-from .conf import settings
+from ..conf import settings
+from .metrics import create_tables, setup_engine
 from .resources import about_system_usage
 from .run import pool
 from .tasks import start_watch_for_messages_task, stop_watch_for_messages_task
@@ -85,8 +85,8 @@ def run():
     app["secret_token"] = secret_token
     print(f"Login to /about pages {protocol}://{hostname}/about/login?token={secret_token}")
 
-    engine = metrics.setup_engine()
-    metrics.create_tables(engine)
+    engine = setup_engine()
+    create_tables(engine)
 
     try:
         if settings.WATCH_FOR_MESSAGES:

src/aleph/vm/orchestrator/tasks.py

@@ -12,11 +12,12 @@
 from aleph_message.models import AlephMessage, ItemHash, ProgramMessage, parse_message
 from yarl import URL
 
-from .conf import settings
+from aleph.vm.utils import create_task_log_exceptions
+
+from ..conf import settings
 from .messages import load_updated_message
 from .pubsub import PubSub
 from .reactor import Reactor
-from .utils import create_task_log_exceptions
 
 logger = logging.getLogger(__name__)
 

src/aleph/vm/orchestrator/views/__init__.py

@@ -14,25 +14,25 @@
 from aleph_message.models import ItemHash
 from pydantic import ValidationError
 
+from aleph.vm.conf import settings
 from aleph.vm.controllers.firecracker.executable import (
     ResourceDownloadError,
     VmSetupError,
 )
 from aleph.vm.controllers.firecracker.program import FileTooLargeError
 from aleph.vm.hypervisors.firecracker.microvm import MicroVMFailedInit
 from aleph.vm.orchestrator import status
-from aleph.vm.orchestrator.conf import settings
 from aleph.vm.orchestrator.metrics import get_execution_records
 from aleph.vm.orchestrator.pubsub import PubSub
 from aleph.vm.orchestrator.resources import Allocation
 from aleph.vm.orchestrator.run import pool, run_code_on_request, start_persistent_vm
-from aleph.vm.orchestrator.utils import (
+from aleph.vm.orchestrator.version import __version__
+from aleph.vm.utils import (
     HostNotFoundError,
     b32_to_b16,
     dumps_for_json,
     get_ref_from_dns,
 )
-from aleph.vm.orchestrator.version import __version__
 from packaging.version import InvalidVersion, Version
 
 logger = logging.getLogger(__name__)