feat: load all dependencies in worker

Author: MHajoha

questionpy_server/worker/__init__.py

@@ -2,18 +2,21 @@
 #  The QuestionPy Server is free software released under terms of the MIT license. See LICENSE.md.
 #  (c) Technische Universität Berlin, innoCampus <info@isis.tu-berlin.de>
 from abc import ABC, abstractmethod
+from collections.abc import Mapping
 from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
 from typing import TypedDict, TypeVar, Unpack
 
 from pydantic import BaseModel
 
+from questionpy_common import PackageNamespaceAndShortName
 from questionpy_common.api.attempt import AttemptModel, AttemptScoredModel, AttemptStartedModel
 from questionpy_common.api.question import LmsPermissions
 from questionpy_common.elements import OptionsFormDefinition
 from questionpy_common.environment import PackagePermissions, RequestInfo
 from questionpy_common.manifest import PackageFile
+from questionpy_server.dependencies import SolutionAndLocation
 from questionpy_server.models import LoadedPackage, QuestionCreated
 from questionpy_server.utils.manifest import Manifest
 from questionpy_server.worker.runtime.messages import MessageToServer, MessageToWorker
@@ -55,15 +58,22 @@ class PackageFileData:
 class WorkerArgs(TypedDict):
     name: str
     """A unique name given to the worker by its pool."""
+
     package: PackageLocation
     """The main package that the worker should load when [start][questionpy_server.worker.Worker.start] is called."""
+
     worker_home: Path
     """An existing directory owned by the worker, with the same lifetime as the worker."""
+
     permissions: PackagePermissions
     """The package permissions."""
+
     environment_variables: dict[str, str]
     """Environment variables to be set in the worker."""
 
+    dependencies: Mapping[PackageNamespaceAndShortName, SolutionAndLocation]
+    """All resolved dependencies in the root package's tree. Does not include the root package itself."""
+
 
 class Worker(ABC):
     """Interface for worker implementations."""
@@ -75,6 +85,7 @@ def __init__(self, **kwargs: Unpack[WorkerArgs]) -> None:
         self.worker_home = kwargs["worker_home"]
         self.permissions = kwargs["permissions"]
         self.environment_variables = kwargs["environment_variables"]
+        self._dependencies = kwargs["dependencies"]
 
         self.state = WorkerState.NOT_RUNNING
         self.loaded_packages: list[LoadedPackage] = []

questionpy_server/worker/impl/_base.py

@@ -112,7 +112,7 @@ async def _initialize(self) -> None:
 
     async def _load_package(self, package_location: PackageLocation, *, main: bool) -> None:
         loaded = await self.send_and_wait_for_response(
-            LoadQPyPackage(location=package_location, main=main),
+            LoadQPyPackage(location=package_location, main=main, dependencies=self._dependencies),
             LoadQPyPackage.Response,
             self.permissions.bootstrap_timeout,
         )
@@ -373,7 +373,7 @@ class LimitTimeUsageMixin(Worker, ABC):
     the cpu limit in real time.
     """
 
-    _real_time_limit_factor = 3
+    _real_time_limit_factor = 1000
 
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(**kwargs)

questionpy_server/worker/pool.py

@@ -17,6 +17,8 @@
 from questionpy_common.environment import PackagePermissions
 from questionpy_common.error import QPyBaseError
 from questionpy_server.dependencies import WorkerDependencyResolver
+from questionpy_server.package import Package
+from questionpy_server.utils.manifest import ComparableManifest, read_manifest_from_location
 from questionpy_server.worker.impl.subprocess import SubprocessWorker
 from questionpy_server.worker.runtime.package_location import (
     DirPackageLocation,
@@ -111,7 +113,7 @@ def _memory_available(self, required_memory: int) -> bool:
     @asynccontextmanager
     async def get_worker(
         self,
-        package: PackageLocation,
+        package: Package | PackageLocation,
         user: str | None,
         context: str,
         permissions: PackagePermissions,
@@ -122,17 +124,25 @@ async def get_worker(
         A context manager is used to ensure that a worker is always given back to the pool.
 
         Args:
-            package: path to QuestionPy package
+            package: The main package to be run, either as a [Package][questionpy_server.package.Package] instance, or a
+                     specific [PackageLocation][questionpy_server.worker.runtime.package_location.PackageLocation].
             user: the user requesting the worker
             context: context within the lms
             permissions: package permissions
             environment_variables: environment variables to be set in the worker
 
         Returns:
-            A worker
+            A new or previously idle worker running the given package.
         """
         self._workers_requested += 1
 
+        if isinstance(package, Package):
+            manifest = package.manifest
+            package_location: PackageLocation = await package.get_zip_package_location()
+        else:
+            manifest = await read_manifest_from_location(package)
+            package_location = package
+
         # Limit the number of running workers.
         async with self._semaphore:
             worker = None
@@ -149,15 +159,20 @@ async def get_worker(
                 async with self._lock, self._condition:
                     await self._condition.wait_for(lambda: self._memory_available(permissions.memory))
                     worker = await self._create_or_reuse_worker(
-                        package, user, context, permissions, environment_variables
+                        package_location=package_location,
+                        manifest=manifest,
+                        user=user,
+                        context=context,
+                        permissions=permissions,
+                        environment_variables=environment_variables,
                     )
                     self._workers_in_use += 1
 
                 yield worker
             finally:
                 if worker:
                     async with self._condition:
-                        await self._handle_idle_worker(package, user, context, worker)
+                        await self._handle_idle_worker(package_location, user, context, worker)
                         self._condition.notify()
                         self._workers_in_use -= 1
 
@@ -219,16 +234,17 @@ def _generate_worker_name(self, package: PackageLocation) -> str:
 
     async def _create_or_reuse_worker(
         self,
-        package: PackageLocation,
+        *,
+        package_location: PackageLocation,
+        manifest: ComparableManifest,
         user: str | None,
         context: str,
         permissions: PackagePermissions,
         environment_variables: dict[str, str],
     ) -> Worker:
         """If possible, get an idle worker or create a new one."""
-        # Since the `PackagePermissions` only dependent on the `user` and `context` the worker
-        # permissions are the same.
-        identifier = _IdleWorkersIdentifier(package, user, context)
+        # Since the `PackagePermissions` only depend on the `user` and `context`, the worker permissions are the same.
+        identifier = _IdleWorkersIdentifier(package_location, user, context)
         if identifier in self._idle_workers:
             # There is an idle worker with this package loaded - reuse the most recent one.
             worker = self._idle_workers[identifier].popleft()
@@ -239,18 +255,21 @@ async def _create_or_reuse_worker(
 
             self._memory_idle -= worker.permissions.memory
         else:
+            dependencies = await self._dependency_resolver.resolve_and_retrieve(manifest)
+
             # We need to create a new worker - free as much memory as needed to start the worker.
             await self._free_memory(permissions.memory)
 
-            name = self._generate_worker_name(package)
+            name = self._generate_worker_name(package_location)
             worker_home = self._working_dir / f"worker-{name}"
             await asyncio.to_thread(worker_home.mkdir)
 
             worker = self._worker_type(
                 name=name,
-                package=package,
+                package=package_location,
                 permissions=permissions,
                 worker_home=worker_home,
+                dependencies=dependencies,
                 environment_variables=environment_variables,
             )
             await worker.start()

questionpy_server/worker/runtime/manager.py

@@ -7,6 +7,7 @@
 from contextlib import contextmanager
 from dataclasses import dataclass
 from graphlib import TopologicalSorter
+from itertools import chain
 from pathlib import Path
 from types import MappingProxyType
 from typing import TYPE_CHECKING, NoReturn, TypeVar, cast
@@ -22,6 +23,7 @@
     set_qpy_environment,
 )
 from questionpy_common.manifest import PackageType
+from questionpy_server.dependencies import SolutionAndLocation, StaticDependencySolution
 from questionpy_server.worker.runtime.connection import WorkerToServerConnection
 from questionpy_server.worker.runtime.messages import (
     CreateQuestionFromOptions,
@@ -87,12 +89,14 @@ def register_on_request_callback(self, callback: OnRequestCallback) -> None:
 type OnMessageCallback[M: MessageToWorker] = Callable[[M], MessageToServer]
 
 
-def _linearize_packages(
-    packages: Mapping[PackageNamespaceAndShortName, ImportablePackage],
+def _linearize_dependencies(
+    solutions: Mapping[PackageNamespaceAndShortName, SolutionAndLocation],
 ) -> Sequence[PackageNamespaceAndShortName]:
     sorter = TopologicalSorter[PackageNamespaceAndShortName]()
-    for nssn, package in packages.items():
-        sorter.add(nssn, *package.dependencies.keys())
+
+    for nssn, (solution, _) in solutions.items():
+        dep_nssns = [PackageNamespaceAndShortName(dep.namespace, dep.short_name) for dep in solution.dependencies.qpy]
+        sorter.add(nssn, *dep_nssns)
 
     return tuple(sorter.static_order())
 
@@ -159,54 +163,48 @@ def _open_package(location: PackageLocation, worker_home: Path) -> ImportablePac
         # This is a separate method to allow it to be mocked separately.
         return open_qpy_package(location, worker_home)
 
-    def _open_packages_recursively(
-        self,
-        msg: LoadQPyPackage,
-        package_location: PackageLocation,
-        stack: tuple[PackageNamespaceAndShortName, ...] = (),
-    ) -> tuple[PackageNamespaceAndShortName, ImportablePackage]:
+    def on_msg_load_qpy_package(self, msg: LoadQPyPackage) -> MessageToServer:
         if not self._env or not self._worker_home:
             self._raise_not_initialized(msg)
 
-        package = self._open_package(package_location, self._worker_home)
-        nssn = PackageNamespaceAndShortName(package.manifest.namespace, package.manifest.short_name)
-
-        if nssn in stack and self._packages[nssn].manifest.version == package.manifest.version:
-            raise CircularDependencyError(nssn, stack)
-
-        if nssn in self._packages:
-            # For now, we don't support two packages using the same static dependency, even if they would use the same
-            # version. Supporting the latter case would require us to either trust or check that both dependency's
-            # content is identical.
-            err_msg = f"Package '{nssn}' is already loaded. Dependency stack: {stack}"
-            raise DependencyError(err_msg, stack)
+        root_package = self._open_package(msg.location, self._worker_home)
+        root_nssn = root_package.manifest.nssn
+        self._packages[root_nssn] = root_package
 
-        self._packages[nssn] = package
+        linearized = _linearize_dependencies(msg.dependencies)
 
-        new_stack = (*stack, nssn)
+        for nssn in reversed(linearized):
+            solution, package_location = msg.dependencies[nssn]
+            if isinstance(solution, StaticDependencySolution):
+                owner = self._packages.get(solution.owner)
+                if not owner:
+                    # Since we open packages in reverse topological order, this shouldn't happen.
+                    # (Unless the tree passed to us by the server contains errors.)
+                    err_msg = f"Cannot open static dependency '{nssn}' before owner '{solution.owner}'."
+                    raise RuntimeError(err_msg)
 
-        if len(stack) >= MAX_QPY_DEPENDENCY_LEVELS and package.manifest.dependencies.qpy:
-            raise TooDeeplyNestedDependencyError(new_stack)
+                package_location = owner.resolve_static_dependency(nssn)
 
-        for dep_location in package.resolve_static_dependencies():
-            dep_nssn, dep_package = self._open_packages_recursively(msg, dep_location, new_stack)
-            package.dependencies[dep_nssn] = dep_package
-
-        return nssn, package
-
-    def on_msg_load_qpy_package(self, msg: LoadQPyPackage) -> MessageToServer:
-        if not self._env or not self._worker_home:
-            self._raise_not_initialized(msg)
-
-        root_nssn, root_package = self._open_packages_recursively(msg, msg.location, ())
+            # MyPy doesn't narrow the type properly.
+            self._packages[nssn] = self._open_package(cast("PackageLocation", package_location), self._worker_home)
 
         if msg.main:
             self._env = dataclasses.replace(self._env, _main_package=root_package)
             set_qpy_environment(self._env)
 
-        linearized = _linearize_packages(self._packages)
-        for nssn in linearized:
+        for nssn in chain(linearized, (root_nssn,)):
             package = self._packages[nssn]
+
+            # Make the package's dependencies accessible to the package.
+            for dep in package.manifest.dependencies.qpy:
+                dep_nssn = PackageNamespaceAndShortName(dep.namespace, dep.short_name)
+                dep_package = self._packages.get(dep_nssn)
+                if not dep_package:
+                    err_msg = f"Unfulfilled dependency of '{nssn}': '{dep_nssn}'"
+                    raise RuntimeError(err_msg)
+
+                package.dependencies[dep_nssn] = dep_package
+
             if package.state < PackageState.LOADED:
                 package.load()
 

questionpy_server/worker/runtime/messages.py

@@ -18,6 +18,7 @@
 from questionpy_common.environment import PackagePermissions, RequestInfo
 from questionpy_common.error import QPyBaseError
 from questionpy_common.manifest import Manifest
+from questionpy_server.dependencies import SolutionAndLocation
 from questionpy_server.worker.runtime.package_location import PackageLocation
 
 messages_header_struct: Struct = Struct("=LL")
@@ -118,6 +119,9 @@ class LoadQPyPackage(MessageToWorker):
     main: bool
     """Set this package as the main package and execute its entry point."""
 
+    dependencies: dict[PackageNamespaceAndShortName, SolutionAndLocation]
+    """All resolved dependencies in the root package's tree. Does not include the root package itself."""
+
     class Response(MessageToServer):
         """Success message in return to LoadQPyPackage."""
 

questionpy_server/worker/runtime/package.py

@@ -81,7 +81,7 @@ def init(self, env: Environment) -> None:
         """
 
     @abstractmethod
-    def resolve_static_dependencies(self) -> list[PackageLocation]:
+    def resolve_static_dependency(self, nssn: PackageNamespaceAndShortName) -> PackageLocation:
         pass
 
 
@@ -100,12 +100,33 @@ def __repr__(self) -> str:
 
     __str__ = __repr__
 
-    def resolve_static_dependencies(self) -> list[PackageLocation]:
-        return [
-            DirPackageLocation(self.path / "dependencies" / "qpy" / dep.dir_name / DIST_DIR)
-            for dep in self.manifest.dependencies.qpy
-            if isinstance(dep, DistStaticQPyDependency)
-        ]
+    def resolve_static_dependency(self, nssn: PackageNamespaceAndShortName) -> PackageLocation:
+        dep = next(
+            (
+                dep
+                for dep in self.manifest.dependencies.qpy
+                if isinstance(dep, DistStaticQPyDependency)
+                and dep.namespace == nssn.namespace
+                and dep.short_name == nssn.short_name
+                for dep in self.manifest.dependencies.qpy
+            ),
+            None,
+        )
+        if not dep:
+            msg = f"Package '{self.manifest.nssn}' does not provide static dependency '{nssn}'."
+            raise RuntimeError(msg)
+
+        dep_dist_path = (
+            self.path / "dependencies" / "qpy" / f"{dep.namespace}-{dep.short_name}-{dep.version}" / DIST_DIR
+        )
+        if not dep_dist_path.exists():
+            msg = (
+                f"Package '{self.manifest.nssn}' lists static dependency '{nssn}', but '{dep_dist_path}' is not "
+                f"present."
+            )
+            raise RuntimeError(msg)
+
+        return DirPackageLocation(dep_dist_path)
 
     def load(self) -> None:
         for new_path in (
@@ -166,8 +187,8 @@ def __repr__(self) -> str:
 
     __str__ = __repr__
 
-    def resolve_static_dependencies(self) -> list[PackageLocation]:
-        return []
+    def resolve_static_dependency(self, nssn: PackageNamespaceAndShortName) -> PackageLocation:
+        raise NotImplementedError
 
 
 def _package_dir(worker_home: Path, manifest: Manifest) -> Path: